import assert from 'node:assert/strict';
import { randomUUID } from 'node:crypto';
import { mkdtemp, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join, resolve } from 'node:path';
import { setTimeout as delay } from 'node:timers/promises';
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
import { Auth } from '../src/auth.js';
import { Colab } from '../src/colab.js';
if (process.env.COLAB_LIVE_TEST !== '1') {
throw new Error(
'set COLAB_LIVE_TEST=1 to authorize a billed T4 runtime; it will be stopped afterwards',
);
}
const direct = new Colab({ auth: new Auth() });
const notebookId = randomUUID().replaceAll('-', '_').padEnd(44, '.');
const marker = randomUUID().replaceAll('-', '');
const recordDir = await mkdtemp(join(tmpdir(), 'colab-mcp-live-'));
const recordPath = join(recordDir, 'result.json');
const record = { notebook_id: notebookId, started: new Date().toISOString(), checks: [] };
const save = () => writeFile(recordPath, JSON.stringify(record, null, 2), { mode: 0o600 });
console.log(`live test record: ${recordPath}`);
await save();
const baseline = new Set((await direct.list()).runtimes.map((runtime) => runtime.id));
let client;
let stderr = '';
let runtimeId;
let launchAttempted = false;
let cleanupVerified = false;
const cancellation = new AbortController();
const cancel = () => cancellation.abort();
process.once('SIGINT', cancel);
process.once('SIGTERM', cancel);
const cleanupIds = new Set();
const raw = (name, args = {}) =>
client.callTool({ name, arguments: args }, undefined, {
timeout: 180000,
signal: cancellation.signal,
});
const json = (result) => JSON.parse(result.content[0].text);
const text = (result) =>
result.content
.filter((part) => part.type === 'text')
.map((part) => part.text)
.join('\n');
const call = async (name, args = {}) => {
const result = await raw(name, args);
assert.ok(!result.isError, `${name}: ${text(result)}`);
return json(result);
};
const passed = async (label) => {
record.checks.push(label);
await save();
console.log(`passed: ${label}`);
};
async function connect(binary) {
const transport = new StdioClientTransport({
command: binary ? resolve(binary) : process.execPath,
args: binary ? [] : [resolve('src/index.js')],
env: Object.fromEntries(
Object.entries(process.env).filter(([key]) =>
[
'HOME',
'PATH',
'COLAB_CREDENTIALS',
'COLAB_ACCESS_TOKEN',
'COLAB_CLIENT_ID',
'COLAB_CLIENT_SECRET',
].includes(key),
),
),
stderr: 'pipe',
});
transport.stderr.on('data', (data) => {
stderr += data;
});
client = new Client({ name: 'colab-live-smoke', version: '1.0.0' });
await client.connect(transport);
const { tools } = await client.listTools();
assert.equal(tools.length, 11);
assert.ok(tools.every((tool) => tool.annotations?.readOnlyHint === true));
}
try {
await connect();
await passed('stdio initialize and all 10 readOnlyHint annotations');
await call('colab_list');
await call('colab_usage');
launchAttempted = true;
const launched = await call('colab_launch', { accelerator: 'T4', notebook_id: notebookId });
runtimeId = launched.id;
assert.ok(!baseline.has(runtimeId), 'refusing to touch a pre-existing runtime');
cleanupIds.add(runtimeId);
record.runtime_id = runtimeId;
await passed('fresh T4 allocation');
const assignment = (await direct.assignments()).find((item) => item.endpoint === runtimeId);
record.proxy_hostname = new URL(assignment.runtimeProxyInfo.url).hostname;
console.log(`runtime proxy host: ${record.proxy_hostname}`);
await save();
const before = await call('colab_kernels', { id: runtimeId });
record.initial_kernel_count = before.kernels.length;
const execute = (code, extra = {}) =>
raw('colab_execute', { id: runtimeId, code, timeout: 45, ...extra });
const first = await execute(`import os, sys, subprocess, json
_colab_mcp_smoke_counter = 41
gpu = subprocess.check_output(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], text=True).strip()
assert 'T4' in gpu, gpu
print('LIVE_BOOT', json.dumps({'cwd': os.getcwd(), 'gpu': gpu, 'python': sys.version.split()[0]}), flush=True)`);
assert.ok(!first.isError, text(first));
const kernelId = json(first).kernel_id;
record.kernel_id = kernelId;
record.bootstrap_output = json(first)
.outputs.filter((output) => output.type === 'stream')
.map((output) => output.text)
.join('');
assert.match(text(first), /LIVE_BOOT/);
await passed('real Jupyter execution and nvidia-smi confirms T4');
const second = await execute(
"_colab_mcp_smoke_counter += 1; assert _colab_mcp_smoke_counter == 42; print('LIVE_PERSIST_42')",
);
assert.ok(!second.isError, text(second));
assert.equal(json(second).kernel_id, kernelId);
assert.match(text(second), /LIVE_PERSIST_42/);
await passed('Python state persists across websocket connections');
const failure = await execute("raise ValueError('mcp-live-expected')");
assert.equal(failure.isError, true);
assert.match(text(failure), /ValueError/);
await passed('Python exception surfaces as MCP isError');
const filename = `colab-mcp-smoke-${marker}.txt`;
const content = 'héllo λ\n';
await call('colab_write_file', { id: runtimeId, path: filename, content });
assert.equal((await call('colab_read_file', { id: runtimeId, path: filename })).content, content);
const listing = await call('colab_list_files', { id: runtimeId });
assert.ok(listing.entries.some((entry) => entry.name === filename));
const bytes = Buffer.from([0, 255, 1, 128]);
await call('colab_write_file', {
id: runtimeId,
path: `${filename}.bin`,
format: 'base64',
content: bytes.toString('base64'),
});
assert.deepEqual(
Buffer.from(
(await call('colab_read_file', { id: runtimeId, path: `${filename}.bin`, format: 'base64' }))
.content,
'base64',
),
bytes,
);
const paths = await execute(`from pathlib import Path
_candidates = [Path.cwd() / ${JSON.stringify(filename)}, Path('/') / ${JSON.stringify(filename)}, Path('/content') / ${JSON.stringify(filename)}, Path('/root') / ${JSON.stringify(filename)}]
_smoke_file = next(p for p in _candidates if p.is_file() and p.read_text() == ${JSON.stringify(content)})
print('LIVE_CONTENTS_ROOT', str(_smoke_file.parent), flush=True)`);
assert.ok(!paths.isError, text(paths));
record.contents_output = json(paths)
.outputs.filter((output) => output.type === 'stream')
.map((output) => output.text)
.join('');
await passed('text, unicode and binary file round-trips, listing, Python cross-check');
const image = await execute(`from IPython.display import display, Image
import base64
display(Image(data=base64.b64decode('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII='), format='png'))`);
assert.ok(!image.isError, text(image));
assert.ok(image.content.some((part) => part.type === 'image' && part.mimeType === 'image/png'));
await passed('real IPython display becomes MCP image content');
const startedPath = `${filename}.started`;
const pending = execute(
`import time
(_smoke_file.parent / ${JSON.stringify(startedPath)}).write_text('started')
print('LIVE_SLEEP_STARTED', flush=True)
time.sleep(60)`,
{ timeout: 90 },
);
pending.catch(() => {});
let started = false;
for (let attempt = 0; attempt < 40; attempt += 1) {
const result = await raw('colab_read_file', { id: runtimeId, path: startedPath });
if (!result.isError && json(result).content === 'started') {
started = true;
break;
}
await delay(250);
}
assert.ok(started, 'sleep code did not create its ready marker');
await call('colab_interrupt', { id: runtimeId, kernel_id: kernelId });
const interrupted = await pending;
assert.equal(interrupted.isError, true);
assert.match(text(interrupted), /KeyboardInterrupt/);
await passed('explicit interrupt stops a running sleep');
const timed = await execute(
"import time; print('LIVE_TIMEOUT_STARTED', flush=True); time.sleep(30)",
{ timeout: 2 },
);
assert.equal(timed.isError, true);
assert.match(json(timed).error, /timed out/);
assert.ok(!json(timed).error.includes('interrupt failed'), text(timed));
await passed('timeout interrupts running code and returns partial output');
const after = await execute(
"assert _colab_mcp_smoke_counter == 42; print('LIVE_AFTER_INTERRUPT')",
);
assert.ok(!after.isError, text(after));
assert.match(text(after), /LIVE_AFTER_INTERRUPT/);
await passed('kernel remains usable after interrupts');
await client.close();
client = undefined;
await connect(process.argv[2] ?? 'dist/colab-mcp');
const native = await raw('colab_execute', {
id: runtimeId,
kernel_id: kernelId,
code: "assert _colab_mcp_smoke_counter == 42; print('LIVE_NATIVE_PERSIST')",
timeout: 45,
});
assert.ok(!native.isError, text(native));
assert.match(text(native), /LIVE_NATIVE_PERSIST/);
await passed('standalone binary reconnects and preserves Python state');
assert.equal(stderr, '', 'server wrote unexpected stderr');
const stopped = await call('colab_stop', { id: runtimeId });
assert.equal(stopped.stopped, true);
assert.ok(!(await call('colab_list')).runtimes.some((runtime) => runtime.id === runtimeId));
cleanupVerified = true;
await passed('MCP stop and list confirm instance removed');
record.passed = true;
} catch (error) {
record.passed = false;
record.error = error.message;
console.error(`live test failed: ${error.message}`);
process.exitCode = 1;
} finally {
try {
if (launchAttempted && !runtimeId) {
const answer = await direct.api('/tun/m/assign', {
query: { nbh: notebookId, variant: 'GPU', accelerator: 'T4' },
});
const id = (answer.assignment ?? answer).endpoint;
if (id && !baseline.has(id)) cleanupIds.add(id);
}
for (const id of cleanupIds) {
if (!cleanupVerified) await direct.stop(id);
}
const remaining = (await direct.list()).runtimes.filter((runtime) =>
cleanupIds.has(runtime.id),
);
assert.equal(remaining.length, 0, 'test runtime remains active');
record.cleanup_verified = true;
console.log('cleanup verified: test instance absent from Google assignments');
} catch (error) {
record.cleanup_verified = false;
record.cleanup_error = error.message;
console.error(`CLEANUP FAILED: ${error.message}; inspect runtime id in ${recordPath}`);
process.exitCode = 1;
}
await client?.close();
direct.close();
process.off('SIGINT', cancel);
process.off('SIGTERM', cancel);
record.finished = new Date().toISOString();
await save();
}