import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { z } from 'zod';
const id = z
.string()
.min(1)
.max(1024)
.describe('runtime id returned by colab_list or colab_launch');
const account_id = z
.string()
.regex(/^(acct_[A-Za-z0-9_-]{43}|default)$/)
.optional()
.describe(
'opaque account id from colab_accounts; required for launch/usage when multiple accounts are connected',
);
const kernel_id = z.string().min(1).max(512).optional();
const path = z
.string()
.max(4096)
.describe(
'path relative to the Jupyter contents root; call colab_list_files to inspect it; not a local path',
);
const format = z.enum(['text', 'base64']).default('text');
const MIB = 1024 * 1024;
function result(value) {
const images = [];
const text = JSON.stringify(value, (key, data) => {
if (['image/png', 'image/jpeg', 'image/webp'].includes(key) && typeof data === 'string') {
images.push({ type: 'image', mimeType: key, data: data.replaceAll('\n', '') });
return '[returned as MCP image]';
}
return data;
});
if (Buffer.byteLength(text) > 2 * MIB)
throw new Error('tool result exceeded 2 MiB; request a smaller file or directory');
return {
content: [{ type: 'text', text }, ...images],
...(value.status === 'error' ? { isError: true } : {}),
};
}
export function createServer(colab) {
const server = new McpServer(
{ name: 'colab-mcp', version: '0.1.0' },
{
instructions:
'Use colab_accounts to discover opaque account ids. Choose an account_id explicitly for launches when multiple accounts exist; never blindly retry on another account. Runtime ids bind their owning account. Manage Google Colab runtimes. Launching consumes compute units; stop runtimes when finished. Every tool has readOnlyHint=true by user convention, NOT because operations are read-only. Execution and writes mutate state; stopping destroys runtime-local data.',
},
);
const tool = (name, description, inputSchema, run) => {
server.registerTool(
name,
{
description,
inputSchema: z.strictObject(inputSchema),
annotations: { readOnlyHint: true },
},
async (args, extra) => {
try {
return result(await run(args, extra.signal));
} catch (error) {
return { isError: true, content: [{ type: 'text', text: error.message }] };
}
},
);
};
tool(
'colab_accounts',
'List connected account ids without revealing emails. Check colab_usage for an account before choosing it. Discovery does not verify token health.',
{},
() => colab.accounts(),
);
tool(
'colab_list',
'List active runtimes across accounts, or filter by account_id. Runtime handles bind the owning account; preserve the whole id. Partial account failures are returned in errors. Proxy URLs and credentials are never returned.',
{ account_id },
({ account_id }, signal) => colab.list(signal, account_id),
);
tool(
'colab_usage',
'Get compute-unit balance and consumption for an account. Choose account_id explicitly when multiple accounts are connected; balances are not pooled.',
{ account_id },
({ account_id }, signal) => colab.usage(signal, account_id),
);
tool(
'colab_launch',
'Allocate a runtime (default T4 GPU). Consumes compute units. Save its id and stop it when done. Reuse notebook_id to reconcile an uncertain allocation; never blindly retry.',
{
account_id,
accelerator: z.enum(['NONE', 'T4', 'L4', 'G4', 'A100', 'H100', 'V5E1', 'V6E1']).default('T4'),
high_memory: z.boolean().default(false),
notebook_id: z
.string()
.regex(/^[A-Za-z0-9_.-]{44}$/)
.optional(),
},
(args, signal) => colab.launch(args, signal),
);
tool(
'colab_stop',
'Release a runtime. DESTRUCTIVE: all runtime-local files and Python state are lost.',
{ id },
({ id }, signal) => colab.stop(id, signal),
);
tool(
'colab_execute',
'Execute Python/IPython in a persistent kernel. Shell commands work via !command. Mutates runtime state. One execution per runtime at a time; output capped at 1 MiB. Timeout/cancellation attempts interrupt, never automatic re-execution.',
{
id,
code: z
.string()
.min(1)
.max(256 * 1024),
kernel_id,
timeout: z.number().min(1).max(3600).default(60).describe('execution timeout in seconds'),
},
(args, signal) => colab.execute(args, signal),
);
tool(
'colab_kernels',
'List kernel ids and execution state for a runtime. No kernel is created.',
{ id },
({ id }, signal) => colab.kernels(id, signal),
);
tool(
'colab_interrupt',
'Interrupt running code without deleting Python state or stopping the VM.',
{ id, kernel_id },
(args, signal) => colab.interrupt(args, signal),
);
tool(
'colab_list_files',
'List a runtime directory (up to 1000 entries). Paths are relative to the contents root.',
{
id,
path: path.default(''),
},
async (args, signal) => {
const data = await colab.files(args, signal);
if (data.type !== 'directory' || !Array.isArray(data.content))
throw new Error('path is not a directory');
return {
path: data.path,
entries: data.content.slice(0, 1000).map(({ name, path, type, size, last_modified }) => ({
name,
path,
type,
size,
last_modified,
})),
truncated: data.content.length > 1000,
};
},
);
tool(
'colab_read_file',
'Read a runtime file, up to 1 MiB of returned content. Use base64 for binary files. Use Python for larger files.',
{
id,
path,
format,
},
async (args, signal) => {
const data = await colab.files(args, signal);
if (data.type !== 'file')
throw new Error('path is not a regular file; use colab_execute for notebooks');
if (typeof data.content !== 'string')
throw new Error('file response did not contain text or base64');
if (Buffer.byteLength(data.content) > MIB)
throw new Error('file exceeds 1 MiB; use Python to read a slice');
return { path: data.path, format: data.format, content: data.content };
},
);
tool(
'colab_write_file',
'Create or OVERWRITE a runtime file. Parent directory must exist. Content capped at 1 MiB; use base64 for binary.',
{
id,
path,
format,
content: z.string().max(MIB),
},
(args, signal) => {
if (Buffer.byteLength(args.content) > MIB) throw new Error('content exceeds 1 MiB');
if (
args.format === 'base64' &&
!/^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/.test(args.content)
) {
throw new Error('invalid base64 content');
}
return colab.files(args, signal);
},
);
return server;
}