import assert from 'node:assert/strict';
import test from 'node:test';
import { randomBytes } from 'node:crypto';
import { mkdtemp, writeFile, rename, readFile, rm, chmod } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { Client } from '@modelcontextprotocol/sdk/client/index.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, RUNTIME, startFake } from './fixtures/fake-colab.js';
async function setup(t) {
const dir = await mkdtemp(join(tmpdir(), 'colab-accounts-'));
const path = join(dir, 'google.json.accounts.json');
const records = [0, 1].map((i) => ({
account_id: `acct_${randomBytes(32).toString('base64url')}`,
google_subject: `google-subject-${i}`,
email: `private+${i}@example.com`,
type: 'authorized_user',
client_id: `client-${i}`,
client_secret: `secret-${i}`,
refresh_token: `refresh-${i}`,
token: `access-${i}`,
expiry: new Date(Date.now() + 3600000).toISOString(),
}));
const save = async (accounts = records) => {
await writeFile(`${path}.tmp`, JSON.stringify({ version: 1, accounts }), { mode: 0o600 });
await rename(`${path}.tmp`, path);
};
await save();
const fakes = await Promise.all([startFake(), startFake()]);
const tokens = [[], []];
for (const [i, fake] of fakes.entries()) {
fake.routes = (request) => {
if (request.path.startsWith('/tun/m/')) {
tokens[i].push(request.headers.authorization);
assert.equal(request.headers.authorization, `Bearer ${records[i].token}`);
request.headers.authorization = `Bearer ${ACCESS}`;
}
return false;
};
}
let fetchImpl = globalThis.fetch;
const options = {
env: {
COLAB_MANAGED_CREDENTIALS: join(dir, 'google.json'),
COLAB_ACCESS_TOKEN: 'wrong-global-token',
COLAB_CLIENT_ID: 'wrong-client',
},
fetch: (...args) => fetchImpl(...args),
createClient: (auth, id) =>
new Colab({
auth,
baseUrl: fakes[records.findIndex((r) => r.account_id === id)].url,
allowLocalProxy: true,
fetch: (...args) => fetchImpl(...args),
}),
};
const accounts = new Accounts(options);
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
const client = new Client({ name: 'pool-test', version: '0' });
const server = createServer(accounts);
await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]);
const call = async (name, args = {}) => {
const result = await client.callTool({ name, arguments: args });
return {
error: Boolean(result.isError),
text: result.content[0].text,
json: () => JSON.parse(result.content[0].text),
};
};
t.after(async () => {
accounts.close();
await client.close();
await Promise.all(fakes.map((fake) => fake.close()));
await rm(dir, { recursive: true, force: true });
});
return {
records,
fakes,
tokens,
accounts,
options,
call,
save,
path,
setFetch: (value) => {
fetchImpl = value;
},
};
}
test('MCP discovery hides identities; explicit choices route to separate Google accounts', async (t) => {
const { records, fakes, call, tokens } = await setup(t);
const discovered = await call('colab_accounts');
assert.deepEqual(discovered.json(), {
accounts: records.map(({ account_id }) => ({ account_id })),
});
for (const record of records)
for (const secret of [
record.email,
record.google_subject,
record.refresh_token,
record.client_secret,
])
assert.ok(!discovered.text.includes(secret));
assert.equal((await call('colab_launch')).error, true);
assert.equal((await call('colab_usage')).error, true);
assert.equal(fakes.flatMap((fake) => fake.requests).length, 0);
const id = records[1].account_id;
const launched = await call('colab_launch', { account_id: id, accelerator: 'T4' });
assert.equal(launched.error, false);
assert.equal(launched.json().id, `${id}:${RUNTIME}`);
assert.equal(launched.json().account_id, id);
const usage = await call('colab_usage', { account_id: id });
assert.equal(usage.json().currentBalance, 42);
assert.equal(usage.json().account_id, id);
assert.equal(tokens[0].length, 0);
assert.equal(tokens[1].length, 3);
});
test('identical provider runtime ids stay pinned for execute, files, interrupt, kernels and stop', async (t) => {
const { records, call, fakes, tokens } = await setup(t);
const listed = (await call('colab_list')).json();
assert.equal(listed.runtimes.length, 2);
assert.equal(new Set(listed.runtimes.map((runtime) => runtime.id)).size, 2);
const id = `${records[1].account_id}:${RUNTIME}`;
const firstCalls = tokens[0].length;
for (const [name, args] of [
['colab_execute', { id, code: 'hello' }],
['colab_write_file', { id, path: 'test.txt', content: 'hello' }],
['colab_read_file', { id, path: 'test.txt' }],
['colab_list_files', { id }],
['colab_kernels', { id }],
['colab_interrupt', { id, kernel_id: 'k1' }],
['colab_stop', { id }],
])
assert.equal((await call(name, args)).error, false, name);
assert.equal(tokens[0].length, firstCalls);
assert.equal(fakes[0].kernels.length, 0);
assert.equal(fakes[1].kernels.length, 1);
assert.equal(fakes[1].files.get('test.txt').data.toString(), 'hello');
});
test('restart preserves ids; reconnect replaces only that account and reloads tokens', async (t) => {
const { accounts, options, records, save, call, tokens, path } = await setup(t);
const before = await readFile(path);
const restarted = new Accounts(options);
t.after(() => restarted.close());
assert.deepEqual(await restarted.accounts(), await accounts.accounts());
records[0].token = 'reconnected-access';
await save();
assert.equal((await call('colab_usage', { account_id: records[0].account_id })).error, false);
assert.equal(tokens[0].at(-1), 'Bearer reconnected-access');
assert.notDeepEqual(await readFile(path), before);
const saved = await readFile(path);
await call('colab_usage', { account_id: records[1].account_id });
assert.deepEqual(await readFile(path), saved, 'worker does not rewrite the store');
});
test('refresh is account-local, client-pinned and read-only even under concurrent requests', async (t) => {
const { records, call, save, setFetch, path } = await setup(t);
for (const r of records) r.expiry = new Date(0).toISOString();
await save();
const before = await readFile(path);
const refreshes = [];
setFetch(async (url, options) => {
if (String(url) !== 'https://oauth2.googleapis.com/token') return fetch(url, options);
const params = new URLSearchParams(options.body);
const index = records.findIndex(
(record) => record.refresh_token === params.get('refresh_token'),
);
refreshes.push(index);
assert.equal(params.get('client_id'), records[index].client_id);
assert.equal(params.get('client_secret'), records[index].client_secret);
return Response.json({ access_token: records[index].token, expires_in: 3600 });
});
const results = await Promise.all(
[0, 0, 1, 1].map((i) => call('colab_usage', { account_id: records[i].account_id })),
);
assert.ok(results.every((result) => !result.error));
assert.deepEqual(refreshes.sort(), [0, 1]);
assert.deepEqual(await readFile(path), before);
});
test('a failed account does not hide others or trigger launch failover; quota identity fields are private', async (t) => {
const { records, fakes, tokens, call } = await setup(t);
fakes[0].routes = (_, response, json) => json(response, 403, { error: records[0].email });
const listed = await call('colab_list');
assert.equal(listed.error, false);
assert.equal(listed.json().runtimes.length, 1);
assert.equal(listed.json().errors[0].account_id, records[0].account_id);
assert.ok(!listed.text.includes(records[0].email));
const otherCalls = tokens[1].length;
const failed = await call('colab_launch', { account_id: records[0].account_id });
assert.equal(failed.error, true);
assert.match(failed.text, /notebook_id=/);
assert.equal(tokens[1].length, otherCalls);
fakes[1].routes = (request, response, json) =>
request.path.endsWith('/ccu-info') &&
json(response, 200, {
currentBalance: 42,
email: records[1].email,
userId: records[1].google_subject,
message: `for ${records[1].email.toUpperCase()}`,
});
const usage = await call('colab_usage', { account_id: records[1].account_id });
assert.equal(usage.error, false);
assert.equal(usage.json().email, undefined);
assert.ok(!usage.text.toLowerCase().includes(records[1].email.toLowerCase()));
assert.ok(!usage.text.includes(records[1].google_subject));
});
test('removed accounts never retarget handles; an empty store never resurrects legacy credentials', async (t) => {
const { records, accounts, options, call, save, fakes } = await setup(t);
await call('colab_accounts');
const stale = `${records[0].account_id}:${RUNTIME}`;
await save([records[1]]);
const before = fakes.flatMap((f) => f.requests).length;
assert.equal((await call('colab_stop', { id: stale })).error, true);
assert.equal((await call('colab_stop', { id: RUNTIME })).error, true);
assert.equal(fakes.flatMap((f) => f.requests).length, before);
await save([]);
assert.deepEqual(await accounts.accounts(), { accounts: [] });
assert.equal((await call('colab_launch')).error, true);
const restarted = new Accounts(options);
t.after(() => restarted.close());
assert.deepEqual(await restarted.accounts(), { accounts: [] });
});
test('bad, oversized, non-private or missing-after-load stores fail closed without fallback', async (t) => {
const { path, call, save } = await setup(t);
await call('colab_accounts');
for (const data of [
'{',
'{}',
JSON.stringify({ version: 9, accounts: [] }),
JSON.stringify({ version: 1, accounts: [null] }),
'x'.repeat(1024 * 1024 + 1),
JSON.stringify({ version: 1, accounts: [] }) + ' '.repeat(1024 * 1024),
]) {
await writeFile(path, data);
assert.equal((await call('colab_accounts')).error, true);
assert.equal((await call('colab_launch')).error, true);
}
await save();
await chmod(path, 0o644);
assert.equal((await call('colab_accounts')).error, true);
await chmod(path, 0o600);
await rm(path);
assert.equal((await call('colab_accounts')).error, true);
});
test('unknown bound handles cannot fall through to the legacy default account', async () => {
let called = false;
const accounts = new Accounts({
env: {},
single: {
stop: () => {
called = true;
},
close() {},
},
});
await assert.rejects(
accounts.stop(`acct_${'a'.repeat(43)}:${RUNTIME}`),
/unknown runtime account/,
);
assert.equal(called, false);
});
test('pool rejects duplicate ids and too many accounts without network requests', async (t) => {
const { records, path, call, fakes } = await setup(t);
for (const accounts of [
[records[0], records[0]],
Array.from({ length: 17 }, () => ({
...records[0],
account_id: `acct_${randomBytes(32).toString('base64url')}`,
})),
]) {
await writeFile(path, JSON.stringify({ version: 1, accounts }));
assert.equal((await call('colab_accounts')).error, true);
}
assert.equal(fakes.flatMap((fake) => fake.requests).length, 0);
});
test('execution limit remains global rather than multiplying by account count', async (t) => {
const { records, accounts, fakes } = await setup(t);
// Keep eight distinct runtime executions pending across both accounts.
for (const fake of fakes) {
fake.assignments = Array.from({ length: 5 }, (_, i) => `runtime-${i}`);
fake.execute = () => {};
}
const controller = new AbortController();
const pending = Array.from({ length: 8 }, (_, i) =>
accounts.execute(
{
id: `${records[i % 2].account_id}:runtime-${Math.floor(i / 2)}`,
code: 'wait',
timeout: 60,
},
controller.signal,
),
);
await assert.rejects(
accounts.execute({ id: `${records[0].account_id}:runtime-4`, code: 'ninth' }),
/at most 8/,
);
controller.abort();
await Promise.allSettled(pending);
});
test('uncertain launch stays with its chosen account and preserves reconciliation id', async (t) => {
const { records, fakes, call, setFetch } = await setup(t);
setFetch(async (url, options) => {
if (String(url).startsWith(fakes[0].url) && options.method === 'POST')
throw new TypeError('connection lost');
return fetch(url, options);
});
const notebook_id = 'a'.repeat(44);
const result = await call('colab_launch', { account_id: records[0].account_id, notebook_id });
assert.equal(result.error, true);
assert.match(result.text, /mutation may still have completed/);
assert.ok(result.text.includes(notebook_id));
assert.equal(fakes[1].requests.length, 0);
});