import assert from 'node:assert/strict';
import { fileURLToPath } from 'node:url';
import test from 'node:test';
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import {
StdioClientTransport,
getDefaultEnvironment,
} from '@modelcontextprotocol/sdk/client/stdio.js';
import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js';
import { Accounts } from '../src/accounts.js';
import { Colab } from '../src/colab.js';
import { createServer } from '../src/server.js';
import { ACCESS, PROXY_TOKEN, RUNTIME, startFake } from './fixtures/fake-colab.js';
const TOOLS = [
'colab_accounts',
'colab_execute',
'colab_interrupt',
'colab_kernels',
'colab_launch',
'colab_list',
'colab_list_files',
'colab_read_file',
'colab_stop',
'colab_usage',
'colab_write_file',
];
const NBH = 'a'.repeat(44);
async function setup(t, { fetch } = {}) {
const fake = await startFake();
const colab = new Colab({
auth: { token: async () => ACCESS },
baseUrl: fake.url,
allowLocalProxy: true,
fetch,
});
const [clientSide, serverSide] = InMemoryTransport.createLinkedPair();
const server = createServer(new Accounts({ env: {}, single: colab }));
const client = new Client({ name: 'test', version: '0' });
await Promise.all([server.connect(serverSide), client.connect(clientSide)]);
t.after(async () => {
colab.close();
await client.close();
await fake.close();
});
const call = async (name, args = {}, options) => {
let result;
try {
result = await client.callTool({ name, arguments: args }, undefined, options);
} catch (error) {
return { isError: true, text: error.message, thrown: true };
}
const text = result.content[0].text;
return { isError: Boolean(result.isError), text, result, json: () => JSON.parse(text) };
};
return { fake, colab, client, call };
}
const networkCalls = (fake) => fake.requests.filter((r) => r.path !== '/tun/m/assignments');
async function waitFor(predicate) {
for (let attempt = 0; attempt < 100; attempt += 1) {
if (predicate()) return;
await new Promise((resolve) => setTimeout(resolve, 10));
}
assert.fail('expected fake Colab request did not arrive');
}
test('tools/list is the exact inventory and every tool has readOnlyHint true', async (t) => {
const { client } = await setup(t);
const { tools } = await client.listTools();
assert.deepEqual(tools.map((tool) => tool.name).sort(), TOOLS);
for (const tool of tools) assert.strictEqual(tool.annotations?.readOnlyHint, true, tool.name);
});
// COLAB_TEST_COMMAND="bun" (or a compiled binary) runs the fixture under another runtime
test('real stdio subprocess: clean protocol, launch, execute, stop', async (t) => {
const fake = await startFake();
const fixture = fileURLToPath(new URL('./fixtures/stdio-server.js', import.meta.url));
const [command, ...extra] = process.env.COLAB_TEST_COMMAND?.split(' ') ?? [process.execPath];
const args = [...extra, fixture];
const transport = new StdioClientTransport({
command,
args,
env: { ...getDefaultEnvironment(), FAKE_COLAB_URL: fake.url },
stderr: 'pipe',
});
let stderr = '';
transport.stderr.on('data', (d) => (stderr += d));
const client = new Client({ name: 'test', version: '0' });
t.after(async () => {
await client.close();
await fake.close();
});
await client.connect(transport);
assert.equal(client.getServerVersion().name, 'colab-mcp');
const { tools } = await client.listTools();
assert.equal(tools.length, TOOLS.length);
const text = async (name, args) =>
JSON.parse((await client.callTool({ name, arguments: args })).content[0].text);
const launched = await text('colab_launch', { notebook_id: NBH });
assert.equal(launched.id, RUNTIME);
const run = await client.callTool({
name: 'colab_execute',
arguments: { id: RUNTIME, code: 'print(1)' },
});
assert.ok(!run.isError);
assert.match(run.content[0].text, /out:print\(1\)/);
assert.equal((await text('colab_stop', { id: RUNTIME })).stopped, true);
await client.close();
assert.equal(stderr, '', 'server must not write noise to stderr');
});
test('launch does GET then POST with the XSRF token and the same nbh', async (t) => {
const { fake, call } = await setup(t);
const res = await call('colab_launch', { accelerator: 'A100', high_memory: true });
assert.ok(!res.isError, res.text);
const [get, post, ...rest] = fake.requests;
assert.equal(rest.length, 0);
assert.deepEqual([get.method, post.method], ['GET', 'POST']);
assert.equal(get.headers.authorization, `Bearer ${ACCESS}`);
assert.equal(post.headers['x-goog-colab-token'], 'xsrf-1');
assert.equal(get.query.nbh, post.query.nbh);
assert.equal(get.query.nbh.length, 44);
assert.equal(post.query.accelerator, 'A100');
assert.equal(post.query.shape, 'hm');
const body = res.json();
assert.equal(body.notebook_id, post.query.nbh);
assert.equal(body.id, RUNTIME);
});
test('launch failures are not retried and report the notebook id without raw bodies', async (t) => {
const { fake, call } = await setup(t);
fake.routes = (req, res, json) => {
if (req.method === 'POST' && req.path === '/tun/m/assign')
return json(res, 500, 'RAW-POST-BODY ' + ACCESS);
};
const post = await call('colab_launch', { notebook_id: NBH });
assert.ok(post.isError);
assert.equal(fake.count('GET', '/tun/m/assign'), 1);
assert.equal(fake.count('POST', '/tun/m/assign'), 1);
assert.ok(post.text.includes(NBH));
assert.ok(!/RAW-POST-BODY|fake-access/.test(post.text));
fake.requests.length = 0;
fake.routes = (req, res, json) => {
if (req.method === 'GET' && req.path === '/tun/m/assign') return json(res, 503, 'RAW-GET-BODY');
};
const get = await call('colab_launch', { notebook_id: NBH });
assert.ok(get.isError);
assert.match(get.text, /capacity unavailable/);
assert.equal(fake.requests.length, 1);
assert.ok(!get.text.includes('RAW-GET-BODY'));
});
test('stop does the destructive handshake and treats 404 as already stopped', async (t) => {
const { fake, call } = await setup(t);
const res = await call('colab_stop', { id: RUNTIME });
assert.deepEqual(res.json(), { id: RUNTIME, stopped: true });
const [get, post] = fake.requests;
assert.deepEqual([get.method, post.method], ['GET', 'POST']);
assert.equal(post.headers['x-goog-colab-token'], 'dealloc-1');
assert.equal(post.path, `/tun/m/unassign/${RUNTIME}`);
fake.routes = (req, res, json) => {
if (req.method === 'GET' && req.path.startsWith('/tun/m/unassign/')) return json(res, 404, {});
};
const gone = await call('colab_stop', { id: 'missing' });
assert.ok(!gone.isError);
assert.equal(gone.json().stopped, false);
});
test('list returns public fields only and leaks no proxy url or tokens', async (t) => {
const { fake, call } = await setup(t);
const res = await call('colab_list');
assert.deepEqual(res.json().runtimes, [
{ id: RUNTIME, variant: 'GPU', accelerator: 'T4', status: 'assigned', account_id: 'default' },
]);
for (const secret of [PROXY_TOKEN, ACCESS, fake.url, '127.0.0.1'])
assert.ok(!res.text.includes(secret), secret);
});
test('usage passes the account info through', async (t) => {
const { call } = await setup(t);
assert.deepEqual((await call('colab_usage')).json(), {
currentBalance: 42,
consumptionRateHourly: 1.5,
account_id: 'default',
});
});
test('kernels lists state without creating one', async (t) => {
const { fake, call } = await setup(t);
fake.kernels.push({ id: 'k9', name: 'python3', execution_state: 'idle', connections: 3 });
const res = await call('colab_kernels', { id: RUNTIME });
assert.deepEqual(res.json().kernels, [{ id: 'k9', name: 'python3', execution_state: 'idle' }]);
assert.equal(fake.count('POST', '/api/kernels'), 0);
});
test('two executions reuse one kernel with their own message ids', async (t) => {
const { fake, call } = await setup(t);
const seen = [];
fake.execute = (ctx) => {
seen.push(ctx.msgId);
ctx.send(ctx.frame('stream', 'unrelated-msg', { name: 'stdout', text: 'NOISE' }));
// reply arrives before idle for the first run, after idle for the second
const reply = () => ctx.send(ctx.frame('execute_reply', ctx.msgId, { status: 'ok' }));
const idle = () => ctx.send(ctx.frame('status', ctx.msgId, { execution_state: 'idle' }));
ctx.send(ctx.frame('stream', ctx.msgId, { name: 'stdout', text: ctx.code }));
if (seen.length === 1) {
reply();
idle();
} else {
idle();
reply();
}
};
const a = await call('colab_execute', { id: RUNTIME, code: 'one' });
const b = await call('colab_execute', { id: RUNTIME, code: 'two' });
assert.equal(fake.count('POST', '/api/kernels'), 1);
assert.equal(fake.count('WS', '/api/kernels/k1/channels'), 2);
assert.notEqual(seen[0], seen[1]);
assert.equal(a.json().kernel_id, b.json().kernel_id);
assert.deepEqual(
a.json().outputs.map((o) => o.text),
['one'],
);
assert.deepEqual(
b.json().outputs.map((o) => o.text),
['two'],
);
assert.ok(!a.text.includes('NOISE') && !b.text.includes('NOISE'));
});
test('binary websocket frames are decoded', async (t) => {
const { fake, call } = await setup(t);
fake.binary = true;
const res = await call('colab_execute', { id: RUNTIME, code: 'bin' });
assert.ok(!res.isError, res.text);
assert.equal(res.json().outputs[0].text, 'out:bin');
});
test('python errors come back as isError with the traceback output', async (t) => {
const { call } = await setup(t);
const res = await call('colab_execute', { id: RUNTIME, code: 'boom' });
assert.ok(res.isError);
const body = res.json();
assert.equal(body.status, 'error');
assert.equal(body.outputs.find((o) => o.type === 'error').ename, 'ValueError');
});
test('timeout interrupts and keeps partial output', async (t) => {
const { fake, call } = await setup(t);
fake.execute = (ctx) =>
ctx.send(ctx.frame('stream', ctx.msgId, { name: 'stdout', text: 'partial' }));
const res = await call('colab_execute', { id: RUNTIME, code: 'sleep', timeout: 1 });
assert.ok(res.isError);
const body = res.json();
assert.equal(body.error, 'execution timed out');
assert.deepEqual(
body.outputs.map((o) => o.text),
['partial'],
);
assert.equal(fake.count('POST', '/api/kernels/k1/interrupt'), 1);
});
test('websocket closing early is reported and interrupts', async (t) => {
const { fake, call } = await setup(t);
fake.execute = ({ ws }) => ws.close();
const res = await call('colab_execute', { id: RUNTIME, code: 'x' });
assert.ok(res.isError);
assert.match(res.json().error, /closed before execution completed/);
assert.equal(fake.count('POST', '/api/kernels/k1/interrupt'), 1);
});
test('output is capped at 1 MiB and marked truncated', async (t) => {
const { fake, call } = await setup(t);
fake.execute = (ctx) => {
const text = 'x'.repeat(100_000);
for (let i = 0; i < 20; i += 1)
ctx.send(ctx.frame('stream', ctx.msgId, { name: 'stdout', text }));
ctx.send(ctx.frame('execute_reply', ctx.msgId, { status: 'ok' }));
ctx.send(ctx.frame('status', ctx.msgId, { execution_state: 'idle' }));
};
const res = await call('colab_execute', { id: RUNTIME, code: 'big' });
const body = res.json();
assert.equal(body.truncated, true);
assert.ok(body.outputs.length >= 9 && body.outputs.length < 20);
assert.ok(Buffer.byteLength(res.text) < 1.1 * 1024 * 1024);
});
test('cancelling the tool call interrupts the kernel', async (t) => {
const { fake, call } = await setup(t);
fake.execute = () => {};
const controller = new AbortController();
const pending = call(
'colab_execute',
{ id: RUNTIME, code: 'hang', timeout: 30 },
{ signal: controller.signal },
);
await waitFor(() => fake.count('WS', '/api/kernels/k1/channels') > 0);
controller.abort();
assert.ok((await pending).isError);
await waitFor(() => fake.count('POST', '/api/kernels/k1/interrupt') > 0);
assert.equal(fake.count('POST', '/api/kernels/k1/interrupt'), 1);
});
test('a second execution on a busy runtime is refused', async (t) => {
const { fake, call } = await setup(t);
fake.execute = () => {};
const first = call('colab_execute', { id: RUNTIME, code: 'a', timeout: 1 });
await waitFor(() => fake.count('WS', '/api/kernels/k1/channels') > 0);
const second = await call('colab_execute', { id: RUNTIME, code: 'b' });
assert.ok(second.isError);
assert.match(second.text, /already executing/);
await first;
});
test('invalid input is rejected before any network call', async (t) => {
const { fake, call } = await setup(t);
const cases = [
['colab_list', { extra: 1 }],
['colab_stop', { id: RUNTIME, force: true }],
['colab_execute', { id: RUNTIME, code: '' }],
['colab_execute', { id: RUNTIME, code: 'x', timeout: 0 }],
['colab_launch', { accelerator: 'TPUX' }],
['colab_launch', { notebook_id: 'short' }],
['colab_read_file', { id: RUNTIME, path: 'a', format: 'hex' }],
['colab_write_file', { id: RUNTIME, path: 'a', content: 'x'.repeat(1024 * 1024 + 1) }],
];
for (const [name, args] of cases) {
const res = await call(name, args);
assert.ok(res.isError, `${name} ${JSON.stringify(args).slice(0, 60)}`);
}
assert.equal(fake.requests.length, 0);
const bad = await call('colab_write_file', {
id: RUNTIME,
path: 'a',
format: 'base64',
content: 'not base64!',
});
assert.ok(bad.isError);
assert.equal(networkCalls(fake).length, 0);
});
test('file write then read round-trips text, unicode paths and base64', async (t) => {
const { fake, call } = await setup(t);
const unicode = 'dir ü/файл #1?.txt';
fake.files.set('dir ü', { type: 'directory' });
assert.ok(
(await call('colab_write_file', { id: RUNTIME, path: unicode, content: 'héllo ☃' })).json()
.written,
);
const put = fake.requests.find((r) => r.method === 'PUT');
assert.equal(put.path, `/api/contents/${unicode}`);
assert.deepEqual(put.body, { type: 'file', format: 'text', content: 'héllo ☃' });
assert.equal(put.headers['x-colab-runtime-proxy-token'], PROXY_TOKEN);
const read = await call('colab_read_file', { id: RUNTIME, path: unicode });
assert.equal(read.json().content, 'héllo ☃');
const bytes = Buffer.from([0, 255, 1, 128]);
await call('colab_write_file', {
id: RUNTIME,
path: 'bin.dat',
format: 'base64',
content: bytes.toString('base64'),
});
assert.deepEqual(fake.files.get('bin.dat').data, bytes);
const back = await call('colab_read_file', { id: RUNTIME, path: '/bin.dat', format: 'base64' });
assert.deepEqual(Buffer.from(back.json().content, 'base64'), bytes);
});
test('list_files returns whitelisted entry fields only', async (t) => {
const { fake, call } = await setup(t);
fake.files.set('a.txt', { type: 'file', data: Buffer.from('hi') });
fake.files.set('sub', { type: 'directory' });
const res = await call('colab_list_files', { id: RUNTIME });
const entries = res.json().entries;
assert.deepEqual(entries.map((e) => e.name).sort(), ['a.txt', 'sub']);
assert.ok(!res.text.includes('must-not-leak'));
const sub = await call('colab_list_files', { id: RUNTIME, path: 'a.txt' });
assert.ok(sub.isError);
assert.match(sub.text, /not a directory/);
});
test('path traversal is refused before any contents request', async (t) => {
const { fake, call } = await setup(t);
for (const path of ['../etc/passwd', 'a/../../b', './x', 'a/./b', 'bad\0name']) {
for (const [name, extra] of [
['colab_read_file', {}],
['colab_write_file', { content: 'x' }],
['colab_list_files', {}],
]) {
const res = await call(name, { id: RUNTIME, path, ...extra });
assert.ok(res.isError, `${name} ${JSON.stringify(path)}`);
}
}
assert.equal(fake.requests.filter((r) => r.path.startsWith('/api/contents')).length, 0);
});
test('read refuses files over 1 MiB and failures leak no raw bodies or tokens', async (t) => {
const { fake, call } = await setup(t);
fake.files.set('big.bin', { type: 'file', data: Buffer.alloc(1024 * 1024 + 10, 97) });
const big = await call('colab_read_file', { id: RUNTIME, path: 'big.bin' });
assert.ok(big.isError);
assert.match(big.text, /exceeds 1 MiB/);
const missing = await call('colab_read_file', { id: RUNTIME, path: 'nope.txt' });
assert.ok(missing.isError);
assert.match(missing.text, /not found/);
for (const secret of ['SECRET-RAW-BODY', PROXY_TOKEN, ACCESS])
assert.ok(!missing.text.includes(secret), secret);
});
test('oversized HTTP responses are cut off at 8 MiB', async (t) => {
const { fake, call } = await setup(t);
fake.routes = (req, res) => {
if (req.path !== '/tun/m/ccu-info') return false;
res.writeHead(200, { 'content-type': 'application/json' });
res.end(JSON.stringify({ pad: 'x'.repeat(9 * 1024 * 1024) }));
return true;
};
const res = await call('colab_usage');
assert.ok(res.isError);
assert.match(res.text, /8 MiB/);
});
test('an untrusted runtime proxy url is refused and a failed auth response leaks nothing', async (t) => {
const { fake, call } = await setup(t);
fake.routes = (req, res, json) => {
if (req.path !== '/tun/m/assignments') return false;
return json(res, 200, {
assignments: [
{ endpoint: RUNTIME, runtimeProxyInfo: { url: 'http://evil.example', token: 't' } },
],
});
};
const res = await call('colab_read_file', { id: RUNTIME, path: 'a' });
assert.ok(res.isError);
assert.match(res.text, /untrusted/);
fake.routes = (req, res, json) => json(res, 401, 'RAW-401 ' + ACCESS);
const unauth = await call('colab_usage');
assert.match(unauth.text, /colab-mcp login/);
assert.ok(!/RAW-401|fake-access/.test(unauth.text));
});
test('malformed websocket json (null, non-json) fails the execution without hanging', async (t) => {
const { fake, call } = await setup(t);
for (const raw of ['null', 'not json']) {
fake.execute = ({ ws }) => ws.send(raw);
const res = await call('colab_execute', { id: RUNTIME, code: 'x', timeout: 5 });
assert.ok(res.isError, raw);
assert.match(res.json().error, /invalid Jupyter websocket message/);
}
});
test('an aborted execute_reply is reported as an error', async (t) => {
const { fake, call } = await setup(t);
fake.execute = (ctx) => {
ctx.send(ctx.frame('execute_reply', ctx.msgId, { status: 'aborted' }));
ctx.send(ctx.frame('status', ctx.msgId, { execution_state: 'idle' }));
};
const res = await call('colab_execute', { id: RUNTIME, code: 'x' });
assert.ok(res.isError);
assert.equal(res.json().status, 'error');
});
test('multiple kernels need an explicit kernel_id', async (t) => {
const { fake, call } = await setup(t);
fake.kernels.push({ id: 'k1', name: 'python3' }, { id: 'k2', name: 'python3' });
const ambiguous = await call('colab_execute', { id: RUNTIME, code: 'x' });
assert.ok(ambiguous.isError);
assert.match(ambiguous.text, /multiple kernels/);
assert.equal(
fake.count('WS', '/api/kernels/k1/channels') + fake.count('WS', '/api/kernels/k2/channels'),
0,
);
const unknown = await call('colab_execute', { id: RUNTIME, code: 'x', kernel_id: 'nope' });
assert.ok(unknown.isError);
const ok = await call('colab_execute', { id: RUNTIME, code: 'x', kernel_id: 'k2' });
assert.ok(!ok.isError, ok.text);
assert.equal(fake.count('WS', '/api/kernels/k2/channels'), 1);
const interrupt = await call('colab_interrupt', { id: RUNTIME });
assert.ok(interrupt.isError);
assert.equal(fake.count('POST', '/api/kernels/k2/interrupt'), 0);
});
test('image outputs are returned as MCP image content, not inline text', async (t) => {
const { fake, call } = await setup(t);
const png = 'iVBORw0KGgo\nAAAA';
fake.execute = (ctx) => {
ctx.send(
ctx.frame('display_data', ctx.msgId, {
data: { 'image/png': png, 'text/plain': ['<Figure>'] },
}),
);
ctx.send(ctx.frame('execute_reply', ctx.msgId, { status: 'ok' }));
ctx.send(ctx.frame('status', ctx.msgId, { execution_state: 'idle' }));
};
const { result, text } = await call('colab_execute', { id: RUNTIME, code: 'plot' });
const image = result.content.find((c) => c.type === 'image');
assert.deepEqual([image.mimeType, image.data], ['image/png', 'iVBORw0KGgoAAAA']);
assert.ok(!text.includes('iVBORw0KGgo'));
});
test('a rotated runtime proxy token is used for the timeout interrupt', async (t) => {
const { fake, call } = await setup(t);
fake.execute = (ctx) => {
ctx.send(ctx.frame('stream', ctx.msgId, { name: 'stdout', text: 'partial' }));
fake.proxyToken = 'proxy-rotated';
};
const res = await call('colab_execute', { id: RUNTIME, code: 'sleep', timeout: 1 });
assert.ok(res.isError);
assert.ok(!res.json().error.includes('interrupt failed'), res.text);
const interrupts = fake.requests.filter((r) => r.path === '/api/kernels/k1/interrupt');
assert.equal(interrupts.at(-1).headers['x-colab-runtime-proxy-token'], 'proxy-rotated');
assert.ok(!res.text.includes('proxy-rotated'));
});
test('proxy accepts Google production colab.dev hosts but rejects lookalikes and insecure URLs', async (t) => {
const { fake, colab } = await setup(t);
let proxyUrl;
fake.routes = (req, res, json) => {
if (req.path === '/tun/m/assignments') {
return json(res, 200, {
assignments: [
{ endpoint: RUNTIME, runtimeProxyInfo: { url: proxyUrl, token: PROXY_TOKEN } },
],
});
}
};
for (const url of [
'https://8080-test.us-central1-1.prod.colab.dev',
'https://test.colab.googleusercontent.com',
]) {
proxyUrl = url;
assert.equal((await colab.proxy(RUNTIME)).url, url);
}
for (const url of [
'http://test.prod.colab.dev',
'https://test.prod.colab.dev.evil.example',
'https://evil.example/prod.colab.dev',
'https://user:secret@test.prod.colab.dev',
'https://test.prod.colab.dev/?token=secret',
]) {
proxyUrl = url;
await assert.rejects(colab.proxy(RUNTIME), /untrusted|invalid/);
}
});