import { ACCOUNT_ID, readStore } from './account-store.js';
import { Auth, statKey } from './auth.js';
import { Colab } from './colab.js';
export class Accounts {
#path;
#env;
#fetch;
#createClient;
#single;
#clients = new Map();
#identityPattern = null;
#key;
#pooled = false;
#executing = 0;
constructor({ env = process.env, fetch = globalThis.fetch, single, createClient } = {}) {
this.#createClient = createClient ?? ((auth) => new Colab({ auth, fetch }));
this.#env = env;
this.#fetch = fetch;
this.#path = env.COLAB_MANAGED_CREDENTIALS && `${env.COLAB_MANAGED_CREDENTIALS}.accounts.json`;
this.#single = single ?? new Colab({ auth: new Auth({ env, fetch }), fetch });
}
#newClient(id) {
const env = {
...this.#env,
COLAB_ACCESS_TOKEN: '',
COLAB_CREDENTIALS: '',
COLAB_MANAGED_CREDENTIALS: this.#path,
};
return this.#createClient(new Auth({ env, fetch: this.#fetch, accountId: id }), id);
}
async #snapshot() {
const key = this.#path && (await statKey(this.#path));
if (!key) {
if (this.#pooled)
throw new Error('colab account store is missing; restore it through the gateway');
return new Map([['default', this.#single]]);
}
if (key !== this.#key) {
let records;
try {
records = await readStore(this.#path);
} catch {
throw new Error('cannot read colab account store; check the owner management page');
}
const next = new Map();
for (const record of records) {
const id = record.account_id;
next.set(id, this.#clients.get(id) ?? this.#newClient(id));
}
for (const [id, client] of this.#clients) if (!next.has(id)) client.close();
if (!this.#pooled) this.#single.close();
this.#clients = next;
const identities = records
.flatMap((record) => [record.email, record.google_subject])
.filter((value) => typeof value === 'string' && value.length > 0)
.map((value) => value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'));
this.#identityPattern = identities.length ? new RegExp(identities.join('|'), 'gi') : null;
this.#key = key;
this.#pooled = true;
}
return this.#clients;
}
#redact(value) {
return JSON.parse(
JSON.stringify(value, (key, data) => {
if (/^(email|email_address|emailAddress|google_subject|sub|userId|user_id)$/i.test(key))
return undefined;
if (typeof data !== 'string') return data;
if (this.#identityPattern) data = data.replace(this.#identityPattern, '[private account]');
return data;
}),
);
}
async #run(client, method, ...args) {
try {
return await client[method](...args);
} catch (error) {
throw new Error(this.#redact(error.message));
}
}
#select(clients, accountId) {
if (accountId === undefined && clients.size === 1) return [...clients][0];
if (accountId !== undefined && clients.has(accountId))
return [accountId, clients.get(accountId)];
if (clients.size === 0)
throw new Error('no connected colab accounts; use the owner management page');
throw new Error('choose an account_id from colab_accounts; no automatic account selection');
}
#publicRuntime(runtime, accountId) {
return {
...this.#redact(runtime),
id: accountId === 'default' ? runtime.id : `${accountId}:${runtime.id}`,
account_id: accountId,
};
}
async accounts() {
return { accounts: [...(await this.#snapshot()).keys()].map((account_id) => ({ account_id })) };
}
async list(signal, accountId) {
const clients = await this.#snapshot();
const selected = accountId === undefined ? [...clients] : [this.#select(clients, accountId)];
const runtimes = [],
errors = [];
// bound simultaneous control-plane responses, including their JSON buffers.
for (let i = 0; i < selected.length; i += 4) {
signal?.throwIfAborted();
await Promise.all(
selected.slice(i, i + 4).map(async ([id, client]) => {
try {
const listed = await this.#run(client, 'list', signal);
runtimes.push(...listed.runtimes.map((runtime) => this.#publicRuntime(runtime, id)));
} catch (error) {
if (selected.length === 1) throw error;
errors.push({ account_id: id, error: error.message });
}
}),
);
}
return { runtimes, ...(errors.length ? { errors } : {}) };
}
async usage(signal, accountId) {
const [id, client] = this.#select(await this.#snapshot(), accountId);
return { ...this.#redact(await this.#run(client, 'usage', signal)), account_id: id };
}
async launch({ account_id, ...options }, signal) {
const [id, client] = this.#select(await this.#snapshot(), account_id);
return this.#publicRuntime(await this.#run(client, 'launch', options, signal), id);
}
async #runtime(method, args, signal) {
const clients = await this.#snapshot();
let accountId = 'default',
endpoint = args.id;
if (this.#pooled || args.id.startsWith('acct_')) {
const separator = args.id.indexOf(':');
accountId = args.id.slice(0, separator);
endpoint = args.id.slice(separator + 1);
if (!ACCOUNT_ID.test(accountId) || !endpoint)
throw new Error('invalid runtime handle; call colab_list');
}
const client = clients.get(accountId);
if (!client)
throw new Error('unknown runtime account; call colab_list; accounts are never substituted');
const argument = ['stop', 'kernels'].includes(method) ? endpoint : { ...args, id: endpoint };
const result = await this.#run(client, method, argument, signal);
return result.id === undefined ? result : { ...result, id: args.id };
}
stop(id, signal) {
return this.#runtime('stop', { id }, signal);
}
kernels(id, signal) {
return this.#runtime('kernels', { id }, signal);
}
files(args, signal) {
return this.#runtime('files', args, signal);
}
interrupt(args, signal) {
return this.#runtime('interrupt', args, signal);
}
async execute(args, signal) {
if (this.#executing >= 8)
throw new Error('at most 8 concurrent runtime executions are allowed');
this.#executing++;
try {
return await this.#runtime('execute', args, signal);
} finally {
this.#executing--;
}
}
close() {
this.#single.close();
for (const client of this.#clients.values()) client.close();
}
}