import { randomUUID } from 'node:crypto';
import WebSocket from 'ws';
const API = 'https://colab.research.google.com';
const RESPONSE_LIMIT = 8 * 1024 * 1024;
const OUTPUT_LIMIT = 1024 * 1024;
export class ColabError extends Error {
constructor(message, status) {
super(message);
this.status = status;
}
}
async function jsonResponse(response) {
if (!response.ok) {
await response.body?.cancel();
const reasons = {
401: 'authentication expired; run colab-mcp login',
403: 'forbidden; check account permissions and Colab entitlement',
404: 'runtime, kernel or file not found',
412: 'runtime allocation limit reached',
429: 'rate limited',
503: 'capacity unavailable',
};
throw new ColabError(
`colab HTTP ${response.status}: ${reasons[response.status] ?? 'request failed'}`,
response.status,
);
}
const chunks = [];
let size = 0;
for await (const chunk of response.body ?? []) {
size += chunk.length;
if (size > RESPONSE_LIMIT) {
throw new ColabError('colab response exceeded the 8 MiB limit');
}
chunks.push(chunk);
}
const text = Buffer.concat(chunks)
.toString('utf8')
.replace(/^\)\]\}'[^\n]*\n/, '')
.trim();
if (!text) return {};
try {
return JSON.parse(text);
} catch {
throw new ColabError('colab returned invalid JSON');
}
}
function publicRuntime(item) {
if (!item || typeof item.endpoint !== 'string' || !item.endpoint) {
throw new ColabError('colab returned an invalid runtime assignment');
}
return {
id: item.endpoint,
variant: item.variant,
accelerator: item.accelerator,
machineShape: item.machineShape,
status: item.status ?? 'assigned',
};
}
function contentsPath(path) {
const parts = path.replace(/^\/+|\/+$/g, '').split('/');
if (parts.some((part) => part === '.' || part === '..' || part.includes('\0'))) {
throw new ColabError('use a path relative to the Jupyter contents root, without . or ..');
}
return parts.map(encodeURIComponent).join('/');
}
function decodeMessage(bytes, binary) {
if (!binary) return JSON.parse(bytes.toString());
// Jupyter's default binary framing: offset table, JSON, then optional buffers.
if (bytes.length < 8) throw new Error('short Jupyter frame');
const count = bytes.readUInt32BE(0);
if (count < 1 || count > 1024 || bytes.length < 4 * (count + 1)) {
throw new Error('invalid Jupyter buffer count');
}
const start = bytes.readUInt32BE(4);
const end = count > 1 ? bytes.readUInt32BE(8) : bytes.length;
if (start < 4 * (count + 1) || end < start || end > bytes.length) {
throw new Error('invalid Jupyter buffer offsets');
}
return JSON.parse(bytes.subarray(start, end).toString());
}
export class Colab {
constructor({ auth, fetch = globalThis.fetch, baseUrl = API, allowLocalProxy = false } = {}) {
this.auth = auth;
this.fetch = fetch;
this.baseUrl = baseUrl;
this.allowLocalProxy = allowLocalProxy;
this.busy = new Set();
this.sockets = new Set();
}
async request(url, { method = 'GET', headers = {}, body, signal, timeout = 30000 } = {}) {
const deadline = AbortSignal.timeout(timeout);
try {
return await jsonResponse(
await this.fetch(url, {
method,
redirect: 'error',
headers: { Accept: 'application/json', 'X-Colab-Client-Agent': 'colab-cli', ...headers },
body: body === undefined ? undefined : JSON.stringify(body),
signal: signal ? AbortSignal.any([deadline, signal]) : deadline,
}),
);
} catch (error) {
if (error instanceof ColabError) throw error;
if (signal?.aborted)
throw new ColabError('colab request cancelled; a mutation may still have completed');
throw new ColabError(
'colab request failed or timed out; a mutation may still have completed; inspect state before retrying',
);
}
}
async api(path, { query = {}, headers = {}, ...options } = {}) {
const token = await this.auth.token();
const url = new URL(path, this.baseUrl);
url.search = new URLSearchParams({ authuser: '0', ...query }).toString();
return this.request(url, {
...options,
headers: { Authorization: `Bearer ${token}`, 'X-Colab-Tunnel': 'Google', ...headers },
});
}
async assignments(signal) {
const data = await this.api('/tun/m/assignments', { signal });
const items = Array.isArray(data) ? data : data.assignments;
if (!Array.isArray(items)) throw new ColabError('colab returned an invalid assignments list');
return items;
}
async list(signal) {
return { runtimes: (await this.assignments(signal)).map(publicRuntime) };
}
async usage(signal) {
return this.api('/tun/m/ccu-info', { signal });
}
async launch({ accelerator = 'T4', high_memory = false, notebook_id }, signal) {
const variant =
accelerator === 'NONE' ? 'DEFAULT' : accelerator.startsWith('V') ? 'TPU' : 'GPU';
const nbh = notebook_id ?? randomUUID().replaceAll('-', '_').padEnd(44, '.');
const query = { nbh, variant, accelerator };
if (high_memory && !['L4', 'V5E1', 'V6E1'].includes(accelerator)) query.shape = 'hm';
try {
let answer = await this.api('/tun/m/assign', { query, signal, timeout: 120000 });
if (!answer.endpoint && !answer.assignment?.endpoint) {
if (typeof answer.token !== 'string')
throw new ColabError('colab did not return an allocation token');
answer = await this.api('/tun/m/assign', {
method: 'POST',
query,
signal,
timeout: 120000,
headers: { 'X-Goog-Colab-Token': answer.token },
});
}
return { ...publicRuntime(answer.assignment ?? answer), notebook_id: nbh };
} catch (error) {
throw new ColabError(
`${error.message}; notebook_id=${nbh}; list runtimes before retrying`,
error.status,
);
}
}
async stop(id, signal) {
const path = `/tun/m/unassign/${encodeURIComponent(id)}`;
try {
const answer = await this.api(path, { signal });
if (typeof answer.token !== 'string')
throw new ColabError('colab did not return a deallocation token');
await this.api(path, {
method: 'POST',
signal,
headers: { 'X-Goog-Colab-Token': answer.token },
});
return { id, stopped: true };
} catch (error) {
if (error.status === 404)
return { id, stopped: false, message: 'already stopped or not found' };
throw error;
}
}
async proxy(id, signal) {
const assignment = (await this.assignments(signal)).find((item) => item.endpoint === id);
if (!assignment) throw new ColabError('no runtime with that id; call colab_list');
const proxy = assignment.runtimeProxyInfo;
if (!proxy?.url || !proxy.token)
throw new ColabError('runtime proxy is not ready; try again after it starts');
const url = new URL(proxy.url);
const local = this.allowLocalProxy && url.protocol === 'http:' && url.hostname === '127.0.0.1';
const googleHost = ['.googleusercontent.com', '.prod.colab.dev'].some((suffix) =>
url.hostname.endsWith(suffix),
);
if (!local && (url.protocol !== 'https:' || !googleHost)) {
throw new ColabError('colab returned an untrusted runtime proxy URL');
}
if (url.username || url.password || url.search || url.hash)
throw new ColabError('invalid runtime proxy URL');
return { url: url.href.replace(/\/$/, ''), token: proxy.token };
}
proxyUrl(proxy, path, query = {}) {
const url = new URL(`${proxy.url}${path}`);
url.search = new URLSearchParams({
authuser: '0',
'colab-runtime-proxy-token': proxy.token,
...query,
}).toString();
return url;
}
runtime(proxy, path, { query, headers, ...options } = {}) {
return this.request(this.proxyUrl(proxy, path, query), {
...options,
headers: {
'X-Colab-Runtime-Proxy-Token': proxy.token,
'Content-Type': 'application/json',
...headers,
},
});
}
async files({ id, path = '', content, format = 'text' }, signal) {
const apiPath = `/api/contents/${contentsPath(path)}`;
const proxy = await this.proxy(id, signal);
if (content !== undefined) {
await this.runtime(proxy, apiPath, {
method: 'PUT',
signal,
body: { type: 'file', format, content },
});
return { id, path, written: true };
}
return this.runtime(proxy, apiPath, { query: { content: '1', format }, signal });
}
async kernelList(proxy, signal) {
const kernels = await this.runtime(proxy, '/api/kernels', { signal });
if (
!Array.isArray(kernels) ||
kernels.some((kernel) => typeof kernel?.id !== 'string' || !kernel.id)
) {
throw new ColabError('invalid Jupyter kernel list');
}
return kernels;
}
async kernel(proxy, kernelId, signal) {
const kernels = await this.kernelList(proxy, signal);
if (kernelId) {
if (!kernels.some((kernel) => kernel.id === kernelId))
throw new ColabError('kernel not found');
return kernelId;
}
if (kernels.length > 1) throw new ColabError('multiple kernels; pass kernel_id explicitly');
if (kernels.length === 1) return kernels[0].id;
const kernel = await this.runtime(proxy, '/api/kernels', {
method: 'POST',
body: { name: 'python3' },
signal,
});
if (typeof kernel.id !== 'string' || !kernel.id)
throw new ColabError('invalid Jupyter kernel id');
return kernel.id;
}
async kernels(id, signal) {
const proxy = await this.proxy(id, signal);
const kernels = await this.kernelList(proxy, signal);
return {
kernels: kernels.map(({ id, name, execution_state, last_activity }) => ({
id,
name,
execution_state,
last_activity,
})),
};
}
async interrupt({ id, kernel_id }, signal) {
const proxy = await this.proxy(id, signal);
const kernels = await this.kernelList(proxy, signal);
const target = kernel_id ?? (kernels.length === 1 ? kernels[0].id : undefined);
if (!target) throw new ColabError('pass kernel_id when there is not exactly one kernel');
await this.runtime(proxy, `/api/kernels/${encodeURIComponent(target)}/interrupt`, {
method: 'POST',
signal,
});
return { id, kernel_id: target, interrupted: true };
}
async execute({ id, code, kernel_id, timeout = 60 }, signal) {
if (this.busy.has(id))
throw new ColabError('this runtime is already executing; wait or call colab_interrupt');
if (this.busy.size >= 8)
throw new ColabError('at most 8 concurrent runtime executions are allowed');
this.busy.add(id);
try {
const proxy = await this.proxy(id, signal);
const kernelId = await this.kernel(proxy, kernel_id, signal);
return await this.executeSocket(id, proxy, kernelId, code, timeout, signal);
} finally {
this.busy.delete(id);
}
}
executeSocket(id, proxy, kernelId, code, timeout, signal) {
const session = randomUUID();
const msgId = randomUUID();
const url = this.proxyUrl(proxy, `/api/kernels/${encodeURIComponent(kernelId)}/channels`, {
session_id: session,
});
url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
return new Promise((resolve) => {
const ws = new WebSocket(url, {
headers: {
'X-Colab-Runtime-Proxy-Token': proxy.token,
'X-Colab-Client-Agent': 'colab-cli',
},
maxPayload: RESPONSE_LIMIT,
handshakeTimeout: 15000,
followRedirects: false,
});
this.sockets.add(ws);
const outputs = [];
let bytes = 0;
let truncated = false;
let reply;
let idle = false;
let settled = false;
let sent = false;
const finish = async (error) => {
if (settled) return;
settled = true;
clearTimeout(timer);
signal?.removeEventListener('abort', cancel);
this.sockets.delete(ws);
ws.terminate();
if (error && sent) {
try {
const cleanupSignal = AbortSignal.timeout(10000);
const freshProxy = await this.proxy(id, cleanupSignal);
await this.runtime(
freshProxy,
`/api/kernels/${encodeURIComponent(kernelId)}/interrupt`,
{
method: 'POST',
signal: cleanupSignal,
},
);
} catch {
error += '; interrupt failed; code may still be running';
}
}
resolve({
kernel_id: kernelId,
outputs,
truncated,
status: error || reply?.status !== 'ok' ? 'error' : 'ok',
...(error ? { error } : {}),
});
};
const cancel = () => void finish('execution cancelled');
const timer = setTimeout(() => void finish('execution timed out'), timeout * 1000);
signal?.addEventListener('abort', cancel, { once: true });
ws.on('open', () => {
if (settled) return;
sent = true;
ws.send(
JSON.stringify({
header: {
msg_id: msgId,
username: 'colab-mcp',
session,
msg_type: 'execute_request',
version: '5.3',
date: new Date().toISOString(),
},
parent_header: {},
metadata: {},
channel: 'shell',
buffers: [],
content: {
code,
silent: false,
store_history: true,
user_expressions: {},
allow_stdin: false,
stop_on_error: true,
},
}),
);
});
ws.on('message', (data, binary) => {
if (settled) return;
let message;
try {
message = decodeMessage(data, binary);
if (!message || typeof message !== 'object' || Array.isArray(message)) {
throw new Error('invalid message object');
}
} catch {
void finish('invalid Jupyter websocket message');
return;
}
if (message.parent_header?.msg_id !== msgId) return;
const type = message.header?.msg_type ?? message.msg_type;
const content = message.content ?? {};
if (type === 'execute_reply') reply = content;
if (type === 'status' && content.execution_state === 'idle') idle = true;
if (
[
'stream',
'execute_result',
'display_data',
'update_display_data',
'error',
'clear_output',
].includes(type)
) {
const output = { type, ...content };
const size = Buffer.byteLength(JSON.stringify(output));
if (bytes + size <= OUTPUT_LIMIT && outputs.length < 2000) {
bytes += size;
outputs.push(output);
} else truncated = true;
}
if (idle && reply) void finish();
});
ws.on(
'error',
(error) =>
void finish(
error.code === 'WS_ERR_UNSUPPORTED_MESSAGE_LENGTH'
? 'Jupyter websocket message exceeded the 8 MiB limit'
: 'Jupyter websocket failed; refresh runtime state before retrying',
),
);
ws.on('close', () => void finish('Jupyter websocket closed before execution completed'));
if (signal?.aborted) cancel();
});
}
close() {
for (const socket of this.sockets) socket.terminate();
}
}