From 7ccee3a2f521d06954254b1a7d4e591893eed70a Mon Sep 17 00:00:00 2001 From: israel Date: Wed, 2 Sep 2026 16:53:05 +0100 Subject: [PATCH] fix(platform): revoke the gateway key on every session teardown A sandbox session mints an LLM gateway virtual key and hands its value into the container. revokeVirtualKey had one caller where 0.4 had four, so the TTL expiry sweep, admin destroy, phantom heal and the task-agent lane all tore sessions down and left the key live. The mint body states the invariant twice: the gateway has no native TTL, and reset_duration '1M' is safe BECAUSE teardown revokes the key long before any reset matters. It did not. So each abandoned key is a permanent bearer credential with a self-refilling monthly allowance against the org's own provider keys, one per expired, deadline-failed or destroyed session. All teardown edges now go through one seam, domains/sandbox/gateway-keys.ts. The token flip IS the election: UPDATE ... WHERE revoked_at_ms IS NULL RETURNING hands each live key to exactly one caller, so a watchdog sweeping twice or a destroy racing the expiry revokes once. An exec-scoped teardown narrows to that turn's key, so a sibling turn on the same standing session keeps its own credential. Best-effort per key, as 0.4 was: an unreachable gateway must never wedge a teardown. But the failure logs at console.error naming the key id, because that key stays spendable and only an operator can delete it by hand. --- .../domains/sandbox/gateway-keys.test.ts | 134 ++++++++ .../backend/domains/sandbox/gateway-keys.ts | 129 ++++++++ .../backend/domains/sandbox/sessions.ts | 16 +- .../backend/domains/sandbox/watchdogs.ts | 29 +- .../backend/domains/tasks/agent-runs.ts | 35 ++- .../platform/backend/integration-check.ts | 289 ++++++++++++++++++ 6 files changed, 621 insertions(+), 11 deletions(-) create mode 100644 services/platform/backend/domains/sandbox/gateway-keys.test.ts create mode 100644 services/platform/backend/domains/sandbox/gateway-keys.ts diff --git a/services/platform/backend/domains/sandbox/gateway-keys.test.ts b/services/platform/backend/domains/sandbox/gateway-keys.test.ts new file mode 100644 index 0000000000..ddb4773f46 --- /dev/null +++ b/services/platform/backend/domains/sandbox/gateway-keys.test.ts @@ -0,0 +1,134 @@ +/** + * The election itself is the SQL `WHERE revoked_at_ms IS NULL` flip, so it is + * the integration harness that proves it against a real Postgres. What these + * cover is the glue around it, where the consequences are just as sharp: + * + * - a gateway failure must NOT throw. An unreachable gateway wedging a + * teardown is worse than a leaked key, so the failure is counted and + * logged loudly instead. + * - an exec-scoped teardown must NOT clear the session row's parked key — + * that key belongs to the standing session, and a sibling turn is still + * spending against it. + * - the same key id arriving from both the token table and the session row + * must be revoked once. + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const { revokeVirtualKey } = vi.hoisted(() => ({ revokeVirtualKey: vi.fn() })); + +vi.mock('../../core/node_only/sandbox/llm_gateway_admin.ts', () => ({ + revokeVirtualKey, +})); + +const { revokeSessionGatewayKeys } = await import('./gateway-keys.ts'); + +/** + * A tagged-template stand-in for `postgres`: answers each query with the + * next queued result and records how many ran, so a test can assert that + * the session-row statement did not. + */ +function fakeSql(results: unknown[][]) { + const queries: string[] = []; + const sql = (strings: TemplateStringsArray) => { + queries.push(strings.join('?')); + return Promise.resolve(results[queries.length - 1] ?? []); + }; + return { sql: sql as never, queries }; +} + +const ARGS = { organizationId: 'org-1', sessionId: 'sess-1' }; + +describe('revokeSessionGatewayKeys', () => { + beforeEach(() => { + vi.clearAllMocks(); + revokeVirtualKey.mockResolvedValue(undefined); + }); + + it('revokes a claimed token key and the key parked on the session row', async () => { + const { sql } = fakeSql([[{ keyId: 'k1' }], [{ keyId: 'k2' }]]); + + const out = await revokeSessionGatewayKeys(sql, ARGS); + + expect(out).toEqual({ revoked: 2, failed: 0 }); + expect( + revokeVirtualKey.mock.calls + .map(([id]) => id as string) + .toSorted((a, b) => a.localeCompare(b)), + ).toEqual(['k1', 'k2']); + }); + + it('revokes a key claimed from both places only once', async () => { + const { sql } = fakeSql([[{ keyId: 'k1' }], [{ keyId: 'k1' }]]); + + const out = await revokeSessionGatewayKeys(sql, ARGS); + + expect(out).toEqual({ revoked: 1, failed: 0 }); + expect(revokeVirtualKey).toHaveBeenCalledTimes(1); + }); + + it('leaves the session row alone for an exec-scoped teardown', async () => { + const { sql, queries } = fakeSql([[{ keyId: 'k1' }]]); + + const out = await revokeSessionGatewayKeys(sql, { + ...ARGS, + execId: 'exec-9', + }); + + expect(out).toEqual({ revoked: 1, failed: 0 }); + // One statement only: the token claim. The `sandbox_sessions` clear must + // not run — a sibling turn on the standing session still holds that key. + expect(queries).toHaveLength(1); + expect(queries[0]).toContain('sandbox_session_tokens'); + expect(queries.join('\n')).not.toContain('UPDATE app.sandbox_sessions'); + }); + + it('skips the gateway entirely when nothing was claimed', async () => { + const { sql } = fakeSql([[], []]); + + const out = await revokeSessionGatewayKeys(sql, ARGS); + + expect(out).toEqual({ revoked: 0, failed: 0 }); + expect(revokeVirtualKey).not.toHaveBeenCalled(); + }); + + it('never throws when the gateway fails, and says the key leaked', async () => { + const error = vi.spyOn(console, 'error').mockImplementation(() => {}); + revokeVirtualKey.mockRejectedValueOnce(new Error('gateway down')); + const { sql } = fakeSql([[{ keyId: 'k1' }, { keyId: 'k2' }], []]); + + const out = await revokeSessionGatewayKeys(sql, ARGS); + + // Counted, not thrown: an unreachable gateway must not wedge teardown. + expect(out).toEqual({ revoked: 1, failed: 1 }); + // console.error, not warn — the key stays spendable until an operator + // deletes it by hand, so this has to be loud. + expect(error).toHaveBeenCalledTimes(1); + expect(String(error.mock.calls[0]?.[0])).toContain('LEAKED gateway key'); + error.mockRestore(); + }); + + it('keeps revoking the remaining keys after one fails', async () => { + const error = vi.spyOn(console, 'error').mockImplementation(() => {}); + revokeVirtualKey.mockRejectedValueOnce(new Error('gateway down')); + const { sql } = fakeSql([ + [{ keyId: 'k1' }, { keyId: 'k2' }, { keyId: 'k3' }], + [], + ]); + + const out = await revokeSessionGatewayKeys(sql, ARGS); + + expect(out).toEqual({ revoked: 2, failed: 1 }); + expect(revokeVirtualKey).toHaveBeenCalledTimes(3); + error.mockRestore(); + }); + + it('ignores token rows that minted no gateway key', async () => { + const { sql } = fakeSql([[{ keyId: null }, { keyId: 'k1' }], []]); + + const out = await revokeSessionGatewayKeys(sql, ARGS); + + expect(out).toEqual({ revoked: 1, failed: 0 }); + expect(revokeVirtualKey).toHaveBeenCalledExactlyOnceWith('k1'); + }); +}); diff --git a/services/platform/backend/domains/sandbox/gateway-keys.ts b/services/platform/backend/domains/sandbox/gateway-keys.ts new file mode 100644 index 0000000000..fc569c89cb --- /dev/null +++ b/services/platform/backend/domains/sandbox/gateway-keys.ts @@ -0,0 +1,129 @@ +import type { Sql } from 'postgres'; + +import { revokeVirtualKey } from '../../core/node_only/sandbox/llm_gateway_admin.ts'; + +/** + * Session-teardown credential reclaim — the 0.5 twin of 0.4's + * `node_only/sandbox/session_teardown.ts` + the revoke half of + * `session_admin_actions.destroySandbox`. + * + * WHY this exists as its own seam: the gateway has NO native TTL on a + * virtual key, and `mintVirtualKey` deliberately gives every key a + * `reset_duration: '1M'` budget window because teardown is supposed to + * delete it long before any reset matters. A key that outlives its session + * is therefore a permanent bearer credential with a self-refilling monthly + * allowance against the org's own provider keys. Every teardown edge — TTL + * expiry, admin destroy, phantom heal, a deadline-failed agent run — has to + * come through here, and there is exactly ONE copy of the behaviour so the + * edges cannot drift. + * + * THE ELECTION IS THE TOKEN FLIP. `UPDATE … WHERE revoked_at_ms IS NULL + * RETURNING llm_gateway_key_id` hands each live key to exactly one caller, + * so a watchdog sweeping the same session twice, or a destroy racing the + * expiry sweep, revokes once and the second pass is a no-op. (The gateway + * itself is idempotent too — `revokeVirtualKey` treats a 404 for an unknown + * key as success — so even a genuine double DELETE cannot fail a teardown.) + * + * FAILURE POSTURE: best-effort per key, like 0.4. A gateway that is + * unreachable must never wedge a teardown, so the HTTP failure is caught — + * but it is caught at `console.error` with the key id, because the key stays + * spendable and only an operator can delete it by hand. + */ + +/** Delete the given virtual keys on the gateway. Never throws: one key's + * failure is logged (loudly — that key stays spendable) and the rest run. */ +async function revokeGatewayKeys( + keyIds: string[], + context: string, +): Promise<{ revoked: number; failed: number }> { + let revoked = 0; + let failed = 0; + for (const keyId of keyIds) { + try { + await revokeVirtualKey(keyId); + revoked += 1; + } catch (error) { + failed += 1; + // NOT a warning: the key survives with a self-refilling monthly + // budget against this org's provider keys, and nothing else tracks + // it — an operator has to delete it on the gateway by hand. + console.error( + `[sandbox] LEAKED gateway key ${keyId} (${context}): revoke failed, the key stays spendable until an operator deletes it:`, + error, + ); + } + } + return { revoked, failed }; +} + +/** + * Claim and revoke the gateway virtual keys a torn-down session still holds. + * + * Scope: + * - no `execId` — the WHOLE session (TTL expiry, destroy, phantom heal): + * every unrevoked session token is marked revoked and every gateway key + * among them deleted, plus the key id parked on the session row itself + * (`app.sandbox_sessions.llm_gateway_key_id`, cleared here so a later + * sweep can tell a revoked key from a live one — nothing reads that + * column downstream, unlike the token table's copy, which the run + * provenance ledger matches turns by and which therefore keeps its id + * and carries `revoked_at_ms` as the mark). + * - with `execId` — ONE turn of a STANDING session (a deadline-failed + * task-agent run): only that exec's minted key, so a sibling turn still + * running on the same `pa-` session keeps its own credential. + */ +export async function revokeSessionGatewayKeys( + sql: Sql, + args: { organizationId: string; sessionId: string; execId?: string }, +): Promise<{ revoked: number; failed: number }> { + const execId = args.execId ?? null; + const now = Date.now(); + // The flip IS the claim — see the module note. An exec-scoped teardown + // narrows to the key that exec's op row minted; a token row with no + // gateway key (the subscription lane mints none) never matches it. + const claimed = await sql<{ keyId: string | null }[]>` + UPDATE app.sandbox_session_tokens SET revoked_at_ms = ${now} + WHERE session_id = ${args.sessionId} AND org_id = ${args.organizationId} + AND revoked_at_ms IS NULL + AND (${execId}::text IS NULL OR llm_gateway_key_id = ( + SELECT minted_key_id FROM app.sandbox_session_ops + WHERE session_id = ${args.sessionId} AND exec_id = ${execId} + LIMIT 1 + )) + RETURNING llm_gateway_key_id AS "keyId" + `; + const keyIds = new Set( + claimed + .map((row) => row.keyId) + .filter((keyId): keyId is string => keyId !== null), + ); + if (execId === null) { + // `RETURNING` answers with the NEW row, so the id has to be read before + // the clear — one statement, so the `FOR UPDATE` re-check makes this a + // claim too: a racing sweep sees the already-cleared row and skips it. + const parked = await sql<{ keyId: string }[]>` + WITH live AS ( + SELECT id, llm_gateway_key_id AS key_id + FROM app.sandbox_sessions + WHERE session_id = ${args.sessionId} + AND org_id = ${args.organizationId} + AND llm_gateway_key_id IS NOT NULL + FOR UPDATE + ), cleared AS ( + UPDATE app.sandbox_sessions SET llm_gateway_key_id = NULL + WHERE id IN (SELECT id FROM live) + RETURNING id + ) + SELECT key_id AS "keyId" FROM live + WHERE id IN (SELECT id FROM cleared) + `; + for (const row of parked) keyIds.add(row.keyId); + } + if (keyIds.size === 0) return { revoked: 0, failed: 0 }; + return revokeGatewayKeys( + [...keyIds], + execId === null + ? `session ${args.sessionId}` + : `session ${args.sessionId} exec ${execId}`, + ); +} diff --git a/services/platform/backend/domains/sandbox/sessions.ts b/services/platform/backend/domains/sandbox/sessions.ts index 5219890e69..efe55992fb 100644 --- a/services/platform/backend/domains/sandbox/sessions.ts +++ b/services/platform/backend/domains/sandbox/sessions.ts @@ -17,6 +17,7 @@ import { sessionIdForWorkflowExecution } from '../../core/sandbox/session_naming import { toJson } from '../../db/sql.ts'; import { readGovernancePolicyForOrg } from '../../lib/org-config.ts'; import { wakeParkedAgentRuns } from '../tasks/agent-runs.ts'; +import { revokeSessionGatewayKeys } from './gateway-keys.ts'; /** * The sandbox session substrate over PG — the 0.5 twin of @@ -455,11 +456,24 @@ export async function resumeSessionSlot( }); } -/** Terminal: mark destroyed + revoke tokens + drop the checkpoint + ticket. */ +/** Terminal: revoke the gateway keys + mark destroyed + revoke tokens + drop + * the checkpoint + ticket. The single bottom of EVERY destroy — the admin + * Destroy, the watchdog's phantom heal, `provisionSession`'s heal, the owner + * cascades — so credential reclaim cannot be missed on one of them. */ export async function markSessionDestroyed( sql: Sql, args: { organizationId: string; sessionId: string }, ): Promise { + // Credentials FIRST: the gateway key outlives the row (no native TTL), and + // the token flip below is what elects a single revoker — running it after + // the flip would find nothing to revoke. Best-effort by construction, so a + // down gateway cannot wedge the destroy; see `gateway-keys.ts`. + await revokeSessionGatewayKeys(sql, args).catch((error: unknown) => { + console.error( + `[sandbox] gateway key reclaim for destroyed ${args.sessionId} failed:`, + error, + ); + }); return sql .begin(async (tx) => { const now = Date.now(); diff --git a/services/platform/backend/domains/sandbox/watchdogs.ts b/services/platform/backend/domains/sandbox/watchdogs.ts index dd44fba9f6..65023d903c 100644 --- a/services/platform/backend/domains/sandbox/watchdogs.ts +++ b/services/platform/backend/domains/sandbox/watchdogs.ts @@ -1,6 +1,7 @@ import type { Sql } from 'postgres'; import { wakeParkedAgentRuns } from '../tasks/agent-runs.ts'; +import { revokeSessionGatewayKeys } from './gateway-keys.ts'; import { reconcileSession } from './service.ts'; import { reapStaleAdmissionTickets } from './sessions.ts'; @@ -10,9 +11,10 @@ import { reapStaleAdmissionTickets } from './sessions.ts'; * `recoverStuckAdmissionTickets`: * * - EXPIRE: unpinned sessions past their TTL among the compute-holding - * statuses flip to `expired` (freeing their slots) and the parked-run - * wake fires for their orgs. The spawner's own reaper collects the - * container on its TTL — the row must not wait for it. + * statuses flip to `expired` (freeing their slots), their gateway + * virtual keys are revoked, and the parked-run wake fires for their + * orgs. The spawner's own reaper collects the container on its TTL — + * the row must not wait for it. * - RECONCILE: a bounded batch of the stalest compute-holding rows is * checked against the spawner; a container gone spawner-side settles the * row as destroyed (phantom heal). Requires a reachable spawner — when @@ -29,12 +31,29 @@ export async function runSandboxWatchdog( } = {}, ): Promise<{ expired: number; healed: number; reaped: number }> { const now = Date.now(); - const expired = await sql<{ orgId: string }[]>` + const expired = await sql<{ orgId: string; sessionId: string }[]>` UPDATE app.sandbox_sessions SET status = 'expired' WHERE status IN ('creating', 'active', 'degraded') AND pinned = false AND expires_at_ms < ${now} - RETURNING org_id AS "orgId" + RETURNING org_id AS "orgId", session_id AS "sessionId" `; + // Reclaim the CREDENTIALS of every session this sweep just expired: the + // gateway key has no TTL of its own, so an expired row that keeps its key + // leaves a spendable credential live forever (0.4's + // `teardownExpiredSessions`). Deliberately NOT a container destroy — + // expiry is a lifetime cap, not a user Destroy, and the workspace is the + // user's state; the spawner's idle reaper collects the container. + for (const row of expired) { + await revokeSessionGatewayKeys(sql, { + organizationId: row.orgId, + sessionId: row.sessionId, + }).catch((error: unknown) => { + console.error( + `[watchdog] gateway key reclaim after expiring ${row.sessionId} failed:`, + error, + ); + }); + } for (const orgId of new Set(expired.map((row) => row.orgId))) { await wakeParkedAgentRuns(sql, orgId).catch((error: unknown) => { console.warn('[watchdog] capacity wake after expiry failed:', error); diff --git a/services/platform/backend/domains/tasks/agent-runs.ts b/services/platform/backend/domains/tasks/agent-runs.ts index ca0b0ceabe..606c047da2 100644 --- a/services/platform/backend/domains/tasks/agent-runs.ts +++ b/services/platform/backend/domains/tasks/agent-runs.ts @@ -4,6 +4,7 @@ import type { Sql, TransactionSql } from 'postgres'; import { AUTO_RETRY_MAX_ATTEMPTS } from '../../core/tasks/task_auto_retry.ts'; import { addJobInTx } from '../../jobs/enqueue.ts'; +import { revokeSessionGatewayKeys } from '../sandbox/gateway-keys.ts'; import { recordTaskAgentRunLedgerEntry } from './run-ledger.ts'; /** @@ -203,6 +204,14 @@ export async function settleAgentRun( }); } +/** + * Fail exactly once (the watchdog's deadline pass — the drive chain's own + * failures go through the shim's `markTaskAgentRunFailed`). The turn died + * without reaching `releaseTurnKey`, so its gateway key is reclaimed here: + * the winning flip IS the election, so the revoke fires once even when two + * sweeps race. Scoped to THIS exec — a sibling turn on the same standing + * `pa-` session keeps its own key. + */ export async function failAgentRun( sql: Sql, args: { @@ -214,17 +223,21 @@ export async function failAgentRun( }, ): Promise { const now = Date.now(); - return sql.begin(async (tx) => { - const rows = await tx<{ id: string }[]>` + const failed = await sql.begin(async (tx) => { + const rows = await tx<{ sessionId: string }[]>` UPDATE app.project_agent_runs SET status = 'failed', error = ${args.error.slice(0, 2000)}, api_error_status = ${args.apiErrorStatus ?? null}, settled_at_ms = ${now}, updated_at_ms = ${now} WHERE id = ${args.runId} AND org_id = ${args.organizationId} AND exec_id = ${args.execId} AND status IN ('queued', 'running') - RETURNING id + RETURNING session_id AS "sessionId" `; - if (rows.length === 0) return false; + const run = rows[0]; + if (run === undefined) return null; + // Inside the election's transaction: the provenance entry still reads + // the turn's token row by key id, which the revoke below leaves in + // place (it marks `revoked_at_ms`, it does not drop the id). await recordTaskAgentRunLedgerEntry(tx, { runId: args.runId, organizationId: args.organizationId, @@ -232,8 +245,20 @@ export async function failAgentRun( settledAt: now, error: args.error.slice(0, 2000), }); - return true; + return run.sessionId; + }); + if (failed === null) return false; + await revokeSessionGatewayKeys(sql, { + organizationId: args.organizationId, + sessionId: failed, + execId: args.execId, + }).catch((error: unknown) => { + console.error( + `[task-agent] gateway key reclaim for deadline-failed run ${args.runId} failed:`, + error, + ); }); + return true; } export async function cancelAgentRun( diff --git a/services/platform/backend/integration-check.ts b/services/platform/backend/integration-check.ts index 90377e8227..c63727ec32 100644 --- a/services/platform/backend/integration-check.ts +++ b/services/platform/backend/integration-check.ts @@ -33704,6 +33704,294 @@ async function checkOrganizationLifecycle( ); } +/** + * TEARDOWN CREDENTIAL RECLAIM — every edge that ends a sandbox session has to + * delete its gateway virtual key (fake bifrost records the DELETEs). + * + * The gateway has NO native TTL and `mintVirtualKey` gives every key a + * `reset_duration: '1M'` budget window precisely because teardown deletes it + * first, so a key that outlives its session is a permanent bearer credential + * with a self-refilling monthly allowance against the org's own provider + * keys. Three edges reach the same helper, each asserted on the key ids the + * gateway actually saw: + * + * - EXPIRY: the sandbox watchdog's TTL pass (0.4 `teardownExpiredSessions`). + * - DESTROY/HEAL: `markSessionDestroyed`, the bottom of the admin Destroy, + * the reconcile phantom heal and the owner cascades (0.4 + * `destroySandbox` + `teardownThreadSessions`) — driven TWICE, because + * the token flip is the election and a second teardown must revoke + * nothing rather than fail. + * - DEADLINE: the task-agent watchdog's overdue fail (0.4 `releaseTurnKey`, + * which a died drive chain never reaches) — scoped to the failed exec, so + * a sibling turn on the same standing session keeps its own key. + * + * Plus the FAILURE POSTURE: one key's revoke answers 503, and the teardown + * still settles the row (the leak is logged, never rethrown) — a gateway + * outage must not leave sessions undestroyable. + */ +async function checkSandboxGatewayKeyReclaim( + sql: Sql, + base: string, + ctx: { cookie: string; orgId: string }, +): Promise { + const { cookie, orgId } = ctx; + const { createServer } = await import('node:http'); + const deleted: string[] = []; + const attempted: string[] = []; + const gateway = createServer((req, res) => { + const url = req.url ?? ''; + res.setHeader('content-type', 'application/json'); + if ( + req.method === 'DELETE' && + url.startsWith('/api/governance/virtual-keys/') + ) { + const keyId = decodeURIComponent(url.split('/').at(-1) ?? ''); + attempted.push(keyId); + if (keyId === 'vk-gk-boom') { + // The one key whose revoke fails — the posture lane below. + res.statusCode = 503; + res.end('{}'); + return; + } + deleted.push(keyId); + res.end('{}'); + return; + } + res.statusCode = 404; + res.end('{}'); + }); + await new Promise((resolve) => { + gateway.listen(0, '127.0.0.1', resolve); + }); + const address = gateway.address(); + const port = + address !== null && typeof address === 'object' ? address.port : 0; + const previousGatewayUrl = process.env.SANDBOX_LLM_GATEWAY_URL; + const previousGatewayPassword = + process.env.SANDBOX_LLM_GATEWAY_ADMIN_PASSWORD; + process.env.SANDBOX_LLM_GATEWAY_URL = `http://127.0.0.1:${port}`; + // The admin client refuses to run anonymous — deliberately, since the + // management API is reachable from every sandbox session. The fake gateway + // below does not check the value, but without one set every revoke fails + // authentication and the check would pass its fail-open lane while proving + // nothing about the revoke itself. + process.env.SANDBOX_LLM_GATEWAY_ADMIN_PASSWORD ??= 'itest-gateway-admin'; + const now = Date.now(); + const seedSession = async ( + sessionId: string, + args: { expiresAt: number; keyId?: string; ownerType?: string }, + ): Promise => { + await sql` + INSERT INTO app.sandbox_sessions ( + org_id, session_id, status, owner_type, owner_id, created_by, + llm_gateway_key_id, created_at_ms, expires_at_ms + ) VALUES ( + ${orgId}, ${sessionId}, 'active', ${args.ownerType ?? 'render'}, + ${sessionId}, 'itest:gk', ${args.keyId ?? null}, ${now}, + ${args.expiresAt} + ) + `; + }; + const seedToken = async ( + sessionId: string, + tokenHash: string, + keyId: string, + ): Promise => { + await sql` + INSERT INTO app.sandbox_session_tokens ( + org_id, session_id, token_hash, llm_gateway_key_id, scope, + created_at_ms, expires_at_ms + ) VALUES ( + ${orgId}, ${sessionId}, ${tokenHash}, ${keyId}, + ${sql.json({ agentKind: 'claude-code', allowedModels: [], connectorGrants: [], budgetCents: 100 })}, + ${now}, ${now + 3_600_000} + ) + `; + }; + const tokenRevoked = async (tokenHash: string): Promise => { + const rows = await sql<{ revokedAt: number | null }[]>` + SELECT revoked_at_ms::float8 AS "revokedAt" + FROM app.sandbox_session_tokens WHERE token_hash = ${tokenHash} + `; + return rows[0]?.revokedAt != null; + }; + const sessionKeyId = async (sessionId: string): Promise => { + const rows = await sql<{ keyId: string | null }[]>` + SELECT llm_gateway_key_id AS "keyId" FROM app.sandbox_sessions + WHERE session_id = ${sessionId} AND org_id = ${orgId} + `; + return rows[0]?.keyId ?? null; + }; + const count = (keyId: string): number => + deleted.filter((id) => id === keyId).length; + + try { + // --- edge 1: the TTL expiry sweep ------------------------------------ + await seedSession('gk-ttl-gone', { + expiresAt: now - 3_600_000, + keyId: 'vk-gk-row', + }); + await seedToken('gk-ttl-gone', 'gk-hash-ttl', 'vk-gk-token'); + await seedSession('gk-ttl-live', { + expiresAt: now + 3_600_000, + keyId: 'vk-gk-live', + }); + await seedToken('gk-ttl-live', 'gk-hash-live', 'vk-gk-live-token'); + const sandboxWatchdogs = await import('./domains/sandbox/watchdogs.ts'); + await sandboxWatchdogs.runSandboxWatchdog(sql, { skipReconcile: true }); + const expiryOk = + count('vk-gk-row') === 1 && + count('vk-gk-token') === 1 && + (await sessionKeyId('gk-ttl-gone')) === null && + (await tokenRevoked('gk-hash-ttl')) && + count('vk-gk-live') === 0 && + count('vk-gk-live-token') === 0 && + (await sessionKeyId('gk-ttl-live')) === 'vk-gk-live' && + !(await tokenRevoked('gk-hash-live')); + + // --- edge 2: destroy / phantom heal, driven twice -------------------- + const sessions = await import('./domains/sandbox/sessions.ts'); + await seedSession('gk-destroy', { + expiresAt: now + 3_600_000, + keyId: 'vk-gk-destroy-row', + }); + await seedToken('gk-destroy', 'gk-hash-destroy', 'vk-gk-destroy'); + const destroyedFirst = await sessions.markSessionDestroyed(sql, { + organizationId: orgId, + sessionId: 'gk-destroy', + }); + const destroyedTwice = await sessions.markSessionDestroyed(sql, { + organizationId: orgId, + sessionId: 'gk-destroy', + }); + const destroyOk = + destroyedFirst && + !destroyedTwice && + count('vk-gk-destroy') === 1 && + count('vk-gk-destroy-row') === 1 && + (await sessionKeyId('gk-destroy')) === null && + (await tokenRevoked('gk-hash-destroy')); + + // --- edge 3: the task-agent deadline fail (exec-scoped) -------------- + const post = (route: string, body?: unknown): Promise => + fetch(`${base}${route}`, { + method: 'POST', + headers: { 'content-type': 'application/json', cookie, origin: base }, + ...(body !== undefined ? { body: JSON.stringify(body) } : {}), + }); + const project = z.object({ projectId: z.string() }).safeParse( + await ( + await post(`/api/app/projects?orgId=${orgId}`, { + name: 'Gateway keys', + }) + ).json(), + ); + const projectId = project.success ? project.data.projectId : ''; + const task = z.object({ taskId: z.string() }).safeParse( + await ( + await post(`/api/app/tasks?orgId=${orgId}`, { + projectId, + title: 'Gateway keys', + }) + ).json(), + ); + const taskId = task.success ? task.data.taskId : ''; + await seedSession('pa-gk-agent', { + expiresAt: now + 24 * 3_600_000, + ownerType: 'project_agent', + }); + await sql` + INSERT INTO app.project_agent_runs ( + org_id, project_id, task_id, agent_id, exec_id, session_id, status, + harness, model, started_by, started_at_ms, launched_at_ms, + deadline_at_ms, updated_at_ms + ) VALUES ( + ${orgId}, ${projectId}, ${taskId}, 'gk-agent', 'exec-gk-overdue', + 'pa-gk-agent', 'running', 'claude-code', 'itest-model', 'itest:gk', + ${now - 13 * 3_600_000}, ${now - 13 * 3_600_000}, ${now - 3_600_000}, + ${now} + ) + `; + for (const [execId, keyId] of [ + ['exec-gk-overdue', 'vk-gk-run'], + ['exec-gk-sibling', 'vk-gk-sibling'], + ]) { + await sql` + INSERT INTO app.sandbox_session_ops ( + org_id, session_id, exec_id, kind, status, minted_key_id, + heartbeat_at_ms, started_at_ms + ) VALUES ( + ${orgId}, 'pa-gk-agent', ${execId ?? ''}, 'agent-run', 'running', + ${keyId ?? ''}, ${now}, ${now} + ) + `; + await seedToken('pa-gk-agent', `gk-hash-${execId ?? ''}`, keyId ?? ''); + } + const taskWatchdogs = await import('./domains/tasks/watchdogs.ts'); + await taskWatchdogs.runTaskAgentWatchdog(sql); + const deadlineOk = + count('vk-gk-run') === 1 && + (await tokenRevoked('gk-hash-exec-gk-overdue')) && + count('vk-gk-sibling') === 0 && + !(await tokenRevoked('gk-hash-exec-gk-sibling')); + + // --- posture: a failing gateway must not wedge the teardown ---------- + await seedSession('gk-boom', { expiresAt: now + 3_600_000 }); + await seedToken('gk-boom', 'gk-hash-boom', 'vk-gk-boom'); + const boomDestroyed = await sessions.markSessionDestroyed(sql, { + organizationId: orgId, + sessionId: 'gk-boom', + }); + const boomStatus = await sql<{ status: string }[]>` + SELECT status FROM app.sandbox_sessions + WHERE session_id = 'gk-boom' AND org_id = ${orgId} + `; + const postureOk = + boomDestroyed && + boomStatus[0]?.status === 'destroyed' && + attempted.includes('vk-gk-boom') && + !deleted.includes('vk-gk-boom'); + + record( + 'sandbox teardown revokes the session gateway key (expiry/destroy/deadline)', + expiryOk && destroyOk && deadlineOk && postureOk, + `expiry(gone=${count('vk-gk-row')}/${count('vk-gk-token')} col=${await sessionKeyId('gk-ttl-gone')} live=${count('vk-gk-live')}/${count('vk-gk-live-token')}), destroy(first=${destroyedFirst} again=${destroyedTwice} keys=${count('vk-gk-destroy')}/${count('vk-gk-destroy-row')}), deadline(run=${count('vk-gk-run')} sibling=${count('vk-gk-sibling')}), posture(destroyed=${boomDestroyed} status=${boomStatus[0]?.status} attempted=${attempted.includes('vk-gk-boom')}), deleted=${deleted.join(',')}`, + ); + // Hand the seeded slots back: the sibling op stays `running` on purpose + // (that is what keeps its key alive), and the watchdog's running-op + // guard therefore leaves the standing session holding a project slot — + // which the org cap would charge to whatever check runs next. + await sql` + UPDATE app.sandbox_session_ops SET + status = 'cancelled', finished_at_ms = ${Date.now()}, + finalized_at_ms = coalesce(finalized_at_ms, ${Date.now()}) + WHERE session_id = 'pa-gk-agent' AND status = 'running' + `; + await sql` + UPDATE app.sandbox_sessions SET + status = 'destroyed', destroyed_at_ms = ${Date.now()} + WHERE org_id = ${orgId} + AND session_id IN ('gk-ttl-gone', 'gk-ttl-live', 'pa-gk-agent', + 'gk-boom') + `; + } finally { + if (previousGatewayUrl === undefined) { + delete process.env.SANDBOX_LLM_GATEWAY_URL; + if (previousGatewayPassword === undefined) { + delete process.env.SANDBOX_LLM_GATEWAY_ADMIN_PASSWORD; + } else { + process.env.SANDBOX_LLM_GATEWAY_ADMIN_PASSWORD = + previousGatewayPassword; + } + } else { + process.env.SANDBOX_LLM_GATEWAY_URL = previousGatewayUrl; + } + await new Promise((resolve) => { + gateway.close(() => resolve()); + }); + } +} + async function main(): Promise { const databaseUrl = process.env.DATABASE_URL; if (!databaseUrl) { @@ -33989,6 +34277,7 @@ async function main(): Promise { await checkBackfillRecovery(sql, authCtx); await checkWatchdogs(sql, baseUrl, authCtx); await checkDocumentWriteGuards(sql, baseUrl, authCtx); + await checkSandboxGatewayKeyReclaim(sql, baseUrl, authCtx); await checkLoginThrottleAndAuditChain( sql, baseUrl,