From 322d7809a2f7c8bad6ad85d8a58e928db2689ccd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 8 Sep 2026 22:35:16 -0400 Subject: [PATCH 1/7] feat: Add Bounded Concurrent Native Workspace Leases --- packages/code/README.md | 48 ++ packages/code/src/cli-slots.test.ts | 63 ++ packages/code/src/cli.ts | 309 +++++-- packages/code/src/native-pool.ts | 117 +++ packages/code/src/protocol.ts | 45 +- packages/code/src/worker-slots.test.ts | 73 ++ packages/code/src/worker.ts | 437 ++++++++-- packages/code/src/workspace-guards.ts | 49 ++ service/src/bridge/admission.ts | 27 +- service/src/bridge/concurrent-store.test.ts | 135 +++ service/src/bridge/concurrent-worker.test.ts | 201 +++++ service/src/bridge/index.ts | 7 +- service/src/bridge/router.ts | 861 ++++++++++--------- service/src/bridge/slots.test.ts | 83 ++ service/src/bridge/slots.ts | 133 +++ service/src/bridge/store.ts | 582 +++++++++---- service/src/config.ts | 313 +++++-- 17 files changed, 2645 insertions(+), 838 deletions(-) create mode 100644 packages/code/src/cli-slots.test.ts create mode 100644 packages/code/src/native-pool.ts create mode 100644 packages/code/src/worker-slots.test.ts create mode 100644 packages/code/src/workspace-guards.ts create mode 100644 service/src/bridge/concurrent-store.test.ts create mode 100644 service/src/bridge/concurrent-worker.test.ts create mode 100644 service/src/bridge/slots.test.ts create mode 100644 service/src/bridge/slots.ts diff --git a/packages/code/README.md b/packages/code/README.md index 7e7c0aaf..8ccfc5de 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -520,3 +520,51 @@ with `librechat-code reset-workspace `. The command uses the configured worker credentials, registers a fresh incarnation, and only clears the server fence when no assignment is active. Run it while the normal worker process is stopped, then restart the normal worker after the command exits. + +### Opt-in concurrent native workspaces + +Code API defaults to **one execution slot**. To allow independent native roots +to execute concurrently, configure `CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS=2` +on every Code API replica and start an updated worker with: + +```sh +librechat-code run \ + --worker-dir /projects/first \ + --workspace second=/projects/second \ + --workspace-lease-slots 2 \ + --allow-workspace-writes \ + --allow-workspace-commands +``` + +Keep the existing URL, pairing/identity, and network policy configuration. +The primary root keeps its configured workspace ID (default `primary`). Repeat +`--workspace id=path` to add named roots, up to the protocol's 32-root limit. +Roots must already exist and must not overlap or alias one another. Commands +retain the selected root's sandbox boundary, not a shared parent-directory grant. + +`LIBRECHAT_CODE_WORKSPACE_LEASE_SLOTS` is the equivalent worker setting. Both +ceilings must be integers from 1 to 8; the lower ceiling wins. An older Code API +without the negotiation receipt keeps the worker on the serial protocol. Deploy +the updated API to all replicas before enabling slots on workers. A capacity +change while work is active fails closed; stop and drain the worker before +changing it. + +Different roots can run concurrently; requests targeting the **same root remain +serialized**, even across chats or agents. This is root-level exclusion, not +file-level locking. Assign separate project/worktree roots for independent work. +The admission queue remains bounded at 32 requests per worker. An idle SRT process +cache is bounded by the local slot setting and evicts only idle executors. Runtime +sandbox assignments continue through the exclusive legacy lane; this does not +enable concurrent Docker/NsJail sessions or bypass any approval/network policy. + +An uncertain mutation or executor failure leaves an assignment-owned local guard +and a server-side fence for that root. Healthy roots can continue. The worker +does not replay the failed command. To recover a quarantined native root: + +1. Stop the worker and inspect or restore the affected directory. +2. Run `librechat-code clear-workspace-quarantine --worker-dir /projects/second --workspace-id second` using the same deployment/identity configuration. +3. Run the normal worker command with all its root/slot options plus `--reset-workspace-quarantine second`. This verifies the local guard is cleared, resets the server fence, then exits. +4. Restart the normal worker command without the reset option. + +The workspace selector in LibreChat must preserve these registered IDs. Adding +roots here does not grant a principal access or change an agent's selected root. diff --git a/packages/code/src/cli-slots.test.ts b/packages/code/src/cli-slots.test.ts new file mode 100644 index 00000000..cc69ddfc --- /dev/null +++ b/packages/code/src/cli-slots.test.ts @@ -0,0 +1,63 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { spawnSync } from 'node:child_process'; +import { mkdtemp, mkdir, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { fileURLToPath } from 'node:url'; + +test('CLI rejects aliased and overlapping workspace roots before connecting', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'byom-cli-roots-')); + t.after(() => rm(root, { recursive: true, force: true })); + await mkdir(join(root, '..nested')); + for (const extra of [root, join(root, '..nested')]) { + const result = spawnSync( + process.execPath, + [ + fileURLToPath(new URL('./cli.js', import.meta.url)), + 'run', + '--worker-dir', + root, + '--workspace', + `second=${extra}`, + ], + { + encoding: 'utf8', + timeout: 3000, + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'http://127.0.0.1:1', + LIBRECHAT_CODE_WORKER_TOKEN: 'fixture', + LIBRECHAT_CODE_WORKER_ID: 'fixture-worker', + LIBRECHAT_CODE_COMMAND_SANDBOX: 'native-srt', + }, + }, + ); + assert.notEqual(result.status, 0); + assert.match(result.stderr, /must not overlap or alias/); + } +}); + +test('CLI bounds requested workspace slots before connecting', () => { + const result = spawnSync( + process.execPath, + [ + fileURLToPath(new URL('./cli.js', import.meta.url)), + 'run', + '--workspace-lease-slots', + '9', + ], + { + encoding: 'utf8', + timeout: 3000, + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'http://127.0.0.1:1', + LIBRECHAT_CODE_WORKER_TOKEN: 'fixture', + LIBRECHAT_CODE_WORKER_ID: 'fixture-worker', + }, + }, + ); + assert.notEqual(result.status, 0); + assert.match(result.stderr, /cannot exceed 8/); +}); diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 67ec396c..bf9355f2 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -1,8 +1,8 @@ #!/usr/bin/env node import { createHash, createHmac, randomBytes } from 'node:crypto'; import { readFileSync } from 'node:fs'; -import { realpath } from 'node:fs/promises'; -import { basename, resolve } from 'node:path'; +import { realpath, stat } from 'node:fs/promises'; +import { basename, resolve, relative, isAbsolute, sep } from 'node:path'; import { pairBridgeWorker } from './pairing.js'; import { startFileRelay } from './relay.js'; @@ -28,6 +28,10 @@ import { } from './runtime.js'; import { RuntimeWorkspaceCommandSandbox } from './workspace-runtime.js'; import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; +import { NativeWorkspaceCommandPool } from './native-pool.js'; +import { workspaceMutationGuard } from './workspace-guards.js'; +import type { NativeProcessSandboxOptions } from './native-process.js'; +import type { LocalWorkspaceConfig } from './workspace.js'; import { GITHUB_ALLOWED_DOMAINS, GITHUB_CREDENTIAL_ENV_NAME, @@ -154,9 +158,7 @@ function githubCredentials(): { 'Configure either GitHub App authentication or a GitHub token, not both', ); } - const configuredHostValue = nonEmpty( - process.env.LIBRECHAT_CODE_GITHUB_HOST, - ); + const configuredHostValue = nonEmpty(process.env.LIBRECHAT_CODE_GITHUB_HOST); const configuredHost = configuredHostValue ? normalizeGitHubHost(configuredHostValue) : undefined; @@ -467,21 +469,108 @@ async function run( workspaceRoot: canonicalWorkerDirectory, }) : undefined; + const workspaceLeaseSlots = positiveInteger( + 'LIBRECHAT_CODE_WORKSPACE_LEASE_SLOTS', + option(args, '--workspace-lease-slots') ?? + process.env.LIBRECHAT_CODE_WORKSPACE_LEASE_SLOTS, + 1, + ); + if (workspaceLeaseSlots > 8) + throw new Error('Workspace lease slots cannot exceed 8'); + const roots: LocalWorkspaceConfig[] = canonicalWorkerDirectory + ? [ + { + id: workspaceId, + root: canonicalWorkerDirectory, + writable: allowWorkspaceWrites, + name: + option(args, '--workspace-name') ?? + process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? + (useDefaultWorkspace + ? workspaceId + : defaultWorkspaceName(workerDirectory!, workspaceId)), + }, + ] + : []; + for (let i = 0; i < args.length; i++) { + if ( + args[i] === '--workspace' && + (!args[i + 1] || args[i + 1].startsWith('--')) + ) { + throw new Error('--workspace requires id=path'); + } + const value = + args[i] === '--workspace' + ? args[++i] + : args[i].startsWith('--workspace=') + ? args[i].slice('--workspace='.length) + : undefined; + if (value === undefined) continue; + const separator = value.indexOf('='); + if ( + separator < 1 || + separator === value.length - 1 || + !canonicalWorkerDirectory || + commandSandboxMode !== 'native-srt' + ) { + throw new Error( + 'Additional --workspace id=path roots require a primary workspace and native-srt', + ); + } + roots.push({ + id: value.slice(0, separator), + root: await realpath(value.slice(separator + 1)), + writable: allowWorkspaceWrites, + }); + } + // Aliases and nested grants are not independent execution domains. + if (roots.length > 32) + throw new Error('At most 32 workspace roots may be registered'); + const rootIdentities = await Promise.all( + roots.map((root) => stat(root.root)), + ); + const normalized = roots.map((root) => + process.platform === 'linux' ? root.root : root.root.toLowerCase(), + ); + for (let i = 0; i < roots.length; i++) + for (let j = 0; j < i; j++) { + const inside = (a: string, b: string): boolean => { + const path = relative(a, b); + return ( + path === '' || + (path !== '..' && !path.startsWith(`..${sep}`) && !isAbsolute(path)) + ); + }; + if ( + (rootIdentities[i].dev === rootIdentities[j].dev && + rootIdentities[i].ino === rootIdentities[j].ino) || + inside(normalized[i], normalized[j]) || + inside(normalized[j], normalized[i]) + ) { + throw new Error( + 'Workspace roots must not overlap or alias one another', + ); + } + } + if ( + workspaceLeaseSlots > 1 && + (!allowWorkspaceCommands || commandSandboxMode !== 'native-srt') + ) { + throw new Error('Concurrent workspace leases require native-srt commands'); + } + const rootQuarantinePaths = new Map( + roots.map((root) => [ + root.id, + workspaceQuarantinePath({ + codeApiUrl, + workerId, + workspaceRoot: root.root, + }), + ]), + ); let workspaceTools: WorkspaceToolExecutor | undefined = workerDirectory ? await LocalWorkspaceTools.create({ - workspaces: [ - { - id: workspaceId, - name: - option(args, '--workspace-name') ?? - process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? - (useDefaultWorkspace - ? workspaceId - : defaultWorkspaceName(workerDirectory, workspaceId)), - root: workerDirectory, - writable: allowWorkspaceWrites, - }, - ], + workspaces: roots, }) : undefined; if (allowWorkspaceCommands && !canonicalWorkerDirectory) { @@ -657,47 +746,58 @@ async function run( endpoint: sandboxEndpoint, statefulWorkspace, }); + const nativeOptions: NativeProcessSandboxOptions = { + workspaceRoot: canonicalWorkerDirectory!, + protectedPaths: [ + identityPath, + ...rootQuarantinePaths.values(), + github.privateKeyPath, + ].filter((path): path is string => path != null), + allowedDomains: commandAllowedDomains, + ...(github.provider + ? { + maskedEnvironment: { + variables: [ + { + name: GITHUB_CREDENTIAL_ENV_NAME, + extract: '^(.+)$', + injectHosts: [github.host], + }, + ], + async resolve(signal?: AbortSignal) { + return gitHubCredentialEnvironment( + await github.provider!.getCredential(signal), + ); + }, + wrapCommand(command: string, platform: NodeJS.Platform) { + return wrapGitHubCredentialCommand( + command, + github.host, + platform, + ); + }, + }, + } + : {}), + }; const nativeCommandSandbox = allowWorkspaceCommands && commandSandboxMode === 'native-srt' - ? new NativeProcessWorkspaceCommandSandbox({ - workspaceRoot: canonicalWorkerDirectory!, - protectedPaths: [ - identityPath, - mutationQuarantinePath, - github.privateKeyPath, - ].filter((path): path is string => path != null), - allowedDomains: commandAllowedDomains, - ...(github.provider - ? { - maskedEnvironment: { - variables: [ - { - name: GITHUB_CREDENTIAL_ENV_NAME, - extract: '^(.+)$', - injectHosts: [github.host], - }, - ], - async resolve(signal?: AbortSignal) { - return gitHubCredentialEnvironment( - await github.provider!.getCredential(signal), - ); - }, - wrapCommand(command: string, platform: NodeJS.Platform) { - return wrapGitHubCredentialCommand( - command, - github.host, - platform, - ); - }, - }, - } - : {}), - }) + ? roots.length > 1 || workspaceLeaseSlots > 1 + ? new NativeWorkspaceCommandPool( + new Map( + roots.map((root) => [ + root.id, + { ...nativeOptions, workspaceRoot: root.root }, + ]), + ), + workspaceLeaseSlots, + ) + : new NativeProcessWorkspaceCommandSandbox(nativeOptions) : undefined; if (allowWorkspaceCommands && workspaceTools) { workspaceTools = new SandboxWorkspaceTools({ workspaceTools, - commandWorkspaces: [workspaceId], + commandWorkspaces: roots.map((root) => root.id), commandSandbox: nativeCommandSandbox ?? new RuntimeWorkspaceCommandSandbox({ @@ -726,6 +826,9 @@ async function run( ) .digest('hex'), ...(fileRelayEnabled ? { requiresReadyConfirmation: true } : {}), + ...(workspaceLeaseSlots > 1 + ? { workspaceLeaseSlots, requiresReadyConfirmation: true } + : {}), ...(workspaceTools ? { workspaceTools: workspaceTools.capabilities } : {}), }; if (!isValidBridgeWorkerCapabilities(capabilities)) { @@ -752,44 +855,62 @@ async function run( runtimeSupervisor, capabilities, workspaceTools, - workspaceMutationQuarantine: mutationQuarantinePath + ...(workspaceLeaseSlots > 1 || roots.length > 1 ? { - async assertAvailable() { - const record = await loadWorkspaceMutationQuarantine( - mutationQuarantinePath, - ); - if (record != null) { - throw new BridgeProtocolError( - `Workspace mutations are quarantined since ${record.quarantinedAt}: ${record.reason}. Inspect or restore the workspace, then run librechat-code clear-workspace-quarantine`, - undefined, - 'WORKER_QUARANTINED', - ); - } - }, - async arm(reason) { - await saveWorkspaceMutationQuarantine(mutationQuarantinePath, { - version: 1, - workerId, - workspaceId, - ownerId: incarnationId, - quarantinedAt: new Date().toISOString(), - reason, - }); - }, - async clear() { - await clearWorkspaceMutationQuarantine( - mutationQuarantinePath, - incarnationId, - ); - }, - async quarantine() { - await assertWorkspaceMutationQuarantineOwner( - mutationQuarantinePath, - incarnationId, - ); - }, + workspaceQuarantines: new Map( + roots.map((root) => [ + root.id, + workspaceMutationGuard( + rootQuarantinePaths.get(root.id)!, + workerId, + root.id, + incarnationId, + ), + ]), + ), } - : undefined, + : {}), + workspaceMutationQuarantine: + mutationQuarantinePath && + workspaceLeaseSlots === 1 && + roots.length === 1 + ? { + async assertAvailable() { + const record = await loadWorkspaceMutationQuarantine( + mutationQuarantinePath, + ); + if (record != null) { + throw new BridgeProtocolError( + `Workspace mutations are quarantined since ${record.quarantinedAt}: ${record.reason}. Inspect or restore the workspace, then run librechat-code clear-workspace-quarantine`, + undefined, + 'WORKER_QUARANTINED', + ); + } + }, + async arm(reason) { + await saveWorkspaceMutationQuarantine(mutationQuarantinePath, { + version: 1, + workerId, + workspaceId, + ownerId: incarnationId, + quarantinedAt: new Date().toISOString(), + reason, + }); + }, + async clear() { + await clearWorkspaceMutationQuarantine( + mutationQuarantinePath, + incarnationId, + ); + }, + async quarantine() { + await assertWorkspaceMutationQuarantineOwner( + mutationQuarantinePath, + incarnationId, + ); + }, + } + : undefined, onIdentityChange: pairedIdentity && identityPath ? async (identity) => { @@ -832,6 +953,16 @@ async function run( ); return; } + const resetNativeRoot = option(args, '--reset-workspace-quarantine'); + if (resetNativeRoot != null) { + await worker.refreshCredential(controller.signal); + await worker.register(controller.signal); + await worker.resetNativeWorkspace(resetNativeRoot, controller.signal); + process.stdout.write( + `librechat-code: reset acknowledged for native workspace ${resetNativeRoot}\n`, + ); + return; + } await worker.run(controller.signal); } finally { try { diff --git a/packages/code/src/native-pool.ts b/packages/code/src/native-pool.ts new file mode 100644 index 00000000..82507d65 --- /dev/null +++ b/packages/code/src/native-pool.ts @@ -0,0 +1,117 @@ +import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; +import { WorkspaceToolError } from './workspace.js'; +import type { NativeProcessSandboxOptions } from './native-process.js'; +import type { + WorkspaceExecuteCommandRequest, + WorkspaceExecuteCommandResult, +} from './protocol.js'; + +interface Entry { + sandbox: NativeProcessWorkspaceCommandSandbox; + busy: boolean; +} + +/** Bounded persistent executor cache. Each child owns one root's SRT policy; + * only idle children may be evicted, and ambiguous failures are never retried. + */ +export class NativeWorkspaceCommandPool { + readonly mutationFailuresAreAtomic = true as const; + private readonly entries = new Map(); + private allocation: Promise = Promise.resolve(); + private closing = false; + constructor( + private readonly roots: ReadonlyMap, + private readonly capacity: number, + ) { + if ( + !Number.isSafeInteger(capacity) || + capacity < 1 || + capacity > 8 || + roots.size === 0 + ) { + throw new Error('Native executor capacity must be between 1 and 8'); + } + } + + private allocate(root: string): Promise { + const pending = this.allocation.then(async () => { + const options = this.roots.get(root); + if (this.closing || !options) + throw new WorkspaceToolError( + 'Native workspace unavailable', + 'REGISTRATION_INVALID', + ); + let entry = this.entries.get(root); + if (entry?.busy) + throw new WorkspaceToolError( + 'Native workspace already executing', + 'COMMAND_UNAVAILABLE', + ); + if (!entry) { + if (this.entries.size >= this.capacity) { + const idle = [...this.entries].find( + ([, candidate]) => !candidate.busy, + ); + if (!idle) + throw new WorkspaceToolError( + 'Native executor capacity reached', + 'COMMAND_UNAVAILABLE', + ); + await idle[1].sandbox.close(); + this.entries.delete(idle[0]); + } + entry = { + sandbox: new NativeProcessWorkspaceCommandSandbox(options), + busy: false, + }; + } + entry.busy = true; + // Map insertion order is the idle eviction order. + this.entries.delete(root); + this.entries.set(root, entry); + return entry; + }); + this.allocation = pending.catch(() => undefined); + return pending; + } + + async prepare(): Promise { + const entry = await this.allocate(this.roots.keys().next().value!); + try { + await entry.sandbox.prepare(); + } finally { + entry.busy = false; + } + } + + async execute( + request: WorkspaceExecuteCommandRequest, + signal?: AbortSignal, + ): Promise { + const entry = await this.allocate(request.workspaceId); + try { + if (signal?.aborted) + throw new WorkspaceToolError( + 'Command cancelled before dispatch', + 'EXECUTION_ABORTED', + ); + return await entry.sandbox.execute(request, signal); + } finally { + entry.busy = false; + } + } + + async close(): Promise { + this.closing = true; + await this.allocation; + const results = await Promise.allSettled( + [...this.entries.values()].map((entry) => entry.sandbox.close()), + ); + this.entries.clear(); + const errors = results + .filter((result) => result.status === 'rejected') + .map((result) => result.reason); + if (errors.length) + throw new AggregateError(errors, 'Native executor pool shutdown failed'); + } +} diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 37288d63..07088f94 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -390,14 +390,11 @@ const WORKSPACE_COMMAND_RESULT_KEYS = new Set([ 'truncated', 'timedOut', ]); -const WORKSPACE_SEARCH_MATCH_KEYS = new Set([ - 'path', - 'line', - 'column', - 'text', -]); +const WORKSPACE_SEARCH_MATCH_KEYS = new Set(['path', 'line', 'column', 'text']); export interface BridgeWorkerCapabilities { + /** Opt-in protocol: maximum concurrently leased independent workspace roots. */ + workspaceLeaseSlots?: number; statefulWorkspace: boolean; sandboxProfile: string; runtimes: string[]; @@ -414,6 +411,8 @@ export interface BridgeWorkerRegistration { } export interface BridgeWorkerRegistrationResponse { + /** Absent on legacy servers. Workers must not parallelize without this receipt. */ + workspaceLeaseSlots?: number; protocolVersion: BridgeProtocolVersion; workerId: string; incarnationId: string; @@ -464,6 +463,7 @@ export interface BridgeSandboxRequest { } export interface BridgeAssignment { + workspaceLeaseSlot?: number; protocolVersion: BridgeProtocolVersion; assignmentId: string; workerId: string; @@ -551,7 +551,8 @@ export function isWorkspaceToolErrorCode( } export type BridgeSettlement = - BridgeFulfilledSettlement | BridgeRejectedSettlement; + | BridgeFulfilledSettlement + | BridgeRejectedSettlement; export interface BridgeSettlementResponse { protocolVersion: BridgeProtocolVersion; @@ -608,7 +609,10 @@ function normalizePortableRelativePath(value: string): string { } /** Compare path segments in ripgrep's sorted, depth-first traversal order. */ -export function comparePortableRelativePaths(left: string, right: string): number { +export function comparePortableRelativePaths( + left: string, + right: string, +): number { const encoder = new TextEncoder(); const leftSegments = left.split('/'); const rightSegments = right.split('/'); @@ -638,9 +642,14 @@ function isWithinRequestedPath(candidate: string, requested?: string): boolean { ); } -function isValidWorkspaceEditRequest(request: Record): boolean { +function isValidWorkspaceEditRequest( + request: Record, +): boolean { const hasBatch = request.edits !== undefined; - if (hasBatch && (request.oldText !== undefined || request.newText !== undefined)) { + if ( + hasBatch && + (request.oldText !== undefined || request.newText !== undefined) + ) { return false; } const edits = hasBatch @@ -746,7 +755,8 @@ export function isWorkspaceToolRequest( isSafePortableRelativePath(request.path)) && (request.afterPath === undefined || (isSafePortableRelativePath(request.afterPath) && - normalizePortableRelativePath(request.afterPath) === request.afterPath && + normalizePortableRelativePath(request.afterPath) === + request.afterPath && isWithinRequestedPath(request.afterPath, request.path))) && (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && @@ -833,11 +843,16 @@ export function isWorkspaceToolResult( const maxLines = request.maxLines ?? 200; const content = typeof result.content === 'string' ? result.content : null; const reportedLineCount = - Number.isSafeInteger(result.endLine) && Number(result.endLine) >= startLine - 1 + Number.isSafeInteger(result.endLine) && + Number(result.endLine) >= startLine - 1 ? Number(result.endLine) - startLine + 1 : -1; const actualLineCount = - content === null ? -1 : content.length === 0 ? reportedLineCount : content.split('\n').length; + content === null + ? -1 + : content.length === 0 + ? reportedLineCount + : content.split('\n').length; return ( hasOnlyKeys(result, WORKSPACE_READ_RESULT_KEYS) && result.path === request.path && @@ -1125,6 +1140,10 @@ export function isValidBridgeWorkerCapabilities( if (typeof value !== 'object' || value === null) return false; const capabilities = value as Record; return ( + (capabilities.workspaceLeaseSlots === undefined || + (Number.isSafeInteger(capabilities.workspaceLeaseSlots) && + Number(capabilities.workspaceLeaseSlots) >= 1 && + Number(capabilities.workspaceLeaseSlots) <= 8)) && typeof capabilities.statefulWorkspace === 'boolean' && typeof capabilities.sandboxProfile === 'string' && capabilities.sandboxProfile.trim().length > 0 && diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts new file mode 100644 index 00000000..d31aebed --- /dev/null +++ b/packages/code/src/worker-slots.test.ts @@ -0,0 +1,73 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { BridgeWorker } from './worker.js'; +import type { BridgeWorkspaceToolCapabilities } from './protocol.js'; + +const capabilities: BridgeWorkspaceToolCapabilities = { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'a' }, { id: 'b' }], +}; +for (const receipt of [undefined, 1]) { + test(`worker keeps serial lease wire format for receipt ${receipt}`, async () => { + const controller = new AbortController(); + let leases = 0; + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + incarnationId: 'incarnation-test-slots', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'fixture', + runtimes: [], + workspaceLeaseSlots: 2, + requiresReadyConfirmation: true, + workspaceTools: capabilities, + }, + workspaceTools: { + capabilities, + async execute() { + throw new Error('must not execute'); + }, + }, + workspaceQuarantines: new Map( + ['a', 'b'].map((root) => [ + root, + { + async assertAvailable() {}, + async arm() {}, + async clear() {}, + async quarantine() {}, + }, + ]), + ), + fetchImpl: async (url, init) => { + const path = new URL(String(url)).pathname; + if (path.endsWith('/register')) + return Response.json({ + protocolVersion: 1, + workerId: 'worker', + incarnationId: 'incarnation-test-slots', + registrationGeneration: 1, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60000, + ...(receipt === undefined ? {} : { workspaceLeaseSlots: receipt }), + }); + if (path.endsWith('/lease')) { + leases++; + assert.equal( + JSON.parse(String(init?.body)).workspaceLeaseSlot, + undefined, + ); + controller.abort(); + } + return Response.json({ protocolVersion: 1, ready: true }); + }, + }); + await worker.run(controller.signal); + assert.equal(leases, 1); + await assert.rejects(worker.lease(undefined, 0), /negotiated capacity/); + }); +} diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 1fda70f7..037ecd68 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -35,6 +35,8 @@ export interface BridgeWorkerOptions { capabilities: BridgeWorkerCapabilities; workspaceTools?: WorkspaceToolExecutor; workspaceMutationQuarantine?: WorkspaceMutationQuarantine; + /** Required per-root durable guards when opting into concurrent workspace leases. */ + workspaceQuarantines?: ReadonlyMap; leaseWaitMs?: number; leaseTransportGraceMs?: number; registrationTransportTimeoutMs?: number; @@ -59,9 +61,13 @@ export interface BridgeWorkerOptions { export interface WorkspaceMutationQuarantine { assertAvailable(): Promise; - arm(reason: string): Promise; - clear(): Promise; - quarantine(reason: string, cause?: unknown): Promise; + arm(reason: string, assignmentId?: string): Promise; + clear(assignmentId?: string): Promise; + quarantine( + reason: string, + cause?: unknown, + assignmentId?: string, + ): Promise; } export interface BridgeWorkerIdentity { @@ -141,21 +147,23 @@ function workspaceCapabilitiesMatch( advertised.writeFileModes?.length === executor.writeFileModes?.length && (advertised.writeFileModes?.every( (mode, index) => mode === executor.writeFileModes?.[index], - ) ?? executor.writeFileModes == null) && + ) ?? + executor.writeFileModes == null) && advertised.editFileModes?.length === executor.editFileModes?.length && (advertised.editFileModes?.every( (mode, index) => mode === executor.editFileModes?.[index], - ) ?? executor.editFileModes == null) && - advertised.editFileFeatures?.length === - executor.editFileFeatures?.length && + ) ?? + executor.editFileModes == null) && + advertised.editFileFeatures?.length === executor.editFileFeatures?.length && (advertised.editFileFeatures?.every( (feature, index) => feature === executor.editFileFeatures?.[index], - ) ?? executor.editFileFeatures == null) && - advertised.listFileFeatures?.length === - executor.listFileFeatures?.length && + ) ?? + executor.editFileFeatures == null) && + advertised.listFileFeatures?.length === executor.listFileFeatures?.length && (advertised.listFileFeatures?.every( (feature, index) => feature === executor.listFileFeatures?.[index], - ) ?? executor.listFileFeatures == null) && + ) ?? + executor.listFileFeatures == null) && advertised.workspaces.length === executor.workspaces.length && advertised.workspaces.every( (workspace, index) => @@ -167,7 +175,8 @@ function workspaceCapabilitiesMatch( (operation, operationIndex) => operation === executor.workspaces[index]?.operations?.[operationIndex], - ) ?? executor.workspaces[index]?.operations == null), + ) ?? + executor.workspaces[index]?.operations == null), ) ); } @@ -179,8 +188,7 @@ function registrationCompatibleCapabilities( if ( workspaceTools == null || (workspaceTools.operations.every( - (operation) => - operation === 'read_file' || operation === 'search_text', + (operation) => operation === 'read_file' || operation === 'search_text', ) && workspaceTools.workspaces.every( (workspace) => workspace.operations == null, @@ -189,8 +197,7 @@ function registrationCompatibleCapabilities( return capabilities; } const operations = workspaceTools.operations.filter( - (operation) => - operation === 'read_file' || operation === 'search_text', + (operation) => operation === 'read_file' || operation === 'search_text', ); if (operations.length === 0) { const { workspaceTools: _workspaceTools, ...compatible } = capabilities; @@ -199,7 +206,9 @@ function registrationCompatibleCapabilities( const workspaces = workspaceTools.workspaces.flatMap((workspace) => { if ( workspace.operations != null && - !operations.every((operation) => workspace.operations?.includes(operation)) + !operations.every((operation) => + workspace.operations?.includes(operation), + ) ) { return []; } @@ -344,9 +353,29 @@ export class BridgeWorker { private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; private lastRegisteredAtMs = 0; private mutationGuardArmed = false; + private readonly activeWorkspaceAssignments = new Map< + string, + { id: string; done: Promise } + >(); + private readonly armedWorkspaces = new Set(); + private negotiatedWorkspaceSlots = 1; + private concurrentRunning = false; + private registrationInFlight?: Promise; + private credentialInFlight?: Promise; private serverClockOffsetMs = MAX_PROOF_CLOCK_SKEW_MS; constructor(private readonly options: BridgeWorkerOptions) { + const requestedSlots = options.capabilities.workspaceLeaseSlots; + if ( + requestedSlots !== undefined && + (!Number.isSafeInteger(requestedSlots) || + requestedSlots < 1 || + requestedSlots > 8) + ) { + throw new BridgeProtocolError( + 'Workspace lease slots must be an integer from 1 to 8', + ); + } if (!options.token && !options.identity) { throw new BridgeProtocolError( 'Bridge worker requires a static token or paired identity', @@ -358,7 +387,9 @@ export class BridgeWorker { ); } if (options.runtimeSupervisor == null && !options.sandboxEndpoint?.trim()) { - throw new BridgeProtocolError('Bridge worker requires a runtime supervisor'); + throw new BridgeProtocolError( + 'Bridge worker requires a runtime supervisor', + ); } if ( (options.workspaceTools == null) !== @@ -381,12 +412,26 @@ export class BridgeWorker { operation === 'edit_file' || operation === 'execute_command', ) === true && - options.workspaceMutationQuarantine == null + options.workspaceMutationQuarantine == null && + options.workspaceQuarantines == null ) { throw new BridgeProtocolError( 'Workspace mutation capabilities require durable quarantine storage', ); } + if ((options.capabilities.workspaceLeaseSlots ?? 1) > 1) { + if ( + options.capabilities.requiresReadyConfirmation !== true || + options.workspaceQuarantines == null || + options.capabilities.workspaceTools?.workspaces.some( + (root) => !options.workspaceQuarantines!.has(root.id), + ) !== false + ) { + throw new BridgeProtocolError( + 'Concurrent workspaces require per-root durable guards and readiness confirmation', + ); + } + } this.fetchImpl = options.fetchImpl ?? fetch; this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); this.runtimeSupervisor = @@ -413,6 +458,20 @@ export class BridgeWorker { private async registerWithPolicy( signal: AbortSignal | undefined, allowActiveMutation: boolean, + ): Promise { + if (this.registrationInFlight) return await this.registrationInFlight; + const pending = this.registerOwned(signal, allowActiveMutation); + this.registrationInFlight = pending; + try { + return await pending; + } finally { + this.registrationInFlight = undefined; + } + } + + private async registerOwned( + signal: AbortSignal | undefined, + allowActiveMutation: boolean, ): Promise { if (!allowActiveMutation) { try { @@ -503,6 +562,21 @@ export class BridgeWorker { 'Code API registered a different worker incarnation', ); } + const slots = registration.workspaceLeaseSlots ?? 1; + if ( + !Number.isSafeInteger(slots) || + slots < 1 || + slots > (this.options.capabilities.workspaceLeaseSlots ?? 1) || + slots > 8 || + (this.concurrentRunning && slots !== this.negotiatedWorkspaceSlots) + ) { + throw new BridgeProtocolError( + 'Code API workspace slot negotiation changed or exceeded local policy', + undefined, + 'WORKER_FENCED', + ); + } + this.negotiatedWorkspaceSlots = slots; const registeredAtMs = Date.parse(registration.registeredAt); if (Number.isFinite(registeredAtMs)) { this.serverClockOffsetMs = registeredAtMs - registrationStartedAtMs; @@ -571,7 +645,54 @@ export class BridgeWorker { ); } - async lease(signal?: AbortSignal): Promise { + async resetNativeWorkspace( + workspaceId: string, + signal?: AbortSignal, + ): Promise { + const guard = this.options.workspaceQuarantines?.get(workspaceId); + if ( + !guard || + this.activeWorkspaceAssignments.size > 0 || + !this.options.capabilities.workspaceTools?.workspaces.some( + (root) => root.id === workspaceId, + ) + ) { + throw new BridgeProtocolError( + 'Native workspace reset requires an idle registered root', + ); + } + // The operator must have inspected/restored the root and cleared its + // machine-local guard before the remote fence can be removed. + await guard.assertAvailable(); + await this.timedRequest( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + runtimeSessionId: `native-workspace:${workspaceId}`, + confirmDiscarded: true, + }, + this.options.resetTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + signal, + ); + } + + async lease( + signal?: AbortSignal, + workspaceLeaseSlot?: number, + ): Promise { + if ( + workspaceLeaseSlot !== undefined && + (!Number.isSafeInteger(workspaceLeaseSlot) || + workspaceLeaseSlot < 0 || + this.negotiatedWorkspaceSlots <= 1 || + workspaceLeaseSlot >= this.negotiatedWorkspaceSlots) + ) { + throw new BridgeProtocolError( + 'Workspace lease slot exceeds negotiated capacity', + ); + } const waitMs = Math.min( MAX_LEASE_WAIT_MS, Math.max(0, this.options.leaseWaitMs ?? DEFAULT_LEASE_WAIT_MS), @@ -601,6 +722,7 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, waitMs, incarnationId: this.incarnationId, + ...(workspaceLeaseSlot === undefined ? {} : { workspaceLeaseSlot }), }, leaseController.signal, ); @@ -610,7 +732,8 @@ export class BridgeWorker { } if ( response.assignment != null && - response.assignment.incarnationId !== this.incarnationId + (response.assignment.incarnationId !== this.incarnationId || + response.assignment.workspaceLeaseSlot !== workspaceLeaseSlot) ) { throw new BridgeProtocolError( 'Code API leased an assignment for a different worker incarnation', @@ -703,6 +826,10 @@ export class BridgeWorker { try { await this.refreshCredential(signal); await this.register(signal); + if (this.negotiatedWorkspaceSlots > 1) { + await this.runConcurrent(signal); + return; + } const assignment = await this.lease(signal); reconnectAttempt = 0; if (!assignment) continue; @@ -734,13 +861,128 @@ export class BridgeWorker { } } + private async runConcurrent(signal?: AbortSignal): Promise { + const controller = new AbortController(); + const abort = (): void => controller.abort(signal?.reason); + signal?.addEventListener('abort', abort, { once: true }); + if (signal?.aborted) abort(); + let failure: unknown; + const fail = (error: unknown): void => { + failure ??= error; + controller.abort(error); + }; + this.concurrentRunning = true; + const heartbeat = this.maintainRegistration(controller.signal).catch(fail); + const lane = async (slot?: number): Promise => { + let retries = 0; + while (!controller.signal.aborted) { + let assignment: BridgeAssignment | undefined; + try { + assignment = await this.lease(controller.signal, slot); + retries = 0; + if (assignment == null) continue; + await this.executeAndSettle(assignment, controller.signal); + } catch (error) { + if ( + error instanceof BridgeWorkspaceQuarantinedError && + assignment?.workspaceLeaseSlot !== undefined + ) { + // The durable local guard is retained. A distinct receipt tells + // Code API to release this slot without declaring the root clean. + try { + await this.timedRequest( + this.assignmentUrl(assignment, 'quarantine'), + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected', + error: + 'Workspace quarantined after an uncertain execution or settlement; inspect it before resetting.', + }, + this.options.leaseAckTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + ); + this.options.onError?.(error); + continue; + } catch (quarantineError) { + fail(quarantineError); + return; + } + } + if (controller.signal.aborted) return; + if ( + assignment != null || + (error instanceof BridgeProtocolError && + (error.status === 401 || + error.status === 403 || + error.code === 'WORKER_FENCED' || + error.code === 'WORKER_QUARANTINED')) + ) { + fail(error); + return; + } + this.options.onError?.(error); + await abortableDelay( + reconnectDelayMs( + retries++, + this.options.reconnectDelayMs, + this.options.reconnectMaxDelayMs, + this.options.reconnectRandom, + ), + controller.signal, + ); + } + } + }; + try { + // The legacy lane serves run-code requests only when the aggregate lock + // excludes workspace slots. It never increases simultaneous executions. + await Promise.all([ + lane(), + ...Array.from({ length: this.negotiatedWorkspaceSlots }, (_, i) => + lane(i), + ), + ]); + } finally { + controller.abort(); + await heartbeat; + this.concurrentRunning = false; + signal?.removeEventListener('abort', abort); + } + if (failure != null) throw failure; + } + async refreshCredential( signal?: AbortSignal, - validThroughMs = - Date.now() + + validThroughMs = Date.now() + this.serverClockOffsetMs + (this.options.credentialRefreshWindowMs ?? CREDENTIAL_REFRESH_WINDOW_MS), transportTimeoutMs = Number.POSITIVE_INFINITY, + ): Promise { + while (this.credentialInFlight) { + await this.credentialInFlight; + // A longer-lived caller may still need another refresh after this one. + } + const pending = this.refreshCredentialOwned( + signal, + validThroughMs, + transportTimeoutMs, + ); + this.credentialInFlight = pending; + try { + await pending; + } finally { + if (this.credentialInFlight === pending) + this.credentialInFlight = undefined; + } + } + + private async refreshCredentialOwned( + signal: AbortSignal | undefined, + validThroughMs: number, + transportTimeoutMs: number, ): Promise { const identity = this.options.identity; if (identity == null) return; @@ -815,8 +1057,7 @@ export class BridgeWorker { error instanceof BridgeProtocolError && (error.status === 401 || error.status === 403); const credentialRemainingMs = - Date.parse(identity.expiresAt) - - (Date.now() + serverClockOffsetMs); + Date.parse(identity.expiresAt) - (Date.now() + serverClockOffsetMs); if (terminal || credentialRemainingMs <= 0) throw error; await abortableDelay( Math.min( @@ -833,6 +1074,70 @@ export class BridgeWorker { assignment: BridgeAssignment, signal?: AbortSignal, ): Promise { + const root = + assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) + ? assignment.request.workspaceId + : undefined; + const waitingAt = Date.now(); + while (root != null && this.activeWorkspaceAssignments.has(root)) { + const active = this.activeWorkspaceAssignments.get(root)!; + if (active.id === assignment.assignmentId) + throw new BridgeProtocolError( + 'Code API replayed an active workspace assignment', + undefined, + 'WORKER_FENCED', + ); + // A settlement can commit remotely before the local durable guard clears. + // Keep the next lane out of the root until that cleanup has finished. + await active.done; + if (signal?.aborted) + throw signal.reason ?? new DOMException('aborted', 'AbortError'); + } + let release!: () => void; + if (root != null) + this.activeWorkspaceAssignments.set(root, { + id: assignment.assignmentId, + done: new Promise((resolve) => { + release = resolve; + }), + }); + const adjusted = + assignment.remainingMs === undefined + ? assignment + : { + ...assignment, + remainingMs: Math.max( + 0, + assignment.remainingMs - (Date.now() - waitingAt), + ), + }; + try { + await this.executeOwned(adjusted, signal); + } finally { + if (root != null) { + this.activeWorkspaceAssignments.delete(root); + release(); + } + } + } + + private workspaceGuard( + assignment: BridgeAssignment, + ): WorkspaceMutationQuarantine | undefined { + return assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) + ? (this.options.workspaceQuarantines?.get( + assignment.request.workspaceId, + ) ?? this.options.workspaceMutationQuarantine) + : this.options.workspaceMutationQuarantine; + } + + private async executeOwned( + assignment: BridgeAssignment, + signal?: AbortSignal, + ): Promise { + const guard = this.workspaceGuard(assignment); if (signal?.aborted === true) { throw signal.reason instanceof Error ? signal.reason @@ -897,8 +1202,10 @@ export class BridgeWorker { } const heartbeatController = new AbortController(); let heartbeatError: unknown; - const heartbeat = this.maintainRegistration( - heartbeatController.signal, + const heartbeat = ( + this.concurrentRunning + ? Promise.resolve() + : this.maintainRegistration(heartbeatController.signal) ).catch((error) => { heartbeatError = error; executionController.abort(); @@ -914,7 +1221,9 @@ export class BridgeWorker { let settlement: BridgeSettlement; let ambiguousSandboxError: unknown; let ambiguousWorkspaceMutationError: unknown; - let workspaceMutationGuardError: BridgeWorkspaceQuarantinedError | undefined; + let workspaceMutationGuardError: + | BridgeWorkspaceQuarantinedError + | undefined; let sandboxRejectedExecution = false; let sandboxStarted = false; let workspaceMutationArmed = false; @@ -941,6 +1250,14 @@ export class BridgeWorker { throw new BridgeProtocolError('Invalid workspace tool request'); } const workspaceRequest = assignment.request; + try { + await guard?.assertAvailable(); + } catch (error) { + throw new BridgeWorkspaceQuarantinedError( + 'Workspace is quarantined', + error, + ); + } const advertised = this.activeCapabilities.workspaceTools; if (advertised == null) { throw new BridgeProtocolError( @@ -983,7 +1300,8 @@ export class BridgeWorker { workspaceRequest.operation === 'preview_edit' || workspaceRequest.operation === 'edit_file' ) { - const mode = workspaceRequest.edits === undefined ? 'single' : 'batch'; + const mode = + workspaceRequest.edits === undefined ? 'single' : 'batch'; const modes = advertised.editFileModes; if ( (modes == null && mode !== 'single') || @@ -1019,8 +1337,10 @@ export class BridgeWorker { if (isMutation) { this.mutationGuardArmed = true; try { - await this.options.workspaceMutationQuarantine!.arm( + this.armedWorkspaces.add(workspaceRequest.workspaceId); + await guard!.arm( `Workspace mutation ${workspaceRequest.operation} is pending settlement`, + assignment.assignmentId, ); workspaceMutationArmed = true; } catch (error) { @@ -1040,10 +1360,8 @@ export class BridgeWorker { !advertised.listFileFeatures?.includes('after_path') && 'nextAfterPath' in payload ) { - const { - nextAfterPath: _nextAfterPath, - ...compatiblePayload - } = payload; + const { nextAfterPath: _nextAfterPath, ...compatiblePayload } = + payload; payload = compatiblePayload; } workspaceMutationApplied = isMutation; @@ -1176,11 +1494,10 @@ export class BridgeWorker { error instanceof WorkspaceToolError ? { errorCode: error.code } : {}), - error: - (error instanceof Error - ? error.message - : 'Sandbox execution failed' - ).slice(0, MAX_SETTLEMENT_ERROR_LENGTH), + error: (error instanceof Error + ? error.message + : 'Sandbox execution failed' + ).slice(0, MAX_SETTLEMENT_ERROR_LENGTH), }; } @@ -1190,12 +1507,14 @@ export class BridgeWorker { credentialController.abort(); await credentialMaintenance; try { - if (workspaceMutationGuardError != null) throw workspaceMutationGuardError; + if (workspaceMutationGuardError != null) + throw workspaceMutationGuardError; if (ambiguousWorkspaceMutationError != null) { throw await this.quarantineWorkspace( undefined, 'Worker stopped after a workspace mutation completed without a fulfilled settlement', ambiguousWorkspaceMutationError, + assignment, ); } if (ambiguousSandboxError != null) { @@ -1203,6 +1522,7 @@ export class BridgeWorker { assignment.runtimeSessionId, `Stateful workspace ${assignment.runtimeSessionId} was quarantined after an ambiguous sandbox execution`, ambiguousSandboxError, + assignment, ); } const knownCleanStatefulRejection = @@ -1242,7 +1562,13 @@ export class BridgeWorker { } if (workspaceMutationArmed) { try { - await this.options.workspaceMutationQuarantine!.clear(); + await guard!.clear(assignment.assignmentId); + if ( + assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) + ) { + this.armedWorkspaces.delete(assignment.request.workspaceId); + } this.mutationGuardArmed = false; } catch (error) { throw new BridgeWorkspaceQuarantinedError( @@ -1282,7 +1608,9 @@ export class BridgeWorker { return await lease.execute({ body, headers, signal }); } if (lease.endpoint == null) { - throw new BridgeProtocolError('Runtime lease does not provide an execution transport'); + throw new BridgeProtocolError( + 'Runtime lease does not provide an execution transport', + ); } const endpoint = lease.endpoint.replace(/\/+$/, ''); const response = await this.fetchImpl(`${endpoint}/execute`, { @@ -1307,6 +1635,7 @@ export class BridgeWorker { assignment.runtimeSessionId, `Stateful workspace ${assignment.runtimeSessionId} could not release its runtime lease`, error, + assignment, ); } } @@ -1315,13 +1644,15 @@ export class BridgeWorker { runtimeSessionId: string | undefined, message: string, cause?: unknown, + assignment?: BridgeAssignment, ): Promise { if (runtimeSessionId == null) { try { - await this.options.workspaceMutationQuarantine?.quarantine( - message, - cause, - ); + await ( + assignment == null + ? this.options.workspaceMutationQuarantine + : this.workspaceGuard(assignment) + )?.quarantine(message, cause, assignment?.assignmentId); return new BridgeWorkspaceQuarantinedError(message, cause); } catch (error) { return new BridgeWorkspaceQuarantinedError( @@ -1355,15 +1686,16 @@ export class BridgeWorker { Math.floor(this.registrationTtlMs / 2), ); await this.delay( - Math.max( - 0, - this.lastRegisteredAtMs + heartbeatIntervalMs - Date.now(), - ), + Math.max(0, this.lastRegisteredAtMs + heartbeatIntervalMs - Date.now()), signal, ); if (signal.aborted) return; try { - await this.registerWithPolicy(signal, this.mutationGuardArmed); + if (this.concurrentRunning) await this.refreshCredential(signal); + await this.registerWithPolicy( + signal, + this.mutationGuardArmed || this.armedWorkspaces.size > 0, + ); } catch (error) { const terminal = error instanceof BridgeProtocolError && @@ -1439,6 +1771,7 @@ export class BridgeWorker { ? `Stateful workspace ${assignment.runtimeSessionId} was quarantined before settlement during shutdown` : 'Worker stopped after a workspace mutation could not be settled during shutdown', signal.reason, + assignment, ); } throw signal.reason instanceof Error @@ -1483,6 +1816,7 @@ export class BridgeWorker { ? `Stateful workspace ${assignment.runtimeSessionId} was quarantined after Code API rejected its fulfilled settlement` : 'Worker stopped after Code API rejected a fulfilled workspace mutation settlement', error, + assignment, ); } throw error; @@ -1509,6 +1843,7 @@ export class BridgeWorker { ? `Stateful workspace ${assignment.runtimeSessionId} was quarantined after ambiguous settlement delivery` : 'Worker stopped after ambiguous workspace mutation settlement delivery', lastError, + assignment, ); } if (lastError instanceof Error) throw lastError; diff --git a/packages/code/src/workspace-guards.ts b/packages/code/src/workspace-guards.ts new file mode 100644 index 00000000..98b6bb02 --- /dev/null +++ b/packages/code/src/workspace-guards.ts @@ -0,0 +1,49 @@ +import { BridgeProtocolError } from './protocol.js'; +import { + assertWorkspaceMutationQuarantineOwner, + clearWorkspaceMutationQuarantine, + loadWorkspaceMutationQuarantine, + saveWorkspaceMutationQuarantine, +} from './storage.js'; +import type { WorkspaceMutationQuarantine } from './worker.js'; + +/** Each path is outside sandbox roots and each pending write has a unique owner. */ +export function workspaceMutationGuard( + path: string, + workerId: string, + workspaceId: string, + incarnationId: string, +): WorkspaceMutationQuarantine { + const owner = (assignmentId?: string): string => { + if (!assignmentId) + throw new Error('Workspace mutation requires an assignment owner'); + return `${incarnationId}:${assignmentId}`; + }; + return { + async assertAvailable() { + const record = await loadWorkspaceMutationQuarantine(path); + if (record != null) + throw new BridgeProtocolError( + `Workspace ${workspaceId} is quarantined; inspect it and clear its quarantine before restarting the worker`, + undefined, + 'WORKSPACE_QUARANTINED', + ); + }, + async arm(reason, assignmentId) { + await saveWorkspaceMutationQuarantine(path, { + version: 1, + workerId, + workspaceId, + ownerId: owner(assignmentId), + quarantinedAt: new Date().toISOString(), + reason, + }); + }, + async clear(assignmentId) { + await clearWorkspaceMutationQuarantine(path, owner(assignmentId)); + }, + async quarantine(_reason, _cause, assignmentId) { + await assertWorkspaceMutationQuarantineOwner(path, owner(assignmentId)); + }, + }; +} diff --git a/service/src/bridge/admission.ts b/service/src/bridge/admission.ts index 1e19286e..4002ed6d 100644 --- a/service/src/bridge/admission.ts +++ b/service/src/bridge/admission.ts @@ -7,15 +7,21 @@ export class BridgeAdmissionQueue { private readonly capacity = 32, ) {} - private keys(workerId: string): [string, string, string] { + private keys(workerId: string): [string, string, string, string] { const prefix = `codeapi:bridge:v1:worker:${encodeURIComponent(workerId)}:admission`; - return [prefix, `${prefix}:deadlines`, `${prefix}:sequence`]; + return [ + prefix, + `${prefix}:deadlines`, + `${prefix}:sequence`, + `${prefix}:workspaces`, + ]; } async enter( workerId: string, id: string, deadlineAtMs: number, + workspaceId?: string, ): Promise { return ( Number( @@ -25,31 +31,34 @@ export class BridgeAdmissionQueue { 'for _, id in ipairs(expired) do', " redis.call('ZREM', KEYS[1], id)", " redis.call('ZREM', KEYS[2], id)", + " redis.call('HDEL', KEYS[4], id)", 'end', "if redis.call('ZSCORE', KEYS[1], ARGV[1]) then return 1 end", "if redis.call('ZCARD', KEYS[1]) >= tonumber(ARGV[4]) then return 0 end", "local sequence = redis.call('INCR', KEYS[3])", "redis.call('ZADD', KEYS[1], sequence, ARGV[1])", "redis.call('ZADD', KEYS[2], ARGV[3], ARGV[1])", + "if ARGV[5] ~= '' then redis.call('HSET', KEYS[4], ARGV[1], ARGV[5]) end", "local latest = redis.call('ZREVRANGE', KEYS[2], 0, 0, 'WITHSCORES')", 'for _, key in ipairs(KEYS) do', " redis.call('PEXPIREAT', key, tonumber(latest[2]) + 30000)", 'end', 'return 1', ].join('\n'), - 3, + 4, ...this.keys(workerId), id, Date.now(), deadlineAtMs, this.capacity, + workspaceId ?? '', ), ) === 1 ); } async isHead(workerId: string, id: string): Promise { - const [order, deadlines] = this.keys(workerId); + const [order, deadlines, , workspaces] = this.keys(workerId); return ( Number( await this.redis.eval( @@ -58,14 +67,16 @@ export class BridgeAdmissionQueue { 'for _, id in ipairs(expired) do', " redis.call('ZREM', KEYS[1], id)", " redis.call('ZREM', KEYS[2], id)", + " redis.call('HDEL', KEYS[3], id)", 'end', "local head = redis.call('ZRANGE', KEYS[1], 0, 0)", 'if head[1] == ARGV[1] then return 1 end', 'return 0', ].join('\n'), - 2, + 3, order, deadlines, + workspaces, id, Date.now(), ), @@ -74,16 +85,18 @@ export class BridgeAdmissionQueue { } async leave(workerId: string, id: string): Promise { - const [order, deadlines] = this.keys(workerId); + const [order, deadlines, , workspaces] = this.keys(workerId); await this.redis.eval( [ "redis.call('ZREM', KEYS[1], ARGV[1])", "redis.call('ZREM', KEYS[2], ARGV[1])", + "redis.call('HDEL', KEYS[3], ARGV[1])", 'return 1', ].join('\n'), - 2, + 3, order, deadlines, + workspaces, id, ); } diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts new file mode 100644 index 00000000..b94aaece --- /dev/null +++ b/service/src/bridge/concurrent-store.test.ts @@ -0,0 +1,135 @@ +import { afterEach, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import { RedisBridgeStore } from './store'; +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import type { CodeBridgeAssignment } from './store'; + +const redis = new RedisMock() as unknown as Redis; +const store = new RedisBridgeStore(redis, 60, 1000, 2); +const workerId = 'concurrent-worker'; +const incarnationId = 'concurrent-incarnation'; +afterEach(async () => { + await redis.flushall(); +}); +async function register() { + const generation = await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + runtimes: [], + sandboxProfile: 'native-srt', + requiresReadyConfirmation: true, + workspaceLeaseSlots: 2, + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['read_file'], + workspaces: [{ id: 'a' }, { id: 'b' }], + }, + }, + }); + await store.confirmReady(workerId, incarnationId, generation); +} +function dispatch(workspaceId: string, signal = new AbortController().signal) { + const promise = store.dispatchWorkspaceTool({ + workerId, + signal, + deadlineAtMs: Date.now() + 3000, + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'read_file', + workspaceId, + path: 'file.txt', + }, + }); + void promise.catch(() => undefined); + return promise; +} +async function settle(assignment: CodeBridgeAssignment) { + await store.acknowledgeLease( + workerId, + incarnationId, + assignment.assignmentId, + assignment.generation, + assignment.leaseToken, + ); + await store.settle(workerId, assignment.assignmentId, { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected', + error: 'fixture clean rejection', + }); +} +test('store routes simultaneous roots through separate acknowledged slots', async () => { + await register(); + const a = dispatch('a'); + const b = dispatch('b'); + const first = await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ); + const second = await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 1, + ); + expect(first?.workspaceLeaseSlot).toBe(0); + expect(second?.workspaceLeaseSlot).toBe(1); + expect(first?.assignmentId).not.toBe(second?.assignmentId); + await settle(first!); + await settle(second!); + await expect(a).resolves.toMatchObject({ status: 'rejected' }); + await expect(b).resolves.toMatchObject({ status: 'rejected' }); + expect( + await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`), + ).toBeNull(); +}); + +test('same-root work waits while another root progresses', async () => { + await register(); + const a = dispatch('a'); + const first = await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ); + const nextA = dispatch('a'); + const b = dispatch('b'); + const second = await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 1, + ); + expect(second?.request).toMatchObject({ workspaceId: 'b' }); + await settle(first!); + await a; + const third = await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ); + expect(third?.request).toMatchObject({ workspaceId: 'a' }); + await settle(second!); + await settle(third!); + await Promise.all([nextA, b]); +}); diff --git a/service/src/bridge/concurrent-worker.test.ts b/service/src/bridge/concurrent-worker.test.ts new file mode 100644 index 00000000..3a359f07 --- /dev/null +++ b/service/src/bridge/concurrent-worker.test.ts @@ -0,0 +1,201 @@ +import { expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import { RedisBridgeStore } from './store'; +import { BridgeWorker } from '../../../packages/code/src/worker'; +import { WorkspaceToolError } from '../../../packages/code/src/workspace'; +import type { WorkspaceMutationQuarantine } from '../../../packages/code/src/worker'; +import type { BridgeWorkspaceToolCapabilities } from '../../../packages/code/src/protocol'; + +test('concurrent worker quarantines one root while another settles', async () => { + const redis = new RedisMock() as unknown as Redis; + const store = new RedisBridgeStore(redis, 60, 1000, 2); + const controller = new AbortController(); + const workerId = 'worker-concurrency'; + const incarnationId = 'incarnation-concurrency'; + const guards = new Map(); + const pending = new Set(); + for (const root of ['a', 'b']) + guards.set(root, { + async assertAvailable() { + if (pending.has(root)) throw new Error('quarantined'); + }, + async arm() { + pending.add(root); + }, + async clear() { + pending.delete(root); + }, + async quarantine() { + expect(pending.has(root)).toBe(true); + }, + }); + const capabilities: BridgeWorkspaceToolCapabilities = { + protocolVersion: 1, + operations: ['execute_command'], + workspaces: [{ id: 'a' }, { id: 'b' }], + }; + let startBoth!: () => void; + const bothStarted = new Promise((resolve) => { + startBoth = resolve; + }); + const started = new Set(); + const errors: unknown[] = []; + let registered!: () => void; + const ready = new Promise((resolve) => { + registered = resolve; + }); + const worker = new BridgeWorker({ + codeApiUrl: 'http://fixture.invalid', + token: 'fixture', + workerId, + incarnationId, + sandboxEndpoint: 'http://sandbox.invalid', + leaseWaitMs: 50, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + requiresReadyConfirmation: true, + workspaceLeaseSlots: 2, + workspaceTools: capabilities, + }, + workspaceQuarantines: guards, + onError: (error) => { + errors.push(error); + }, + workspaceTools: { + capabilities, + mutationFailuresAreAtomic: true, + async execute(request) { + started.add(request.workspaceId); + if (started.size === 2) startBoth(); + await bothStarted; + if (request.workspaceId === 'a') + throw new WorkspaceToolError( + 'uncertain command', + 'COMMAND_UNAVAILABLE', + true, + ); + return { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: 'b', + stdout: 'completed', + stderr: '', + exitCode: 0, + truncated: false, + timedOut: false, + }; + }, + }, + fetchImpl: (async (url, init) => { + const path = new URL(String(url)).pathname; + const body = JSON.parse(String(init?.body)); + const signal = init?.signal ?? undefined; + let result: object; + if (path.endsWith('/register')) { + const generation = await store.register(body); + result = { + protocolVersion: 1, + workerId, + incarnationId, + registrationGeneration: generation, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60000, + workspaceLeaseSlots: 2, + supportedWorkspaceToolOperations: ['execute_command'], + }; + } else if (path.endsWith('/ready')) { + await store.confirmReady( + workerId, + incarnationId, + body.registrationGeneration, + ); + registered(); + result = { protocolVersion: 1, ready: true }; + } else if (path.endsWith('/lease')) { + result = { + protocolVersion: 1, + serverElapsedMs: 0, + assignment: await store.lease( + workerId, + incarnationId, + body.waitMs, + signal, + undefined, + body.workspaceLeaseSlot, + ), + }; + } else { + const id = path.split('/').at(-2)!; + if (path.endsWith('/ack')) { + await store.acknowledgeLease( + workerId, + incarnationId, + id, + body.generation, + body.leaseToken, + signal, + ); + result = { protocolVersion: 1, accepted: true }; + } else if (path.endsWith('/cancellation')) { + result = { + protocolVersion: 1, + cancelled: await store.cancelled( + workerId, + incarnationId, + id, + signal, + ), + }; + } else { + await store.settle( + workerId, + id, + body, + signal, + undefined, + path.endsWith('/quarantine'), + ); + result = { protocolVersion: 1, accepted: true }; + } + } + return Response.json(result); + }) as typeof fetch, + }); + const running = worker.run(controller.signal); + void running.catch(() => undefined); + try { + await ready; + const results = await Promise.all( + ['a', 'b'].map((workspaceId) => + store.dispatchWorkspaceTool({ + workerId, + signal: controller.signal, + deadlineAtMs: Date.now() + 3000, + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId, + command: 'fixture', + }, + }), + ), + ); + expect(results[0]).toMatchObject({ status: 'rejected' }); + expect(results[1]).toMatchObject({ status: 'fulfilled' }); + expect([...pending]).toEqual(['a']); + expect(started.size).toBe(2); + } catch (error) { + throw new AggregateError( + [error, ...errors], + `Started roots: ${[...started].join(',')}`, + ); + } finally { + controller.abort(); + await running; + await redis.flushall(); + redis.disconnect(); + } +}); diff --git a/service/src/bridge/index.ts b/service/src/bridge/index.ts index 08ef1b07..fc409ad2 100644 --- a/service/src/bridge/index.ts +++ b/service/src/bridge/index.ts @@ -4,7 +4,12 @@ import { RedisBridgePairingStore } from './pairing'; import { createBridgeRouter } from './router'; import { RedisBridgeStore } from './store'; -export const bridgeStore = new RedisBridgeStore(connection); +export const bridgeStore = new RedisBridgeStore( + connection, + undefined, + undefined, + env.BRIDGE_MAX_WORKSPACE_LEASE_SLOTS, +); export const bridgePairings = new RedisBridgePairingStore(connection); export default createBridgeRouter({ diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index f428bdbe..b362e1b3 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -94,9 +94,9 @@ function sendStoreError(error: BridgeStoreError, res: Response): void { ? 404 : error.code === 'WORKER_UNAUTHORIZED' ? 403 - : error.code === 'WORKER_BUSY' - ? 503 - : 409; + : error.code === 'WORKER_BUSY' + ? 503 + : 409; res.status(status).json({ error: error.message, code: error.code }); } @@ -143,18 +143,16 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ?.match(/^Bearer\s+(.+)$/i)?.[1] ?.trim() ?? ''; - const adminAuth = ( - req: Request, - res: Response, - next: NextFunction, - ): void => { + const adminAuth = (req: Request, res: Response, next: NextFunction): void => { if (!options.adminToken) { res.status(503).json({ error: 'Code bridge is not configured' }); return; } const token = bearerToken(req); if (!token || !sameToken(token, options.adminToken)) { - res.status(401).json({ error: 'Invalid code bridge administrator token' }); + res + .status(401) + .json({ error: 'Invalid code bridge administrator token' }); return; } next(); @@ -228,73 +226,89 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { const workerAuth = options.authMode === 'paired' ? pairedWorkerAuth : staticWorkerAuth; - router.post('/pairings', adminAuth, asyncRoute(async (req, res) => { - if (options.authMode !== 'paired') { - res.status(409).json({ error: 'Paired worker authentication is disabled' }); - return; - } - const workerId = isRecord(req.body) ? req.body.workerId : undefined; - if ( - typeof workerId !== 'string' || - !validWorkerId(workerId) || - !configuredWorker(workerId) - ) { - res.status(400).json({ error: 'Invalid bridge worker ID' }); - return; - } - const hasBinding = isRecord(req.body) && - Object.prototype.hasOwnProperty.call(req.body, 'binding'); - const binding = isRecord(req.body) ? parseBinding(req.body.binding) : undefined; - if (hasBinding && binding == null) { - res.status(400).json({ error: 'Invalid bridge worker principal binding' }); - return; - } - if (options.allowDynamicWorkers === true && binding == null) { - res.status(400).json({ - error: 'Dynamic bridge workers require a valid principal binding', - }); - return; - } - const pairing = await options.pairings.issue(workerId, binding); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); - })); + router.post( + '/pairings', + adminAuth, + asyncRoute(async (req, res) => { + if (options.authMode !== 'paired') { + res + .status(409) + .json({ error: 'Paired worker authentication is disabled' }); + return; + } + const workerId = isRecord(req.body) ? req.body.workerId : undefined; + if ( + typeof workerId !== 'string' || + !validWorkerId(workerId) || + !configuredWorker(workerId) + ) { + res.status(400).json({ error: 'Invalid bridge worker ID' }); + return; + } + const hasBinding = + isRecord(req.body) && + Object.prototype.hasOwnProperty.call(req.body, 'binding'); + const binding = isRecord(req.body) + ? parseBinding(req.body.binding) + : undefined; + if (hasBinding && binding == null) { + res + .status(400) + .json({ error: 'Invalid bridge worker principal binding' }); + return; + } + if (options.allowDynamicWorkers === true && binding == null) { + res.status(400).json({ + error: 'Dynamic bridge workers require a valid principal binding', + }); + return; + } + const pairing = await options.pairings.issue(workerId, binding); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); + }), + ); - router.post('/pairings/redeem', asyncRoute(async (req, res) => { - if (options.authMode !== 'paired') { - res.status(409).json({ error: 'Paired worker authentication is disabled' }); - return; - } - const redemption = req.body as unknown; - if ( - !isRecord(redemption) || - redemption.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - typeof redemption.workerId !== 'string' || - !validWorkerId(redemption.workerId) || - !configuredWorker(redemption.workerId) || - typeof redemption.code !== 'string' || - redemption.code.length < 16 || - typeof redemption.publicKey !== 'string' || - redemption.publicKey.length > 4096 - ) { - res.status(400).json({ error: 'Invalid bridge pairing redemption' }); - return; - } - try { - const credential = await options.pairings.redeem({ - workerId: redemption.workerId, - code: redemption.code, - publicKey: redemption.publicKey, - }); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); - } catch (error) { - if (error instanceof BridgePairingError) { - const status = error.code === 'PUBLIC_KEY_INVALID' ? 400 : 401; - res.status(status).json({ error: error.message, code: error.code }); + router.post( + '/pairings/redeem', + asyncRoute(async (req, res) => { + if (options.authMode !== 'paired') { + res + .status(409) + .json({ error: 'Paired worker authentication is disabled' }); return; } - throw error; - } - })); + const redemption = req.body as unknown; + if ( + !isRecord(redemption) || + redemption.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof redemption.workerId !== 'string' || + !validWorkerId(redemption.workerId) || + !configuredWorker(redemption.workerId) || + typeof redemption.code !== 'string' || + redemption.code.length < 16 || + typeof redemption.publicKey !== 'string' || + redemption.publicKey.length > 4096 + ) { + res.status(400).json({ error: 'Invalid bridge pairing redemption' }); + return; + } + try { + const credential = await options.pairings.redeem({ + workerId: redemption.workerId, + code: redemption.code, + publicKey: redemption.publicKey, + }); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); + } catch (error) { + if (error instanceof BridgePairingError) { + const status = error.code === 'PUBLIC_KEY_INVALID' ? 400 : 401; + res.status(status).json({ error: error.message, code: error.code }); + return; + } + throw error; + } + }), + ); router.post( '/workers/:workerId/revoke', @@ -330,383 +344,397 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }), ); -router.post( - '/workers/register', - workerAuth, - asyncRoute(async (req, res) => { - const registration = req.body as unknown; - if ( - !isRecord(registration) || - registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - typeof registration.workerId !== 'string' || - !validWorkerId(registration.workerId) || - !validIncarnationId(registration.incarnationId) || - !isValidBridgeWorkerCapabilities(registration.capabilities) - ) { - res.status(400).json({ error: 'Invalid bridge worker registration' }); - return; - } - if ( - !configuredWorker(registration.workerId) - ) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - const authorization = options.authMode === 'paired' - ? ( - res.locals.bridgeWorkerAuthorization as { - identityId: string; - pairingGeneration: number; - credentialId: string; - activeCredentialId: string; - binding?: BridgeWorkerBinding; - } - ) - : undefined; - const trustedRegistration: BridgeWorkerRegistration & { - binding?: BridgeWorkerBinding; - credentialId?: string; - identityId?: string; - } = { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - workerId: registration.workerId, - incarnationId: registration.incarnationId, - capabilities: registration.capabilities, - ...(authorization?.credentialId != null - ? { credentialId: authorization.credentialId } - : {}), - ...(authorization?.identityId != null - ? { identityId: authorization.identityId } - : {}), - ...(authorization?.binding != null - ? { binding: authorization.binding } - : {}), - }; - try { - const registrationGeneration = await options.store.register( - trustedRegistration, - authorization, - ); - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - workerId: registration.workerId, - incarnationId: registration.incarnationId, - registrationGeneration, - registeredAt: new Date().toISOString(), - leaseTtlMs: 60_000, - supportedWorkspaceToolOperations: [ - 'read_file', - 'search_text', - 'list_files', - 'write_file', - 'preview_edit', - 'edit_file', - 'execute_command', - ], - supportedWorkspaceWriteFileModes: ['replace', 'create'], - supportedWorkspaceEditFileModes: ['single', 'batch'], - supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], - supportedWorkspaceListFileFeatures: ['after_path'], - }); - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + router.post( + '/workers/register', + workerAuth, + asyncRoute(async (req, res) => { + const registration = req.body as unknown; + if ( + !isRecord(registration) || + registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof registration.workerId !== 'string' || + !validWorkerId(registration.workerId) || + !validIncarnationId(registration.incarnationId) || + !isValidBridgeWorkerCapabilities(registration.capabilities) + ) { + res.status(400).json({ error: 'Invalid bridge worker registration' }); return; } - throw error; - } - }), -); - -router.post( - '/workers/:workerId/ready', - workerAuth, - asyncRoute(async (req, res) => { - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; - if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isSafeInteger(body.registrationGeneration) || - Number(body.registrationGeneration) < 1 - ) { - res.status(400).json({ - error: 'Invalid bridge worker readiness confirmation', - }); - return; - } - if (!configuredWorker(workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - try { - await options.store.confirmReady( - workerId, - body.incarnationId, - Number(body.registrationGeneration), - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ready: true }); - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + if (!configuredWorker(registration.workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); return; } - throw error; - } - }), -); - -router.post( - '/workers/:workerId/workspaces/reset', - workerAuth, - asyncRoute(async (req, res) => { - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; - if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - typeof body.runtimeSessionId !== 'string' || - body.runtimeSessionId.trim().length === 0 || - body.runtimeSessionId.length > 512 || - body.confirmDiscarded !== true - ) { - res.status(400).json({ - error: 'Workspace reset requires confirmation of local discard', - }); - return; - } - if (!configuredWorker(workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - try { - const resetController = new AbortController(); - const abortReset = (): void => resetController.abort(); - req.once('aborted', abortReset); - res.once('close', abortReset); + const authorization = + options.authMode === 'paired' + ? (res.locals.bridgeWorkerAuthorization as { + identityId: string; + pairingGeneration: number; + credentialId: string; + activeCredentialId: string; + binding?: BridgeWorkerBinding; + }) + : undefined; + const trustedRegistration: BridgeWorkerRegistration & { + binding?: BridgeWorkerBinding; + credentialId?: string; + identityId?: string; + } = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + capabilities: registration.capabilities, + ...(authorization?.credentialId != null + ? { credentialId: authorization.credentialId } + : {}), + ...(authorization?.identityId != null + ? { identityId: authorization.identityId } + : {}), + ...(authorization?.binding != null + ? { binding: authorization.binding } + : {}), + }; try { - await options.store.resetWorkspace( - workerId, - body.incarnationId, - body.runtimeSessionId, - resetController.signal, + const registrationGeneration = await options.store.register( + trustedRegistration, + authorization, ); - if (!resetController.signal.aborted) { - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + registrationGeneration, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + workspaceLeaseSlots: options.store.workspaceLeaseCapacity( + registration.capabilities.workspaceLeaseSlots, + ), + supportedWorkspaceToolOperations: [ + 'read_file', + 'search_text', + 'list_files', + 'write_file', + 'preview_edit', + 'edit_file', + 'execute_command', + ], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], + supportedWorkspaceListFileFeatures: ['after_path'], + }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; } - } finally { - req.off('aborted', abortReset); - res.off('close', abortReset); + throw error; } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + }), + ); + + router.post( + '/workers/:workerId/ready', + workerAuth, + asyncRoute(async (req, res) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isSafeInteger(body.registrationGeneration) || + Number(body.registrationGeneration) < 1 + ) { + res.status(400).json({ + error: 'Invalid bridge worker readiness confirmation', + }); + return; + } + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); return; } - throw error; - } - }), -); - -router.post( - '/workers/:workerId/lease', - workerAuth, - asyncRoute(async (req, res) => { - const requestStartedAtMs = Date.now(); - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; - const requestedWait = Number(body.waitMs ?? 25_000); - if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isFinite(requestedWait) || - requestedWait < 0 - ) { - res.status(400).json({ error: 'Invalid bridge lease request' }); - return; - } - if (!configuredWorker(workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - try { - const leaseController = new AbortController(); - const abortLease = (): void => leaseController.abort(); - req.once('aborted', abortLease); - res.once('close', abortLease); - let assignment: CodeBridgeAssignment | undefined; try { - assignment = await options.store.lease( + await options.store.confirmReady( workerId, body.incarnationId, - Math.min(requestedWait, MAX_LEASE_WAIT_MS), - leaseController.signal, - ( - res.locals.bridgeWorkerAuthorization as - | { identityId: string } - | undefined - )?.identityId, + Number(body.registrationGeneration), ); - if (leaseController.signal.aborted) { - if (assignment != null) await options.store.returnLease(assignment); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ready: true }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); return; } - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), - assignment, + throw error; + } + }), + ); + + router.post( + '/workers/:workerId/workspaces/reset', + workerAuth, + asyncRoute(async (req, res) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + typeof body.runtimeSessionId !== 'string' || + body.runtimeSessionId.trim().length === 0 || + body.runtimeSessionId.length > 512 || + body.confirmDiscarded !== true + ) { + res.status(400).json({ + error: 'Workspace reset requires confirmation of local discard', }); - } finally { - req.off('aborted', abortLease); - res.off('close', abortLease); + return; } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); return; } - throw error; - } - }), -); - -router.post( - '/workers/:workerId/assignments/:assignmentId/ack', - workerAuth, - asyncRoute(async (req, res) => { - const body = isRecord(req.body) ? req.body : {}; - if ( - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isSafeInteger(body.generation) || - Number(body.generation) < 1 || - typeof body.leaseToken !== 'string' || - body.leaseToken.length < 32 - ) { - res.status(400).json({ error: 'Invalid bridge lease acknowledgement' }); - return; - } - try { - const acknowledgementController = new AbortController(); - const abortAcknowledgement = (): void => - acknowledgementController.abort(); - req.once('aborted', abortAcknowledgement); - res.once('close', abortAcknowledgement); try { - await options.store.acknowledgeLease( - req.params.workerId, - body.incarnationId, - req.params.assignmentId, - Number(body.generation), - body.leaseToken, - acknowledgementController.signal, - ); - if (!acknowledgementController.signal.aborted) { - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, accepted: true }); + const resetController = new AbortController(); + const abortReset = (): void => resetController.abort(); + req.once('aborted', abortReset); + res.once('close', abortReset); + try { + await options.store.resetWorkspace( + workerId, + body.incarnationId, + body.runtimeSessionId, + resetController.signal, + ); + if (!resetController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); + } + } finally { + req.off('aborted', abortReset); + res.off('close', abortReset); } - } finally { - req.off('aborted', abortAcknowledgement); - res.off('close', abortAcknowledgement); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), + ); + + router.post( + '/workers/:workerId/lease', + workerAuth, + asyncRoute(async (req, res) => { + const requestStartedAtMs = Date.now(); + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + const requestedWait = Number(body.waitMs ?? 25_000); + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isFinite(requestedWait) || + requestedWait < 0 || + (body.workspaceLeaseSlot !== undefined && + (!Number.isSafeInteger(body.workspaceLeaseSlot) || + Number(body.workspaceLeaseSlot) < 0 || + Number(body.workspaceLeaseSlot) >= 8)) + ) { + res.status(400).json({ error: 'Invalid bridge lease request' }); + return; } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); return; } - throw error; - } - }), -); - -router.post( - '/workers/:workerId/assignments/:assignmentId/settle', - workerAuth, - asyncRoute(async (req, res) => { - const settlement = req.body as unknown; - if (!isSettlement(settlement)) { - res.status(400).json({ error: 'Invalid bridge settlement' }); - return; - } - try { - const settlementController = new AbortController(); - const abortSettlement = (): void => settlementController.abort(); - req.once('aborted', abortSettlement); - res.once('close', abortSettlement); try { - await options.store.settle( - req.params.workerId, - req.params.assignmentId, - settlement, - settlementController.signal, - ( - res.locals.bridgeWorkerAuthorization as - | { identityId: string } - | undefined - )?.identityId, - ); - if (!settlementController.signal.aborted) { + const leaseController = new AbortController(); + const abortLease = (): void => leaseController.abort(); + req.once('aborted', abortLease); + res.once('close', abortLease); + let assignment: CodeBridgeAssignment | undefined; + try { + assignment = await options.store.lease( + workerId, + body.incarnationId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + leaseController.signal, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, + body.workspaceLeaseSlot === undefined + ? undefined + : Number(body.workspaceLeaseSlot), + ); + if (leaseController.signal.aborted) { + if (assignment != null) await options.store.returnLease(assignment); + return; + } res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, - accepted: true, + serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), + assignment, }); + } finally { + req.off('aborted', abortLease); + res.off('close', abortLease); } - } finally { - req.off('aborted', abortSettlement); - res.off('close', abortSettlement); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + }), + ); + + router.post( + '/workers/:workerId/assignments/:assignmentId/ack', + workerAuth, + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isSafeInteger(body.generation) || + Number(body.generation) < 1 || + typeof body.leaseToken !== 'string' || + body.leaseToken.length < 32 + ) { + res.status(400).json({ error: 'Invalid bridge lease acknowledgement' }); return; } - throw error; - } - }), -); - -router.post( - '/workers/:workerId/assignments/:assignmentId/cancellation', - workerAuth, - asyncRoute(async (req, res) => { - const body = isRecord(req.body) ? req.body : {}; - if ( - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) - ) { - res.status(400).json({ error: 'Invalid bridge cancellation request' }); - return; - } - const cancellationController = new AbortController(); - const abortCancellation = (): void => cancellationController.abort(); - req.once('aborted', abortCancellation); - res.once('close', abortCancellation); - try { - const cancelled = await options.store.cancelled( - req.params.workerId, - body.incarnationId, - req.params.assignmentId, - cancellationController.signal, - ); - if (!cancellationController.signal.aborted) { - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); + try { + const acknowledgementController = new AbortController(); + const abortAcknowledgement = (): void => + acknowledgementController.abort(); + req.once('aborted', abortAcknowledgement); + res.once('close', abortAcknowledgement); + try { + await options.store.acknowledgeLease( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + Number(body.generation), + body.leaseToken, + acknowledgementController.signal, + ); + if (!acknowledgementController.signal.aborted) { + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + accepted: true, + }); + } + } finally { + req.off('aborted', abortAcknowledgement); + res.off('close', abortAcknowledgement); + } + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; } - } finally { - req.off('aborted', abortCancellation); - res.off('close', abortCancellation); - } - }), -); + }), + ); + + router.post( + [ + '/workers/:workerId/assignments/:assignmentId/settle', + '/workers/:workerId/assignments/:assignmentId/quarantine', + ], + workerAuth, + asyncRoute(async (req, res) => { + const settlement = req.body as unknown; + if (!isSettlement(settlement)) { + res.status(400).json({ error: 'Invalid bridge settlement' }); + return; + } + try { + const settlementController = new AbortController(); + const abortSettlement = (): void => settlementController.abort(); + req.once('aborted', abortSettlement); + res.once('close', abortSettlement); + try { + await options.store.settle( + req.params.workerId, + req.params.assignmentId, + settlement, + settlementController.signal, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, + req.path.endsWith('/quarantine'), + ); + if (!settlementController.signal.aborted) { + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + accepted: true, + }); + } + } finally { + req.off('aborted', abortSettlement); + res.off('close', abortSettlement); + } + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), + ); + + router.post( + '/workers/:workerId/assignments/:assignmentId/cancellation', + workerAuth, + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) + ) { + res.status(400).json({ error: 'Invalid bridge cancellation request' }); + return; + } + const cancellationController = new AbortController(); + const abortCancellation = (): void => cancellationController.abort(); + req.once('aborted', abortCancellation); + res.once('close', abortCancellation); + try { + const cancelled = await options.store.cancelled( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + cancellationController.signal, + ); + if (!cancellationController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); + } + } finally { + req.off('aborted', abortCancellation); + res.off('close', abortCancellation); + } + }), + ); router.post( '/workers/:workerId/credentials/refresh', @@ -732,6 +760,5 @@ router.post( }), ); - return router; } diff --git a/service/src/bridge/slots.test.ts b/service/src/bridge/slots.test.ts new file mode 100644 index 00000000..8694810c --- /dev/null +++ b/service/src/bridge/slots.test.ts @@ -0,0 +1,83 @@ +import { afterEach, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import { BridgeAdmissionQueue } from './admission'; +import { BridgeWorkspaceSlots } from './slots'; + +const redis = new RedisMock() as unknown as Redis; +const admission = new BridgeAdmissionQueue(redis); +const slots = new BridgeWorkspaceSlots(redis); +const workerId = 'slot-worker'; +const incarnationId = 'slot-incarnation'; +const prefix = `codeapi:bridge:v1:worker:${workerId}`; +afterEach(async () => { + await redis.flushall(); +}); +async function enqueue(assignmentId: string, workspaceId: string) { + await redis.set(`${prefix}:incarnation`, incarnationId); + await redis.set(`${prefix}:workspace-slot-capacity`, '2'); + await admission.enter(workerId, assignmentId, Date.now() + 5000, workspaceId); + return { + workerId, + incarnationId, + assignmentId, + workspaceId, + capacity: 2, + expiresAtMs: Date.now() + 10000, + }; +} + +test('slots admit independent workspaces, skip a busy root, and bound capacity', async () => { + const a = await enqueue('a', 'root-a'); + const a2 = await enqueue('a2', 'root-a'); + const b = await enqueue('b', 'root-b'); + const c = await enqueue('c', 'root-c'); + expect(await slots.reserve(b)).toBeUndefined(); + expect(await slots.reserve(a)).toBe(0); + expect(await slots.reserve(a2)).toBeUndefined(); + expect(await slots.reserve(b)).toBe(1); + expect(await slots.reserve(c)).toBeUndefined(); + expect(await slots.reserve(a)).toBe(0); + await slots.release(workerId, incarnationId, 'a'); + await admission.leave(workerId, 'a'); + expect(await slots.reserve(c)).toBeUndefined(); + expect(await slots.reserve(a2)).toBe(0); +}); + +test('stale slot release cannot erase replacement reservation or aggregate lock', async () => { + const a = await enqueue('a', 'root-a'); + expect(await slots.reserve(a)).toBe(0); + await slots.release(workerId, incarnationId, 'a'); + await admission.leave(workerId, 'a'); + const b = await enqueue('b', 'root-b'); + expect(await slots.reserve(b)).toBe(0); + await slots.release(workerId, incarnationId, 'a'); + expect(await redis.hlen(`${prefix}:workspace-slots`)).toBe(4); + expect(await redis.get(`${prefix}:lock`)).toBe( + `workspace-slots:${incarnationId}`, + ); + await slots.release(workerId, incarnationId, 'b'); + expect(await redis.get(`${prefix}:lock`)).toBeNull(); +}); + +test('legacy locks and older serial admission remain barriers', async () => { + const a = await enqueue('a', 'root-a'); + await redis.set(`${prefix}:lock`, 'legacy-assignment'); + expect(await slots.reserve(a)).toBeUndefined(); + await redis.del(`${prefix}:lock`); + await admission.leave(workerId, 'a'); + await admission.enter(workerId, 'legacy', Date.now() + 5000); + await admission.enter(workerId, 'a', Date.now() + 5000, 'root-a'); + expect(await slots.reserve(a)).toBeUndefined(); + await admission.leave(workerId, 'legacy'); + expect(await slots.reserve(a)).toBe(0); +}); + +test('slots reject invalid capacity and replaced incarnation', async () => { + const a = await enqueue('a', 'root-a'); + for (const capacity of [0, 9, 1.5, NaN]) { + await expect(slots.reserve({ ...a, capacity })).rejects.toThrow('Invalid'); + } + await redis.set(`${prefix}:incarnation`, 'replacement'); + await expect(slots.reserve(a)).rejects.toThrow('replaced'); +}); diff --git a/service/src/bridge/slots.ts b/service/src/bridge/slots.ts new file mode 100644 index 00000000..2a2b128d --- /dev/null +++ b/service/src/bridge/slots.ts @@ -0,0 +1,133 @@ +import type Redis from 'ioredis'; + +/** Hard bound keeps every atomic scheduling scan constant-sized. */ +export const MAX_WORKSPACE_LEASE_SLOTS = 8; + +/** Reservations share the legacy admission queue and aggregate worker lock. + * Thus a serial dispatcher or replacement incarnation cannot race active slots. + * Workspace mutation uncertainty is fenced separately by the assignment store. + */ +export class BridgeWorkspaceSlots { + constructor(private readonly redis: Redis) {} + + private keys(workerId: string): string[] { + const prefix = `codeapi:bridge:v1:worker:${encodeURIComponent(workerId)}`; + return [ + `${prefix}:workspace-slots`, + `${prefix}:lock`, + `${prefix}:lock:incarnation`, + `${prefix}:incarnation`, + `${prefix}:admission`, + `${prefix}:admission:deadlines`, + `${prefix}:admission:workspaces`, + `${prefix}:workspace-slot-capacity`, + ]; + } + + async reserve(args: { + workerId: string; + incarnationId: string; + assignmentId: string; + workspaceId: string; + capacity: number; + expiresAtMs: number; + }): Promise { + if ( + !Number.isSafeInteger(args.capacity) || + args.capacity < 1 || + args.capacity > MAX_WORKSPACE_LEASE_SLOTS || + !args.workspaceId || + !Number.isSafeInteger(args.expiresAtMs) || + args.expiresAtMs <= Date.now() + ) { + throw new Error('Invalid workspace slot reservation'); + } + const result = Number( + await this.redis.eval( + [ + "if redis.call('GET', KEYS[4]) ~= ARGV[1] then return -2 end", + "if (redis.call('GET', KEYS[8]) or '1') ~= ARGV[4] then return -2 end", + "local lock = redis.call('GET', KEYS[2])", + "local owner = 'workspace-slots:' .. ARGV[1]", + 'if lock and lock ~= owner then return -1 end', + 'local busy = {}', + 'local free = nil', + 'local latest = tonumber(ARGV[5])', + `for slot = 0, ${MAX_WORKSPACE_LEASE_SLOTS - 1} do`, + " local entry = redis.call('HMGET', KEYS[1], 'a:' .. slot, 'i:' .. slot, 'w:' .. slot, 'e:' .. slot)", + ' local occupied = entry[1]', + ' if occupied then', + ' if entry[2] ~= ARGV[1] or tonumber(entry[4]) <= tonumber(ARGV[6]) then', + " redis.call('HDEL', KEYS[1], 'a:' .. slot, 'i:' .. slot, 'w:' .. slot, 'e:' .. slot)", + ' occupied = false', + ' else', + ' if entry[1] == ARGV[2] then return slot end', + ' busy[entry[3]] = true', + ' latest = math.max(latest, tonumber(entry[4]))', + ' end', + ' end', + ' if not occupied and free == nil and slot < tonumber(ARGV[4]) then free = slot end', + 'end', + 'if free == nil or busy[ARGV[3]] then return -1 end', + // Expiry removes queue metadata, never a workspace uncertainty fence. + "local expired = redis.call('ZRANGEBYSCORE', KEYS[6], '-inf', ARGV[6])", + 'for _, id in ipairs(expired) do', + " redis.call('ZREM', KEYS[5], id)", + " redis.call('ZREM', KEYS[6], id)", + " redis.call('HDEL', KEYS[7], id)", + 'end', + "local pending = redis.call('ZRANGE', KEYS[5], 0, 31)", + 'local selected = nil', + 'for _, id in ipairs(pending) do', + " local workspace = redis.call('HGET', KEYS[7], id)", + // An older serial request remains a barrier until its dispatcher finishes. + ' if not workspace then return -1 end', + ' if not busy[workspace] then selected = id; break end', + 'end', + 'if selected ~= ARGV[2] then return -1 end', + "redis.call('HSET', KEYS[1], 'a:' .. free, ARGV[2], 'i:' .. free, ARGV[1], 'w:' .. free, ARGV[3], 'e:' .. free, ARGV[5])", + "redis.call('PEXPIREAT', KEYS[1], latest)", + "redis.call('SET', KEYS[2], owner, 'PXAT', latest)", + "redis.call('SET', KEYS[3], ARGV[1], 'PXAT', latest)", + 'return free', + ].join('\n'), + 8, + ...this.keys(args.workerId), + args.incarnationId, + args.assignmentId, + args.workspaceId, + args.capacity, + args.expiresAtMs, + Date.now(), + ), + ); + if (result === -2) + throw new Error('Workspace slot incarnation was replaced'); + return result < 0 ? undefined : result; + } + + async release( + workerId: string, + incarnationId: string, + assignmentId: string, + ): Promise { + await this.redis.eval( + [ + `for slot = 0, ${MAX_WORKSPACE_LEASE_SLOTS - 1} do`, + " local entry = redis.call('HMGET', KEYS[1], 'a:' .. slot, 'i:' .. slot)", + ' if entry[2] == ARGV[1] and entry[1] == ARGV[2] then', + " redis.call('HDEL', KEYS[1], 'a:' .. slot, 'i:' .. slot, 'w:' .. slot, 'e:' .. slot)", + ' end', + 'end', + "if redis.call('HLEN', KEYS[1]) == 0 and redis.call('GET', KEYS[2]) == 'workspace-slots:' .. ARGV[1] then", + " redis.call('DEL', KEYS[1], KEYS[2], KEYS[3])", + 'end', + 'return 1', + ].join('\n'), + 3, + ...this.keys(workerId).slice(0, 3), + incarnationId, + assignmentId, + ); + } +} diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index a00e4f2b..22747de6 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -19,6 +19,7 @@ import { } from '../../../packages/code/src/protocol'; import type { BridgeWorkerBinding } from './pairing'; import { BridgeAdmissionQueue } from './admission'; +import { BridgeWorkspaceSlots } from './slots'; const PREFIX = 'codeapi:bridge:v1'; const POLL_INTERVAL_MS = 100; @@ -33,7 +34,8 @@ export type CodeBridgeSettlement = BridgeSettlement< run?: t.ExecuteResponse['run']; } >; -export type CodeBridgeWorkspaceSettlement = BridgeSettlement; +export type CodeBridgeWorkspaceSettlement = + BridgeSettlement; type AnyCodeBridgeSettlement = | CodeBridgeSettlement | CodeBridgeWorkspaceSettlement; @@ -64,6 +66,11 @@ export class BridgeStoreError extends Error { interface StoredAssignment extends CodeBridgeAssignment { leaseTokenHash: string; workerIdentityId?: string; + workspaceFence?: string; +} + +function assignmentWorkspace(assignment: StoredAssignment): string | undefined { + return assignment.workspaceFence ?? assignment.runtimeSessionId; } export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { @@ -112,7 +119,8 @@ function supportsWorkspaceTool( ) { const mode = request.edits === undefined ? 'single' : 'batch'; const modes = capabilities?.editFileModes; - const supportsMode = modes == null ? mode === 'single' : modes.includes(mode); + const supportsMode = + modes == null ? mode === 'single' : modes.includes(mode); if (request.operation === 'preview_edit') return supportsMode; return ( supportsMode && @@ -173,16 +181,28 @@ function workspaceQuarantineKey( return `${PREFIX}:worker:${encodeURIComponent(workerId)}:workspace:${sessionHash}:quarantined`; } -function queueKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:assignments`; +function queueKey( + workerId: string, + incarnationId: string, + slot?: number, +): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:assignments${slot === undefined ? '' : `:slot:${slot}`}`; } -function leaseClaimKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-claim`; +function leaseClaimKey( + workerId: string, + incarnationId: string, + slot?: number, +): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-claim${slot === undefined ? '' : `:slot:${slot}`}`; } -function leaseAckKey(workerId: string, incarnationId: string): string { - return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-ack`; +function leaseAckKey( + workerId: string, + incarnationId: string, + slot?: number, +): string { + return `${PREFIX}:worker:${encodeURIComponent(workerId)}:incarnation:${incarnationId}:lease-ack${slot === undefined ? '' : `:slot:${slot}`}`; } function generationKey(workerId: string): string { @@ -283,7 +303,22 @@ export class RedisBridgeStore { private readonly redis: Redis, private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, private readonly redisCommandTimeoutMs = DEFAULT_REDIS_COMMAND_TIMEOUT_MS, - ) {} + private readonly maxWorkspaceLeaseSlots = 1, + ) { + if ( + !Number.isSafeInteger(maxWorkspaceLeaseSlots) || + maxWorkspaceLeaseSlots < 1 || + maxWorkspaceLeaseSlots > 8 + ) { + throw new Error( + 'Workspace lease slot ceiling must be an integer from 1 to 8', + ); + } + } + + workspaceLeaseCapacity(requested = 1): number { + return Math.min(this.maxWorkspaceLeaseSlots, requested); + } private async dispatchCommand( command: () => Promise, @@ -327,7 +362,7 @@ export class RedisBridgeStore { this.redis.eval( [ "local registration = redis.call('GET', KEYS[1])", - "if not registration then return { false, false, false, -2 } end", + 'if not registration then return { false, false, false, -2 } end', 'return {', ' registration,', " redis.call('GET', KEYS[2]) or false,", @@ -343,8 +378,17 @@ export class RedisBridgeStore { this.redisCommandTimeoutMs, 'Bridge worker status', )) as [string | null, string | null, string | null, number]; - const [rawRegistration, readyToken, registrationGeneration, leaseExpiresInMs] = snapshot; - if (rawRegistration == null || rawRegistration === '' || leaseExpiresInMs <= 0) { + const [ + rawRegistration, + readyToken, + registrationGeneration, + leaseExpiresInMs, + ] = snapshot; + if ( + rawRegistration == null || + rawRegistration === '' || + leaseExpiresInMs <= 0 + ) { return { online: false, ready: false }; } @@ -364,11 +408,16 @@ export class RedisBridgeStore { return { online: false, ready: false }; } - const requiresConfirmation = registration.capabilities.requiresReadyConfirmation === true; + const requiresConfirmation = + registration.capabilities.requiresReadyConfirmation === true; const ready = !requiresConfirmation || (registrationGeneration != null && - readyToken === workerReadyToken(registration.incarnationId, Number(registrationGeneration))); + readyToken === + workerReadyToken( + registration.incarnationId, + Number(registrationGeneration), + )); return { online: true, ready, @@ -379,12 +428,34 @@ export class RedisBridgeStore { async register( registration: RegisteredBridgeWorker, - authorization?: string | { - identityId?: string; - pairingGeneration?: number; - activeCredentialId?: string; - }, + authorization?: + | string + | { + identityId?: string; + pairingGeneration?: number; + activeCredentialId?: string; + }, ): Promise { + if (registration.capabilities.workspaceLeaseSlots !== undefined) { + if ( + !isValidBridgeWorkerCapabilities(registration.capabilities) || + registration.capabilities.requiresReadyConfirmation !== true + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Concurrent workspaces require readiness negotiation', + ); + } + registration = { + ...registration, + capabilities: { + ...registration.capabilities, + workspaceLeaseSlots: this.workspaceLeaseCapacity( + registration.capabilities.workspaceLeaseSlots, + ), + }, + }; + } const authorizationObject = typeof authorization === 'object' ? authorization : undefined; const expectedActiveCredentialId = @@ -396,12 +467,12 @@ export class RedisBridgeStore { ' local pairingGeneration = redis.call(\'GET\', KEYS[7]) or "0"', ' if pairingGeneration ~= ARGV[5] then return -5 end', ' if ARGV[6] ~= "" then', - ' if redis.call(\'GET\', KEYS[8]) ~= ARGV[6] then return -5 end', + " if redis.call('GET', KEYS[8]) ~= ARGV[6] then return -5 end", ' elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -5', ' end', 'end', 'if ARGV[8] ~= "" then', - ' local stableIdentity = redis.call(\'GET\', KEYS[8])', + " local stableIdentity = redis.call('GET', KEYS[8])", ' if stableIdentity and stableIdentity ~= ARGV[8] then return -4 end', ' if not stableIdentity then', ' if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4 end', @@ -409,29 +480,31 @@ export class RedisBridgeStore { ' end', 'elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4', 'end', - 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', - 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', - 'local current = redis.call(\'GET\', KEYS[4])', - 'if not current and redis.call(\'EXISTS\', KEYS[5]) == 1 then', - ' local owner = redis.call(\'GET\', KEYS[6])', + "if redis.call('EXISTS', KEYS[3]) == 1 then return -2 end", + "if redis.call('EXISTS', KEYS[2]) == 1 then return -1 end", + "local current = redis.call('GET', KEYS[4])", + "if current == ARGV[1] and redis.call('EXISTS', KEYS[5]) == 1 and (redis.call('GET', KEYS[13]) or \"1\") ~= ARGV[10] then return -3 end", + "if not current and redis.call('EXISTS', KEYS[5]) == 1 then", + " local owner = redis.call('GET', KEYS[6])", ' if owner ~= ARGV[1] then return -3 end', 'end', 'if current then', ' if current ~= ARGV[1] then', - ' if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -3 end', - ' redis.call(\'SET\', ARGV[4] .. current .. \':fenced\', \"1\")', + " if redis.call('EXISTS', KEYS[5]) == 1 then return -3 end", + " redis.call('SET', ARGV[4] .. current .. ':fenced', \"1\")", ' end', 'end', 'local registrationGeneration = tonumber(redis.call(\'GET\', KEYS[10]) or \"0\")', - 'local registrationGenerationIncarnation = redis.call(\'GET\', KEYS[11])', + "local registrationGenerationIncarnation = redis.call('GET', KEYS[11])", 'local registrationGenerationChanged = false', 'if registrationGeneration < 1 or registrationGenerationIncarnation ~= ARGV[1] then', - ' registrationGeneration = redis.call(\'INCR\', KEYS[10])', - ' redis.call(\'SET\', KEYS[11], ARGV[1])', + " registrationGeneration = redis.call('INCR', KEYS[10])", + " redis.call('SET', KEYS[11], ARGV[1])", ' registrationGenerationChanged = true', 'end', 'redis.call(\'SET\', KEYS[1], ARGV[2], \"EX\", ARGV[3])', 'redis.call(\'SET\', KEYS[4], ARGV[1], \"EX\", ARGV[3])', + 'redis.call(\'SET\', KEYS[13], ARGV[10], \"EX\", ARGV[3])', 'if ARGV[9] == "1" and registrationGenerationChanged then redis.call(\'DEL\', KEYS[12]) end', 'return registrationGeneration', ].join('\n'); @@ -439,9 +512,12 @@ export class RedisBridgeStore { await boundedCommand( this.redis.eval( script, - 12, + 13, workerKey(registration.workerId), - incarnationFenceKey(registration.workerId, registration.incarnationId), + incarnationFenceKey( + registration.workerId, + registration.incarnationId, + ), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), lockKey(registration.workerId), @@ -452,6 +528,7 @@ export class RedisBridgeStore { workerRegistrationGenerationKey(registration.workerId), workerRegistrationGenerationIncarnationKey(registration.workerId), workerReadyKey(registration.workerId), + `${PREFIX}:worker:${encodeURIComponent(registration.workerId)}:workspace-slot-capacity`, registration.incarnationId, JSON.stringify(registration), String(this.workerTtlSeconds), @@ -462,7 +539,10 @@ export class RedisBridgeStore { authorizationObject?.identityId ?? '', expectedActiveCredentialId ?? '', registration.identityId ?? '', - registration.capabilities.requiresReadyConfirmation === true ? '1' : '0', + registration.capabilities.requiresReadyConfirmation === true + ? '1' + : '0', + String(registration.capabilities.workspaceLeaseSlots ?? 1), ), this.redisCommandTimeoutMs, 'Bridge worker registration', @@ -499,7 +579,9 @@ export class RedisBridgeStore { ); } if (!Number.isSafeInteger(result) || result < 1) { - throw new Error('Bridge worker registration returned an invalid generation'); + throw new Error( + 'Bridge worker registration returned an invalid generation', + ); } return result; } @@ -513,12 +595,12 @@ export class RedisBridgeStore { await boundedCommand( this.redis.eval( [ - 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return -1 end', - 'if redis.call(\'GET\', KEYS[2]) ~= ARGV[1] then return -2 end', - 'if redis.call(\'GET\', KEYS[3]) ~= ARGV[2] then return -2 end', - 'if redis.call(\'GET\', KEYS[4]) ~= ARGV[1] then return -2 end', - 'if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -2 end', - 'if redis.call(\'EXISTS\', KEYS[6]) == 1 then return -3 end', + "if redis.call('EXISTS', KEYS[1]) == 0 then return -1 end", + "if redis.call('GET', KEYS[2]) ~= ARGV[1] then return -2 end", + "if redis.call('GET', KEYS[3]) ~= ARGV[2] then return -2 end", + "if redis.call('GET', KEYS[4]) ~= ARGV[1] then return -2 end", + "if redis.call('EXISTS', KEYS[5]) == 1 then return -2 end", + "if redis.call('EXISTS', KEYS[6]) == 1 then return -3 end", 'redis.call(\'SET\', KEYS[7], ARGV[3], "EX", ARGV[4])', 'return 1', ].join('\n'), @@ -622,12 +704,17 @@ export class RedisBridgeStore { registration: RegisteredBridgeWorker, ) => Promise; }): Promise { - if (args.executionTimeoutMs !== undefined && ( - args.workspaceRequest == null || - !Number.isSafeInteger(args.executionTimeoutMs) || - args.executionTimeoutMs < 1 || args.executionTimeoutMs > 305_000 - )) { - throw new BridgeStoreError('ASSIGNMENT_INVALID', 'Invalid workspace execution budget'); + if ( + args.executionTimeoutMs !== undefined && + (args.workspaceRequest == null || + !Number.isSafeInteger(args.executionTimeoutMs) || + args.executionTimeoutMs < 1 || + args.executionTimeoutMs > 305_000) + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_INVALID', + 'Invalid workspace execution budget', + ); } this.assertDispatchActive(args.signal, args.deadlineAtMs); const dispatchable = await this.dispatchCommand( @@ -692,44 +779,94 @@ export class RedisBridgeStore { const assignmentId = randomBytes(18).toString('base64url'); const leaseToken = randomBytes(32).toString('base64url'); // The lock is acquired before admission finishes; it must outlive the later execution deadline. - const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs + (args.executionTimeoutMs ?? 0)); + const ttlSeconds = assignmentTtlSeconds( + args.deadlineAtMs + (args.executionTimeoutMs ?? 0), + ); const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; + let workspaceLeaseSlot: number | undefined; + const workspaceSlots = + args.workspaceRequest != null && + (registration.capabilities.workspaceLeaseSlots ?? 1) > 1 + ? new BridgeWorkspaceSlots(this.redis) + : undefined; let resultCommitted = false; - const admission = args.workspaceRequest == null - ? undefined - : new BridgeAdmissionQueue(this.redis); + const admission = + args.workspaceRequest == null + ? undefined + : new BridgeAdmissionQueue(this.redis); try { - if (admission != null && !(await this.dispatchCommand( - () => admission.enter(args.workerId, assignmentId, args.deadlineAtMs), - args, - 'Bridge admission enqueue', - ))) { - throw new BridgeStoreError('WORKER_QUEUE_FULL', 'Bridge worker pending request limit reached'); - } - let locked = false; - do { - if (admission != null && !(await this.dispatchCommand( - () => admission.isHead(args.workerId, assignmentId), - args, - 'Bridge admission position', - ))) { - await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); - continue; - } - locked = await this.dispatchCommand( + if ( + admission != null && + !(await this.dispatchCommand( () => - this.acquireLock( + admission.enter( args.workerId, assignmentId, - lockIncarnationId, - ttlSeconds, + args.deadlineAtMs, + workspaceSlots == null + ? undefined + : args.workspaceRequest?.workspaceId, ), args, - 'Bridge assignment lock acquisition', + 'Bridge admission enqueue', + )) + ) { + throw new BridgeStoreError( + 'WORKER_QUEUE_FULL', + 'Bridge worker pending request limit reached', ); + } + let locked = false; + do { + if ( + admission != null && + workspaceSlots == null && + !(await this.dispatchCommand( + () => admission.isHead(args.workerId, assignmentId), + args, + 'Bridge admission position', + )) + ) { + await delay( + Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), + args.signal, + ); + continue; + } + if (workspaceSlots != null) { + workspaceLeaseSlot = await this.dispatchCommand( + () => + workspaceSlots.reserve({ + workerId: args.workerId, + incarnationId: lockIncarnationId, + assignmentId, + workspaceId: args.workspaceRequest!.workspaceId, + capacity: registration.capabilities.workspaceLeaseSlots!, + expiresAtMs: Date.now() + ttlSeconds * 1000, + }), + args, + 'Bridge workspace slot acquisition', + ); + locked = workspaceLeaseSlot !== undefined; + } else { + locked = await this.dispatchCommand( + () => + this.acquireLock( + args.workerId, + assignmentId, + lockIncarnationId, + ttlSeconds, + ), + args, + 'Bridge assignment lock acquisition', + ); + } if (!locked && admission != null) { - await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); + await delay( + Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), + args.signal, + ); } } while (!locked && admission != null); if (!locked) { @@ -750,12 +887,21 @@ export class RedisBridgeStore { current == null || current.registration.incarnationId !== registration.incarnationId || current.registration.identityId !== registration.identityId || - current.registration.binding?.tenantId !== registration.binding?.tenantId + current.registration.binding?.tenantId !== + registration.binding?.tenantId ) { - throw new BridgeStoreError('WORKER_OFFLINE', 'Bridge worker changed while the request was waiting'); + throw new BridgeStoreError( + 'WORKER_OFFLINE', + 'Bridge worker changed while the request was waiting', + ); } - if (!supportsWorkspaceTool(current.registration, args.workspaceRequest!)) { - throw new BridgeStoreError('WORKER_MISMATCH', 'Bridge worker capabilities changed while the request was waiting'); + if ( + !supportsWorkspaceTool(current.registration, args.workspaceRequest!) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Bridge worker capabilities changed while the request was waiting', + ); } } this.assertDispatchActive(args.signal, args.deadlineAtMs); @@ -775,6 +921,12 @@ export class RedisBridgeStore { generation, leaseToken, leaseTokenHash: tokenHash(leaseToken), + ...(workspaceLeaseSlot === undefined + ? {} + : { + workspaceLeaseSlot, + workspaceFence: `native-workspace:${args.workspaceRequest!.workspaceId}`, + }), ...(registration.identityId != null ? { workerIdentityId: registration.identityId } : {}), @@ -835,7 +987,10 @@ export class RedisBridgeStore { } if ( args.workspaceRequest != null && - !supportsWorkspaceTool(replacement.registration, args.workspaceRequest) + !supportsWorkspaceTool( + replacement.registration, + args.workspaceRequest, + ) ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -894,6 +1049,17 @@ export class RedisBridgeStore { } else { await this.cleanupDispatch(args.workerId, assignmentId, assignment); } + if (workspaceSlots != null && assignment == null) { + await boundedCommand( + workspaceSlots.release( + args.workerId, + lockIncarnationId, + assignmentId, + ), + this.redisCommandTimeoutMs, + 'Bridge unassigned slot cleanup', + ); + } } } @@ -903,18 +1069,31 @@ export class RedisBridgeStore { waitMs: number, signal?: AbortSignal, identityId?: string, + slot?: number, ): Promise { + if (slot !== undefined) { + const registration = await this.registration(workerId); + if ( + !Number.isSafeInteger(slot) || + slot < 0 || + slot >= (registration?.capabilities.workspaceLeaseSlots ?? 1) || + (registration?.capabilities.workspaceLeaseSlots ?? 1) <= 1 || + registration?.incarnationId !== incarnationId + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Workspace lease slot was not negotiated', + ); + } + } const deadline = Date.now() + waitMs; let firstPoll = true; - while ( - !signalAborted(signal) && - (firstPoll || Date.now() < deadline) - ) { + while (!signalAborted(signal) && (firstPoll || Date.now() < deadline)) { firstPoll = false; let assignmentId: string | null; try { assignmentId = await this.leaseCommand( - this.claimOrPopLease(workerId, incarnationId, identityId), + this.claimOrPopLease(workerId, incarnationId, identityId, slot), signal, 'Bridge lease claim', ); @@ -938,10 +1117,11 @@ export class RedisBridgeStore { if ( assignment == null || assignment.workerId !== workerId || - assignment.incarnationId !== incarnationId + assignment.incarnationId !== incarnationId || + assignment.workspaceLeaseSlot !== slot ) { await this.leaseCommand( - this.discardLeaseClaim(workerId, incarnationId, assignmentId), + this.discardLeaseClaim(workerId, incarnationId, assignmentId, slot), signal, 'Bridge lease claim discard', ); @@ -964,7 +1144,7 @@ export class RedisBridgeStore { } if (assignment.workerIdentityId !== identityId) { await this.leaseCommand( - this.discardLeaseClaim(workerId, incarnationId, assignmentId), + this.discardLeaseClaim(workerId, incarnationId, assignmentId, slot), signal, 'Bridge unauthorized lease discard', ); @@ -973,7 +1153,7 @@ export class RedisBridgeStore { if (Date.parse(assignment.expiresAt) <= Date.now()) { const acknowledged = (await this.leaseCommand( - this.redis.get(leaseAckKey(workerId, incarnationId)), + this.redis.get(leaseAckKey(workerId, incarnationId, slot)), signal, 'Bridge lease acknowledgement read', )) === assignmentId; @@ -985,7 +1165,7 @@ export class RedisBridgeStore { ); } await this.leaseCommand( - this.discardLeaseClaim(workerId, incarnationId, assignmentId), + this.discardLeaseClaim(workerId, incarnationId, assignmentId, slot), signal, 'Bridge expired lease discard', ); @@ -998,6 +1178,7 @@ export class RedisBridgeStore { const { leaseTokenHash: _leaseTokenHash, workerIdentityId: _workerIdentityId, + workspaceFence: _workspaceFence, ...wireAssignment } = assignment; return { @@ -1012,6 +1193,7 @@ export class RedisBridgeStore { workerId, incarnationId, assignmentId, + slot, ); if (signalAborted(signal)) return undefined; throw error; @@ -1061,8 +1243,8 @@ export class RedisBridgeStore { 'return 1', ].join('\n'), 2, - leaseClaimKey(workerId, incarnationId), - leaseAckKey(workerId, incarnationId), + leaseClaimKey(workerId, incarnationId, assignment.workspaceLeaseSlot), + leaseAckKey(workerId, incarnationId, assignment.workspaceLeaseSlot), assignmentId, String(ttlSeconds), ), @@ -1082,6 +1264,7 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, identityId?: string, + slot?: number, ): Promise { const result = await this.redis.eval( [ @@ -1099,8 +1282,8 @@ export class RedisBridgeStore { 'return assignment', ].join('\n'), 3, - queueKey(workerId, incarnationId), - leaseClaimKey(workerId, incarnationId), + queueKey(workerId, incarnationId, slot), + leaseClaimKey(workerId, incarnationId, slot), workerStableIdentityKey(workerId), identityId ?? '', ); @@ -1111,6 +1294,7 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, assignmentId: string, + slot?: number, ): Promise { await this.redis.eval( [ @@ -1120,8 +1304,8 @@ export class RedisBridgeStore { 'return 0', ].join('\n'), 2, - leaseClaimKey(workerId, incarnationId), - leaseAckKey(workerId, incarnationId), + leaseClaimKey(workerId, incarnationId, slot), + leaseAckKey(workerId, incarnationId, slot), assignmentId, ); } @@ -1131,6 +1315,7 @@ export class RedisBridgeStore { assignment.workerId, assignment.incarnationId, assignment.assignmentId, + assignment.workspaceLeaseSlot, ); } @@ -1138,6 +1323,7 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, assignmentId: string, + slot?: number, ): Promise { await boundedCommand( this.redis.eval( @@ -1153,9 +1339,9 @@ export class RedisBridgeStore { ].join('\n'), 4, assignmentKey(assignmentId), - queueKey(workerId, incarnationId), - leaseClaimKey(workerId, incarnationId), - leaseAckKey(workerId, incarnationId), + queueKey(workerId, incarnationId, slot), + leaseClaimKey(workerId, incarnationId, slot), + leaseAckKey(workerId, incarnationId, slot), assignmentId, ), this.redisCommandTimeoutMs, @@ -1167,11 +1353,12 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, assignmentId: string, + slot?: number, ): Promise { let lastError: unknown; for (let attempt = 0; attempt < 3; attempt += 1) { try { - await this.returnLeaseById(workerId, incarnationId, assignmentId); + await this.returnLeaseById(workerId, incarnationId, assignmentId, slot); return; } catch (error) { lastError = error; @@ -1184,7 +1371,7 @@ export class RedisBridgeStore { private async clearUndeliveredWorkspaceFence( assignment: StoredAssignment, ): Promise { - if (assignment.runtimeSessionId === undefined) return; + if (assignmentWorkspace(assignment) === undefined) return; await this.redis.eval( [ "if redis.call('GET', KEYS[1]) == ARGV[1] then", @@ -1195,7 +1382,7 @@ export class RedisBridgeStore { 1, workspaceQuarantineKey( assignment.workerId, - assignment.runtimeSessionId, + assignmentWorkspace(assignment)!, ), assignment.assignmentId, ); @@ -1207,6 +1394,7 @@ export class RedisBridgeStore { settlement: AnyCodeBridgeSettlement, signal?: AbortSignal, identityId?: string, + quarantineWorkspace = false, ): Promise { const serializedSettlement = JSON.stringify(settlement); const existingSettlement = await this.leaseCommand( @@ -1214,8 +1402,12 @@ export class RedisBridgeStore { signal, 'Bridge settlement existing read', ); - if (existingSettlement === serializedSettlement) return; - if (existingSettlement != null) { + if (existingSettlement === serializedSettlement && !quarantineWorkspace) + return; + if ( + existingSettlement != null && + existingSettlement !== serializedSettlement + ) { throw new BridgeStoreError( 'ASSIGNMENT_FENCED', 'Bridge assignment was already settled with a different result', @@ -1227,6 +1419,7 @@ export class RedisBridgeStore { 'Bridge settlement assignment read', ); if (assignment == null) { + if (existingSettlement === serializedSettlement) return; throw new BridgeStoreError( 'ASSIGNMENT_NOT_FOUND', 'Bridge assignment was not found', @@ -1238,6 +1431,17 @@ export class RedisBridgeStore { 'Bridge assignment belongs to another worker', ); } + if ( + quarantineWorkspace && + (assignment.workspaceFence == null || + assignment.workspaceLeaseSlot === undefined || + settlement.status !== 'rejected') + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_INVALID', + 'Quarantine requires a concurrent workspace assignment', + ); + } const registration = await this.leaseCommand( this.registration(workerId), signal, @@ -1268,38 +1472,49 @@ export class RedisBridgeStore { const settlementKeys = [ assignmentKey(assignmentId), settlementKey(assignmentId), - leaseClaimKey(workerId, assignment.incarnationId), - leaseAckKey(workerId, assignment.incarnationId), + leaseClaimKey( + workerId, + assignment.incarnationId, + assignment.workspaceLeaseSlot, + ), + leaseAckKey( + workerId, + assignment.incarnationId, + assignment.workspaceLeaseSlot, + ), assignmentDeadlineKey(assignmentId), ]; - if (assignment.runtimeSessionId !== undefined) { + if (assignmentWorkspace(assignment) !== undefined) { settlementKeys.push( - workspaceQuarantineKey(workerId, assignment.runtimeSessionId), + workspaceQuarantineKey(workerId, assignmentWorkspace(assignment)!), ); } - const hasWorkspace = assignment.runtimeSessionId !== undefined; + const hasWorkspace = assignmentWorkspace(assignment) !== undefined; settlementKeys.push( `${PREFIX}:stable-identity:${workerId}`, workerIncarnationKey(workerId), ); const script = [ - 'local existing = redis.call(\'GET\', KEYS[2])', + "local existing = redis.call('GET', KEYS[2])", 'if existing then', ' if existing == ARGV[1] then return 2 end', ' return -1', 'end', - 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return 0 end', + "if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end", 'if ARGV[6] == "1" and redis.call(\'GET\', KEYS[6]) ~= ARGV[3] then return -2 end', 'if ARGV[4] ~= "rejected" and redis.call(\'EXISTS\', KEYS[5]) == 0 then return -3 end', 'local stableIdentityKey = KEYS[#KEYS - 1]', 'if ARGV[5] ~= "" then', - ' if redis.call(\'GET\', stableIdentityKey) ~= ARGV[5] then return -4 end', - 'elseif redis.call(\'EXISTS\', stableIdentityKey) == 1 then return -4', + " if redis.call('GET', stableIdentityKey) ~= ARGV[5] then return -4 end", + "elseif redis.call('EXISTS', stableIdentityKey) == 1 then return -4", 'end', - 'if redis.call(\'GET\', KEYS[#KEYS]) ~= ARGV[7] then return -4 end', + "if redis.call('GET', KEYS[#KEYS]) ~= ARGV[7] then return -4 end", 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', - 'if redis.call(\'GET\', KEYS[3]) == ARGV[3] then redis.call(\'DEL\', KEYS[3], KEYS[4]) end', - 'if ARGV[6] == "1" and ARGV[4] == "rejected" then redis.call(\'DEL\', KEYS[6]) end', + "if redis.call('GET', KEYS[3]) == ARGV[3] then redis.call('DEL', KEYS[3], KEYS[4]) end", + 'if ARGV[6] == "1" and ARGV[4] == "rejected" then', + ' if ARGV[8] == "1" then redis.call(\'SET\', KEYS[6], "quarantined:" .. ARGV[3])', + " else redis.call('DEL', KEYS[6]) end", + 'end', 'return 1', ].join('\n'); const accepted = Number( @@ -1315,6 +1530,7 @@ export class RedisBridgeStore { identityId ?? '', hasWorkspace ? '1' : '0', settlement.incarnationId, + quarantineWorkspace ? '1' : '0', ), signal, 'Bridge settlement commit', @@ -1350,6 +1566,14 @@ export class RedisBridgeStore { 'Bridge assignment closed before settlement was committed', ); } + if ( + assignment.workspaceLeaseSlot !== undefined && + settlement.status === 'rejected' + ) { + // The dispatcher may already have timed out and finished its cleanup. + // Release only this settled reservation, retaining a quarantine marker. + await this.cleanupWithRetry(workerId, assignmentId, assignment); + } } async cancelled( @@ -1393,9 +1617,9 @@ export class RedisBridgeStore { const script = [ 'redis.call(\'SET\', KEYS[2], \"1\")', 'if #KEYS == 4 then redis.call(\'SET\', KEYS[4], \"1\") end', - 'local current = redis.call(\'GET\', KEYS[3])', + "local current = redis.call('GET', KEYS[3])", 'if current == ARGV[1] then', - ' return redis.call(\'DEL\', KEYS[1], KEYS[3])', + " return redis.call('DEL', KEYS[1], KEYS[3])", 'end', 'return 0', ].join('\n'); @@ -1408,12 +1632,7 @@ export class RedisBridgeStore { keys.push(workspaceQuarantineKey(workerId, runtimeSessionId)); } await boundedCommand( - this.redis.eval( - script, - keys.length, - ...keys, - incarnationId, - ), + this.redis.eval(script, keys.length, ...keys, incarnationId), this.redisCommandTimeoutMs, 'Bridge worker quarantine', ); @@ -1462,21 +1681,23 @@ export class RedisBridgeStore { workerId: string, ): Promise { const raw = await this.redis.get(workerKey(workerId)); - return raw == null ? undefined : (JSON.parse(raw) as RegisteredBridgeWorker); + return raw == null + ? undefined + : (JSON.parse(raw) as RegisteredBridgeWorker); } private async dispatchableRegistration( workerId: string, ): Promise< - | { registration: RegisteredBridgeWorker; readyToken?: string } - | undefined + { registration: RegisteredBridgeWorker; readyToken?: string } | undefined > { - const [raw, ready, generation, generationIncarnation] = await this.redis.mget( - workerKey(workerId), - workerReadyKey(workerId), - workerRegistrationGenerationKey(workerId), - workerRegistrationGenerationIncarnationKey(workerId), - ); + const [raw, ready, generation, generationIncarnation] = + await this.redis.mget( + workerKey(workerId), + workerReadyKey(workerId), + workerRegistrationGenerationKey(workerId), + workerRegistrationGenerationIncarnationKey(workerId), + ); if (raw == null) return undefined; const registration = JSON.parse(raw) as RegisteredBridgeWorker; if (registration.capabilities.requiresReadyConfirmation !== true) { @@ -1487,7 +1708,8 @@ export class RedisBridgeStore { !Number.isSafeInteger(registrationGeneration) || registrationGeneration < 1 || generationIncarnation !== registration.incarnationId || - ready !== workerReadyToken(registration.incarnationId, registrationGeneration) + ready !== + workerReadyToken(registration.incarnationId, registrationGeneration) ) { return undefined; } @@ -1543,11 +1765,11 @@ export class RedisBridgeStore { settlementKey(assignment.assignmentId), assignmentDeadlineKey(assignment.assignmentId), ]; - if (assignment.runtimeSessionId !== undefined) { + if (assignmentWorkspace(assignment) !== undefined) { closeKeys.push( workspaceQuarantineKey( assignment.workerId, - assignment.runtimeSessionId, + assignmentWorkspace(assignment)!, ), ); } @@ -1555,11 +1777,11 @@ export class RedisBridgeStore { // Keep acknowledged assignment metadata for late clean rejection recovery, // but atomically revoke fulfillment when no settlement has won yet. const closeScript = [ - 'local settlement = redis.call(\'GET\', KEYS[2])', + "local settlement = redis.call('GET', KEYS[2])", 'if settlement then return settlement end', - 'redis.call(\'DEL\', KEYS[3])', - 'if #KEYS == 4 and redis.call(\'GET\', KEYS[4]) == ARGV[1] then return nil end', - 'redis.call(\'DEL\', KEYS[1])', + "redis.call('DEL', KEYS[3])", + "if #KEYS == 4 and redis.call('GET', KEYS[4]) == ARGV[1] then return nil end", + "redis.call('DEL', KEYS[1])", 'return nil', ].join('\n'); const finalSettlement = await boundedCommand( @@ -1593,12 +1815,7 @@ export class RedisBridgeStore { ? 30 : assignmentTtlSeconds(Date.parse(assignment.expiresAt)); await boundedCommand( - this.redis.set( - cancellationKey(assignmentId), - '1', - 'EX', - ttlSeconds, - ), + this.redis.set(cancellationKey(assignmentId), '1', 'EX', ttlSeconds), this.redisCommandTimeoutMs, 'Bridge assignment cancellation', ); @@ -1610,30 +1827,36 @@ export class RedisBridgeStore { readyToken?: string, ): Promise { const script = [ - 'if redis.call(\'GET\', KEYS[1]) ~= ARGV[1] then return 0 end', + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", 'if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[6]) ~= ARGV[7] then return 0 end', - 'if #KEYS == 7 and redis.call(\'EXISTS\', KEYS[7]) == 1 then return -1 end', + "if #KEYS == 7 and redis.call('EXISTS', KEYS[7]) == 1 then return -1 end", 'redis.call(\'SET\', KEYS[2], ARGV[2], \"EX\", ARGV[3])', - 'redis.call(\'RPUSH\', KEYS[3], ARGV[4])', - 'redis.call(\'EXPIRE\', KEYS[3], ARGV[3])', - 'redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])', + "redis.call('RPUSH', KEYS[3], ARGV[4])", + "redis.call('EXPIRE', KEYS[3], ARGV[3])", + ...(assignment.workspaceLeaseSlot === undefined + ? ['redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])'] + : []), 'redis.call(\'SET\', KEYS[5], "1", \"PXAT\", ARGV[6])', - 'if #KEYS == 7 then redis.call(\'SET\', KEYS[7], ARGV[4]) end', + "if #KEYS == 7 then redis.call('SET', KEYS[7], ARGV[4]) end", 'return 1', ].join('\n'); const keys = [ workerIncarnationKey(assignment.workerId), assignmentKey(assignment.assignmentId), - queueKey(assignment.workerId, assignment.incarnationId), + queueKey( + assignment.workerId, + assignment.incarnationId, + assignment.workspaceLeaseSlot, + ), lockIncarnationKey(assignment.workerId), assignmentDeadlineKey(assignment.assignmentId), workerReadyKey(assignment.workerId), ]; - if (assignment.runtimeSessionId !== undefined) { + if (assignmentWorkspace(assignment) !== undefined) { keys.push( workspaceQuarantineKey( assignment.workerId, - assignment.runtimeSessionId, + assignmentWorkspace(assignment)!, ), ); } @@ -1665,7 +1888,7 @@ export class RedisBridgeStore { ttlSeconds: number, ): Promise { const script = [ - 'if redis.call(\'EXISTS\', KEYS[1]) == 1 then return 0 end', + "if redis.call('EXISTS', KEYS[1]) == 1 then return 0 end", 'redis.call(\'SET\', KEYS[1], ARGV[1], \"PX\", ARGV[3])', 'redis.call(\'SET\', KEYS[2], ARGV[2], \"PX\", ARGV[3])', 'return 1', @@ -1704,15 +1927,15 @@ export class RedisBridgeStore { settlement: AnyCodeBridgeSettlement, ): Promise { if ( - assignment.runtimeSessionId === undefined || + assignmentWorkspace(assignment) === undefined || settlement.status !== 'fulfilled' ) { return; } - const runtimeSessionId = assignment.runtimeSessionId; + const runtimeSessionId = assignmentWorkspace(assignment)!; const script = [ - 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', - ' return redis.call(\'DEL\', KEYS[1])', + "if redis.call('GET', KEYS[1]) == ARGV[1] then", + " return redis.call('DEL', KEYS[1])", 'end', 'return 0', ].join('\n'); @@ -1721,10 +1944,7 @@ export class RedisBridgeStore { this.redis.eval( script, 1, - workspaceQuarantineKey( - assignment.workerId, - runtimeSessionId, - ), + workspaceQuarantineKey(assignment.workerId, runtimeSessionId), assignment.assignmentId, ), // Once settlement wins, caller cancellation must not prevent its @@ -1762,14 +1982,26 @@ export class RedisBridgeStore { private async cleanup(assignment: StoredAssignment): Promise { const keys = [ assignmentKey(assignment.assignmentId), - queueKey(assignment.workerId, assignment.incarnationId), - leaseClaimKey(assignment.workerId, assignment.incarnationId), - leaseAckKey(assignment.workerId, assignment.incarnationId), - assignment.runtimeSessionId === undefined + queueKey( + assignment.workerId, + assignment.incarnationId, + assignment.workspaceLeaseSlot, + ), + leaseClaimKey( + assignment.workerId, + assignment.incarnationId, + assignment.workspaceLeaseSlot, + ), + leaseAckKey( + assignment.workerId, + assignment.incarnationId, + assignment.workspaceLeaseSlot, + ), + assignmentWorkspace(assignment) === undefined ? `${assignmentKey(assignment.assignmentId)}:no-workspace` : workspaceQuarantineKey( assignment.workerId, - assignment.runtimeSessionId, + assignmentWorkspace(assignment)!, ), ]; const cleanupScript = [ @@ -1798,7 +2030,7 @@ export class RedisBridgeStore { keys.length, ...keys, assignment.assignmentId, - assignment.runtimeSessionId === undefined ? '0' : '1', + assignmentWorkspace(assignment) === undefined ? '0' : '1', ), this.redisCommandTimeoutMs, 'Bridge assignment cleanup', @@ -1806,7 +2038,13 @@ export class RedisBridgeStore { ); if (cleanupResult !== -1) { await boundedCommand( - this.releaseLock(assignment.workerId, assignment.assignmentId), + assignment.workspaceLeaseSlot === undefined + ? this.releaseLock(assignment.workerId, assignment.assignmentId) + : new BridgeWorkspaceSlots(this.redis).release( + assignment.workerId, + assignment.incarnationId, + assignment.assignmentId, + ), this.redisCommandTimeoutMs, 'Bridge assignment lock release', ); @@ -1818,8 +2056,8 @@ export class RedisBridgeStore { assignmentId: string, ): Promise { const script = [ - 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', - ' return redis.call(\'DEL\', KEYS[1], KEYS[2])', + "if redis.call('GET', KEYS[1]) == ARGV[1] then", + " return redis.call('DEL', KEYS[1], KEYS[2])", 'end', 'return 0', ].join('\n'); diff --git a/service/src/config.ts b/service/src/config.ts index 53848103..d9f363f7 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -8,12 +8,35 @@ import { resolveExecutionProfileSource, } from './execution-profile'; -export const languageConfig: Record = { - [Languages.bash]: { language: 'bash', version: '5.2.0', fileName: 'script.sh' }, - [Languages.js]: { language: 'bun-js', version: '1.3.14', fileName: 'index.js' }, - [Languages.node]: { language: 'node', version: '24.15.0', fileName: 'index.js' }, - [Languages.py]: { language: 'python', version: '3.14.4', fileName: 'main.py' }, - [Languages.ts]: { language: 'bun-ts', version: '1.3.14', fileName: 'main.ts' }, +export const languageConfig: Record< + Languages | string, + t.LanguageConfig | undefined +> = { + [Languages.bash]: { + language: 'bash', + version: '5.2.0', + fileName: 'script.sh', + }, + [Languages.js]: { + language: 'bun-js', + version: '1.3.14', + fileName: 'index.js', + }, + [Languages.node]: { + language: 'node', + version: '24.15.0', + fileName: 'index.js', + }, + [Languages.py]: { + language: 'python', + version: '3.14.4', + fileName: 'main.py', + }, + [Languages.ts]: { + language: 'bun-ts', + version: '1.3.14', + fileName: 'main.ts', + }, }; const languageAliases: Record = { @@ -49,8 +72,12 @@ export function resolveLanguage(lang: string): Languages | undefined { } const defaultJobTimeoutMs = Number(process.env.JOB_TIMEOUT) || 300000; -const defaultMaxFileSize = Number(process.env.MAX_FILE_SIZE) || 25 * 1024 * 1024; -const defaultExecutionManifestTtlSeconds = Math.min(Math.ceil((defaultJobTimeoutMs + 60000) / 1000), 600); +const defaultMaxFileSize = + Number(process.env.MAX_FILE_SIZE) || 25 * 1024 * 1024; +const defaultExecutionManifestTtlSeconds = Math.min( + Math.ceil((defaultJobTimeoutMs + 60000) / 1000), + 600, +); const EGRESS_GRANT_GRACE_MS = 10 * 60 * 1000; /** Object-store listing and marker writes are metadata operations, not * checkpoint transfers. Bound each tightly so the post-exec checkpoint @@ -81,11 +108,13 @@ export function checkpointPipelineBudgetMs( checkpointTimeoutMs, CHECKPOINT_METADATA_TIMEOUT_CAP_MS, ); - return launchTimeoutMs - + 2 * checkpointTimeoutMs - + 2 * metadataTimeoutMs - + POST_EXEC_CHECKPOINT_REGISTRY_COMMANDS - * RUNTIME_SESSION_REDIS_COMMAND_TIMEOUT_MS; + return ( + launchTimeoutMs + + 2 * checkpointTimeoutMs + + 2 * metadataTimeoutMs + + POST_EXEC_CHECKPOINT_REGISTRY_COMMANDS * + RUNTIME_SESSION_REDIS_COMMAND_TIMEOUT_MS + ); } /** BullMQ's `timestamp` is the enqueue time. Anchor the worker deadline to it @@ -112,15 +141,20 @@ export function jobCompletionWaitTimeoutMs( backendCleanupTimeoutMs: number, egressRevokeTimeoutMs: number, ): number { - return jobTimeoutMs - + backendCleanupTimeoutMs - + egressRevokeTimeoutMs - + WORKER_COMPLETION_OVERHEAD_MS; + return ( + jobTimeoutMs + + backendCleanupTimeoutMs + + egressRevokeTimeoutMs + + WORKER_COMPLETION_OVERHEAD_MS + ); } export function parseArnList(raw: string | undefined): string[] | undefined { if (raw == null) return undefined; - const entries = raw.split(',').map((entry) => entry.trim()).filter((entry) => entry.length > 0); + const entries = raw + .split(',') + .map((entry) => entry.trim()) + .filter((entry) => entry.length > 0); return entries.length > 0 ? entries : undefined; } @@ -143,7 +177,10 @@ interface IntegerRange { max?: number; } -const lambdaMicrovmNumericRanges: Record = { +const lambdaMicrovmNumericRanges: Record< + keyof LambdaMicrovmNumericConfig, + IntegerRange +> = { LAMBDA_MICROVM_PORT: { min: 1, max: 65_535 }, LAMBDA_MICROVM_MAX_DURATION_SECONDS: { min: 1, max: 28_800 }, LAMBDA_MICROVM_IDLE_SECONDS: { min: 60, max: 28_800 }, @@ -180,14 +217,20 @@ export function resolveLambdaMicrovmNumericConfig( ): LambdaMicrovmNumericConfig { const read = (name: keyof LambdaMicrovmNumericConfig): number => { const raw = source[name]; - return raw == null || raw.trim() === '' ? lambdaMicrovmNumericDefaults[name] : Number(raw); + return raw == null || raw.trim() === '' + ? lambdaMicrovmNumericDefaults[name] + : Number(raw); }; return { LAMBDA_MICROVM_PORT: read('LAMBDA_MICROVM_PORT'), - LAMBDA_MICROVM_MAX_DURATION_SECONDS: read('LAMBDA_MICROVM_MAX_DURATION_SECONDS'), + LAMBDA_MICROVM_MAX_DURATION_SECONDS: read( + 'LAMBDA_MICROVM_MAX_DURATION_SECONDS', + ), LAMBDA_MICROVM_IDLE_SECONDS: read('LAMBDA_MICROVM_IDLE_SECONDS'), LAMBDA_MICROVM_SUSPEND_SECONDS: read('LAMBDA_MICROVM_SUSPEND_SECONDS'), - LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: read('LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS'), + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: read( + 'LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS', + ), LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: read('LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS'), LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: read('LAMBDA_MICROVM_HEALTH_TIMEOUT_MS'), LAMBDA_MICROVM_LAUNCH_TPS: read('LAMBDA_MICROVM_LAUNCH_TPS'), @@ -199,18 +242,28 @@ export function resolveLambdaMicrovmNumericConfig( export function lambdaMicrovmNumericConfigError( config: LambdaMicrovmNumericConfig, ): string | undefined { - for (const name of Object.keys(lambdaMicrovmNumericRanges) as Array) { + for (const name of Object.keys(lambdaMicrovmNumericRanges) as Array< + keyof LambdaMicrovmNumericConfig + >) { const value = config[name]; const { min, max } = lambdaMicrovmNumericRanges[name]; - if (!Number.isSafeInteger(value) || value < min || (max != null && value > max)) { - const range = max == null ? `at least ${min}` : `between ${min} and ${max}`; + if ( + !Number.isSafeInteger(value) || + value < min || + (max != null && value > max) + ) { + const range = + max == null ? `at least ${min}` : `between ${min} and ${max}`; return `${name} must be a whole number ${range}`; } } return undefined; } -export function resolvePositiveIntEnv(raw: string | undefined, defaultValue: number): number { +export function resolvePositiveIntEnv( + raw: string | undefined, + defaultValue: number, +): number { if (raw == null || raw.trim() === '') { return defaultValue; } @@ -221,8 +274,14 @@ export function resolvePositiveIntEnv(raw: string | undefined, defaultValue: num return parsed; } -export function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, jobTimeoutMs: number): number { - const defaultTtlSeconds = Math.max(1, Math.ceil((jobTimeoutMs + EGRESS_GRANT_GRACE_MS) / 1000)); +export function resolveEgressGrantTtlSeconds( + rawTtlSeconds: string | undefined, + jobTimeoutMs: number, +): number { + const defaultTtlSeconds = Math.max( + 1, + Math.ceil((jobTimeoutMs + EGRESS_GRANT_GRACE_MS) / 1000), + ); if (rawTtlSeconds == null || rawTtlSeconds.trim() === '') { return defaultTtlSeconds; } @@ -235,12 +294,17 @@ export function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, return Math.max(1, Math.ceil(configuredTtlSeconds)); } -const lambdaMicrovmNumericConfig = resolveLambdaMicrovmNumericConfig(process.env); +const lambdaMicrovmNumericConfig = resolveLambdaMicrovmNumericConfig( + process.env, +); export function hostedAppOperationTimeoutMs(): number { - return env.CHECKPOINT_TIMEOUT_MS * 9 - + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS * 7 - + env.HOSTED_APP_START_TIMEOUT_MS + 30_000; + return ( + env.CHECKPOINT_TIMEOUT_MS * 9 + + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS * 7 + + env.HOSTED_APP_START_TIMEOUT_MS + + 30_000 + ); } function configuredNumber(raw: string | undefined, fallback: number): number { @@ -261,39 +325,41 @@ function configuredChoice( export function resolveSandboxBackend( raw: string | undefined, ): 'http' | 'lambda-microvm' | 'remote-bridge' { - return configuredChoice( - raw, - 'CODEAPI_SANDBOX_BACKEND', + return configuredChoice(raw, 'CODEAPI_SANDBOX_BACKEND', 'http', [ 'http', - ['http', 'lambda-microvm', 'remote-bridge'], - ); + 'lambda-microvm', + 'remote-bridge', + ]); } export function resolveRuntimeSessionMode( raw: string | undefined, ): 'stateless' | 'affinity' | 'strict' { - return configuredChoice( - raw, - 'CODEAPI_RUNTIME_SESSION_MODE', + return configuredChoice(raw, 'CODEAPI_RUNTIME_SESSION_MODE', 'stateless', [ 'stateless', - ['stateless', 'affinity', 'strict'], - ); + 'affinity', + 'strict', + ]); } export function resolveBridgeAuthMode( raw: string | undefined, ): 'static' | 'paired' { - return configuredChoice( - raw, - 'CODEAPI_BRIDGE_AUTH_MODE', + return configuredChoice(raw, 'CODEAPI_BRIDGE_AUTH_MODE', 'static', [ 'static', - ['static', 'paired'], - ); + 'paired', + ]); } -const sandboxBackend = resolveSandboxBackend(process.env.CODEAPI_SANDBOX_BACKEND); -const runtimeSessionMode = resolveRuntimeSessionMode(process.env.CODEAPI_RUNTIME_SESSION_MODE); -const bridgeAuthMode = resolveBridgeAuthMode(process.env.CODEAPI_BRIDGE_AUTH_MODE); +const sandboxBackend = resolveSandboxBackend( + process.env.CODEAPI_SANDBOX_BACKEND, +); +const runtimeSessionMode = resolveRuntimeSessionMode( + process.env.CODEAPI_RUNTIME_SESSION_MODE, +); +const bridgeAuthMode = resolveBridgeAuthMode( + process.env.CODEAPI_BRIDGE_AUTH_MODE, +); export const env = { PORT: process.env.SERVICE_PORT ?? 3112, @@ -301,26 +367,61 @@ export const env = { HARDENED_SANDBOX_MODE: process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', INSTANCE_ID: process.env.INSTANCE_ID ?? nanoid(), HTTP_JSON_LIMIT: process.env.CODEAPI_HTTP_JSON_LIMIT ?? '50mb', - SANDBOX_ENDPOINT: process.env.SANDBOX_ENDPOINT ?? 'http://localhost:2000/api/v2', + SANDBOX_ENDPOINT: + process.env.SANDBOX_ENDPOINT ?? 'http://localhost:2000/api/v2', EGRESS_GATEWAY_URL: process.env.EGRESS_GATEWAY_URL ?? '', FILE_SERVER_URL: process.env.FILE_SERVER_URL ?? 'http://localhost:3000', - TOOL_CALL_SERVER_URL: process.env.TOOL_CALL_SERVER_URL ?? 'http://localhost:3033', + TOOL_CALL_SERVER_URL: + process.env.TOOL_CALL_SERVER_URL ?? 'http://localhost:3033', EGRESS_GATEWAY_PORT: Number(process.env.EGRESS_GATEWAY_PORT) || 3190, - EGRESS_GATEWAY_FILE_SERVER_URL: process.env.EGRESS_GATEWAY_FILE_SERVER_URL ?? process.env.FILE_SERVER_URL ?? 'http://localhost:3000', - EGRESS_GATEWAY_TOOL_CALL_SERVER_URL: process.env.EGRESS_GATEWAY_TOOL_CALL_SERVER_URL ?? process.env.TOOL_CALL_SERVER_URL ?? 'http://localhost:3033', - EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES: Number(process.env.EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES) || 1024 * 1024, + EGRESS_GATEWAY_FILE_SERVER_URL: + process.env.EGRESS_GATEWAY_FILE_SERVER_URL ?? + process.env.FILE_SERVER_URL ?? + 'http://localhost:3000', + EGRESS_GATEWAY_TOOL_CALL_SERVER_URL: + process.env.EGRESS_GATEWAY_TOOL_CALL_SERVER_URL ?? + process.env.TOOL_CALL_SERVER_URL ?? + 'http://localhost:3033', + EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES: + Number(process.env.EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES) || 1024 * 1024, // Per-entry / aggregate caps for PTC tool results persisted in `tool_history:` (see replay-state.ts). - PTC_MAX_TOOL_RESULT_BYTES: resolvePositiveIntEnv(process.env.PTC_MAX_TOOL_RESULT_BYTES, 5_000_000), - PTC_MAX_TOOL_HISTORY_TOTAL_BYTES: resolvePositiveIntEnv(process.env.PTC_MAX_TOOL_HISTORY_TOTAL_BYTES, 40_000_000), - EGRESS_GATEWAY_MAX_FILE_BYTES: Number(process.env.EGRESS_GATEWAY_MAX_FILE_BYTES ?? process.env.SANDBOX_MAX_FILE_SIZE) || 10_000_000, - EGRESS_GATEWAY_MAX_PATH_LENGTH: Number(process.env.EGRESS_GATEWAY_MAX_PATH_LENGTH ?? process.env.SANDBOX_MAX_PATH_LENGTH) || 256, - EGRESS_GATEWAY_MAX_NESTING_DEPTH: Number(process.env.EGRESS_GATEWAY_MAX_NESTING_DEPTH ?? process.env.SANDBOX_MAX_NESTING_DEPTH) || 10, - EGRESS_GATEWAY_REQUEST_TIMEOUT_MS: Number(process.env.EGRESS_GATEWAY_REQUEST_TIMEOUT_MS) || 30_000, - EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: Number(process.env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS) || 5_000, - EGRESS_LEDGER_REQUIRED: process.env.CODEAPI_EGRESS_LEDGER_REQUIRED === 'true' || process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', - EGRESS_LEDGER_TTL_GRACE_SECONDS: Number(process.env.CODEAPI_EGRESS_LEDGER_TTL_GRACE_SECONDS) || 300, + PTC_MAX_TOOL_RESULT_BYTES: resolvePositiveIntEnv( + process.env.PTC_MAX_TOOL_RESULT_BYTES, + 5_000_000, + ), + PTC_MAX_TOOL_HISTORY_TOTAL_BYTES: resolvePositiveIntEnv( + process.env.PTC_MAX_TOOL_HISTORY_TOTAL_BYTES, + 40_000_000, + ), + EGRESS_GATEWAY_MAX_FILE_BYTES: + Number( + process.env.EGRESS_GATEWAY_MAX_FILE_BYTES ?? + process.env.SANDBOX_MAX_FILE_SIZE, + ) || 10_000_000, + EGRESS_GATEWAY_MAX_PATH_LENGTH: + Number( + process.env.EGRESS_GATEWAY_MAX_PATH_LENGTH ?? + process.env.SANDBOX_MAX_PATH_LENGTH, + ) || 256, + EGRESS_GATEWAY_MAX_NESTING_DEPTH: + Number( + process.env.EGRESS_GATEWAY_MAX_NESTING_DEPTH ?? + process.env.SANDBOX_MAX_NESTING_DEPTH, + ) || 10, + EGRESS_GATEWAY_REQUEST_TIMEOUT_MS: + Number(process.env.EGRESS_GATEWAY_REQUEST_TIMEOUT_MS) || 30_000, + EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: + Number(process.env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS) || 5_000, + EGRESS_LEDGER_REQUIRED: + process.env.CODEAPI_EGRESS_LEDGER_REQUIRED === 'true' || + process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', + EGRESS_LEDGER_TTL_GRACE_SECONDS: + Number(process.env.CODEAPI_EGRESS_LEDGER_TTL_GRACE_SECONDS) || 300, EGRESS_GRANT_SECRET: process.env.CODEAPI_EGRESS_GRANT_SECRET ?? '', - EGRESS_GRANT_TTL_SECONDS: resolveEgressGrantTtlSeconds(process.env.EGRESS_GRANT_TTL_SECONDS, defaultJobTimeoutMs), + EGRESS_GRANT_TTL_SECONDS: resolveEgressGrantTtlSeconds( + process.env.EGRESS_GRANT_TTL_SECONDS, + defaultJobTimeoutMs, + ), PYTHON_CONCURRENCY: Number(process.env.PYTHON_CONCURRENCY) || 1, OTHER_CONCURRENCY: Number(process.env.OTHER_CONCURRENCY) || 8, JOB_WINDOW: Number(process.env.JOB_WINDOW) || 1000, @@ -359,22 +460,32 @@ export const env = { * working; multi-tenant deploys MUST set this to `true` before any tenant * is multi-homed, otherwise a missing tenantId would silently bucket * cross-tenant requests under the same `'legacy'` prefix. */ - TENANT_ISOLATION_STRICT: process.env.CODEAPI_TENANT_ISOLATION_STRICT === 'true', + TENANT_ISOLATION_STRICT: + process.env.CODEAPI_TENANT_ISOLATION_STRICT === 'true', // Signed execution manifests. Prefer private/public key mode for split-runner // deployments so sandbox-runner receives only a verifier, not a signing secret. - EXECUTION_MANIFEST_PRIVATE_KEY: process.env.CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY ?? '', - EXECUTION_MANIFEST_PUBLIC_KEY: process.env.CODEAPI_EXECUTION_MANIFEST_PUBLIC_KEY ?? '', + EXECUTION_MANIFEST_PRIVATE_KEY: + process.env.CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY ?? '', + EXECUTION_MANIFEST_PUBLIC_KEY: + process.env.CODEAPI_EXECUTION_MANIFEST_PUBLIC_KEY ?? '', // Legacy HMAC fallback for non-split deployments. Do not mount into sandbox-runner. - EXECUTION_MANIFEST_SECRET: process.env.CODEAPI_EXECUTION_MANIFEST_SECRET ?? '', + EXECUTION_MANIFEST_SECRET: + process.env.CODEAPI_EXECUTION_MANIFEST_SECRET ?? '', EXECUTION_MANIFEST_TTL_SECONDS: Math.min( - Number(process.env.EXECUTION_MANIFEST_TTL_SECONDS) || defaultExecutionManifestTtlSeconds, + Number(process.env.EXECUTION_MANIFEST_TTL_SECONDS) || + defaultExecutionManifestTtlSeconds, 600, ), - EXECUTION_MANIFEST_MAX_UPLOAD_BYTES: Number(process.env.EXECUTION_MANIFEST_MAX_UPLOAD_BYTES) || defaultMaxFileSize, - EXECUTION_MANIFEST_MAX_OUTPUT_FILES: Number(process.env.EXECUTION_MANIFEST_MAX_OUTPUT_FILES) || 50, - EXECUTION_MANIFEST_MAX_REQUESTS: Number(process.env.EXECUTION_MANIFEST_MAX_REQUESTS) || 1000, + EXECUTION_MANIFEST_MAX_UPLOAD_BYTES: + Number(process.env.EXECUTION_MANIFEST_MAX_UPLOAD_BYTES) || + defaultMaxFileSize, + EXECUTION_MANIFEST_MAX_OUTPUT_FILES: + Number(process.env.EXECUTION_MANIFEST_MAX_OUTPUT_FILES) || 50, + EXECUTION_MANIFEST_MAX_REQUESTS: + Number(process.env.EXECUTION_MANIFEST_MAX_REQUESTS) || 1000, // Redis - Alternative DNS Lookup for AWS ElastiCache TLS connections - REDIS_USE_ALTERNATIVE_DNS_LOOKUP: process.env.REDIS_USE_ALTERNATIVE_DNS_LOOKUP === 'true', + REDIS_USE_ALTERNATIVE_DNS_LOOKUP: + process.env.REDIS_USE_ALTERNATIVE_DNS_LOOKUP === 'true', /** * Programmatic Tool Calling execution model. * - `replay` (default): Temporal-style replay. Sandbox exits between round-trips; @@ -385,7 +496,9 @@ export const env = { * via a long-polling HTTP callback through the Tool Call Server. Retained as * an explicit opt-in during rollout; scheduled for removal in a follow-up. */ - PTC_MODE: (process.env.PTC_MODE === 'blocking' ? 'blocking' : 'replay') as 'replay' | 'blocking', + PTC_MODE: (process.env.PTC_MODE === 'blocking' ? 'blocking' : 'replay') as + | 'replay' + | 'blocking', PTC_DEBUG: process.env.PTC_DEBUG === 'true', /** * Sandbox execution backend. @@ -397,6 +510,10 @@ export const env = { SANDBOX_BACKEND: sandboxBackend, /** Permit trusted callers to route each execution to a paired worker ID. */ BRIDGE_DYNAMIC_WORKERS: process.env.CODEAPI_BRIDGE_DYNAMIC_WORKERS === 'true', + /** Opt-in independent native workspace concurrency; serial by default. */ + BRIDGE_MAX_WORKSPACE_LEASE_SLOTS: Number( + process.env.CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS ?? 1, + ), /** Outbound worker selected by the remote-bridge backend. */ BRIDGE_WORKER_ID: process.env.CODEAPI_BRIDGE_WORKER_ID ?? '', /** Static compatibility auth or short-lived proof-of-possession credentials. */ @@ -431,14 +548,20 @@ export const env = { ), // Lambda MicroVM backend. Connector lists are comma-separated ARNs. LAMBDA_MICROVM_IMAGE_ARN: process.env.LAMBDA_MICROVM_IMAGE_ARN ?? '', - LAMBDA_MICROVM_IMAGE_VERSION: process.env.LAMBDA_MICROVM_IMAGE_VERSION || undefined, - LAMBDA_MICROVM_EXECUTION_ROLE_ARN: process.env.LAMBDA_MICROVM_EXECUTION_ROLE_ARN || undefined, + LAMBDA_MICROVM_IMAGE_VERSION: + process.env.LAMBDA_MICROVM_IMAGE_VERSION || undefined, + LAMBDA_MICROVM_EXECUTION_ROLE_ARN: + process.env.LAMBDA_MICROVM_EXECUTION_ROLE_ARN || undefined, /* Runtime VM stdout reaches CloudWatch only when RunMicrovm sends a logging * config AND an executionRoleArn is set — pairs with the role above. */ LAMBDA_MICROVM_LOG_GROUP: process.env.LAMBDA_MICROVM_LOG_GROUP || undefined, LAMBDA_MICROVM_REGION: process.env.LAMBDA_MICROVM_REGION || undefined, - LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS: parseArnList(process.env.LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS), - LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS: parseArnList(process.env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS), + LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS: parseArnList( + process.env.LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS, + ), + LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS: parseArnList( + process.env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS, + ), ...lambdaMicrovmNumericConfig, /* CreateMicrovmAuthToken is minted per execute + per checkpoint; share a * fleet-wide budget so concurrent warm-session executes queue instead of @@ -452,14 +575,19 @@ export const env = { process.env.CODEAPI_CHECKPOINT_MAX_BYTES, 512 * 1024 * 1024, ), - CHECKPOINT_TIMEOUT_MS: configuredNumber(process.env.CODEAPI_CHECKPOINT_TIMEOUT_MS, 60_000), - CHECKPOINT_PREFIX: process.env.CODEAPI_CHECKPOINT_PREFIX ?? 'rtsx-checkpoints/', + CHECKPOINT_TIMEOUT_MS: configuredNumber( + process.env.CODEAPI_CHECKPOINT_TIMEOUT_MS, + 60_000, + ), + CHECKPOINT_PREFIX: + process.env.CODEAPI_CHECKPOINT_PREFIX ?? 'rtsx-checkpoints/', /** Dedicated Lambda MicroVM resident-server fleet. This remains an explicit * stateful-stack capability; the ordinary/default HTTP profile never starts * or preserves application processes. */ HOSTED_APPS_ENABLED: process.env.CODEAPI_HOSTED_APPS_ENABLED === 'true', HOSTED_APP_IMAGE_ARN: process.env.LAMBDA_MICROVM_APP_IMAGE_ARN ?? '', - HOSTED_APP_IMAGE_VERSION: process.env.LAMBDA_MICROVM_APP_IMAGE_VERSION || undefined, + HOSTED_APP_IMAGE_VERSION: + process.env.LAMBDA_MICROVM_APP_IMAGE_VERSION || undefined, /* These values are part of the pinned app-host image contract. RunMicrovm * cannot inject environment variables into the image, so exposing overrides * here would only make the control plane call ports the runner never opened. */ @@ -478,8 +606,10 @@ export const env = { 900, ), HOSTED_APP_START_TIMEOUT_MS: 30_000 as number, - HOSTED_APP_CREDENTIAL_KEY: process.env.CODEAPI_HOSTED_APP_CREDENTIAL_KEY ?? '', - HOSTED_APP_PREVIEW_ORIGIN: process.env.CODEAPI_HOSTED_APP_PREVIEW_ORIGIN ?? '', + HOSTED_APP_CREDENTIAL_KEY: + process.env.CODEAPI_HOSTED_APP_CREDENTIAL_KEY ?? '', + HOSTED_APP_PREVIEW_ORIGIN: + process.env.CODEAPI_HOSTED_APP_PREVIEW_ORIGIN ?? '', HOSTED_APP_PREVIEW_SIGNING_KEY: process.env.CODEAPI_HOSTED_APP_PREVIEW_SIGNING_KEY ?? '', }; @@ -502,7 +632,9 @@ type PlanLimits = { * JSON object keyed by the `plan_id` JWT claim. Unknown or absent plan ids * fall back to the default tier, which is the only entry defined in code. */ -export function parsePlanLimits(raw: string | undefined): Record { +export function parsePlanLimits( + raw: string | undefined, +): Record { if (raw == null || raw.trim() === '') { return {}; } @@ -510,10 +642,14 @@ export function parsePlanLimits(raw: string | undefined): Record; } @@ -521,7 +657,8 @@ export function parsePlanLimits(raw: string | undefined): Record Date: Tue, 8 Sep 2026 22:43:43 -0400 Subject: [PATCH 2/7] fix: Bound Workspace Cleanup Waits and Preserve Legacy Heartbeats --- packages/code/src/cli.ts | 149 ++-- packages/code/src/native-pool.test.ts | 112 +++ packages/code/src/native-pool.ts | 24 +- packages/code/src/protocol.ts | 36 +- packages/code/src/worker-slots.test.ts | 59 +- packages/code/src/worker.ts | 64 +- service/src/bridge/concurrent-store.test.ts | 39 ++ service/src/bridge/router.ts | 718 ++++++++++---------- service/src/bridge/slots.test.ts | 10 + service/src/bridge/slots.ts | 8 +- service/src/bridge/store.ts | 303 ++++----- service/src/config.ts | 306 +++------ 12 files changed, 946 insertions(+), 882 deletions(-) create mode 100644 packages/code/src/native-pool.test.ts diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index bf9355f2..a9f59556 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -158,7 +158,9 @@ function githubCredentials(): { 'Configure either GitHub App authentication or a GitHub token, not both', ); } - const configuredHostValue = nonEmpty(process.env.LIBRECHAT_CODE_GITHUB_HOST); + const configuredHostValue = nonEmpty( + process.env.LIBRECHAT_CODE_GITHUB_HOST, + ); const configuredHost = configuredHostValue ? normalizeGitHubHost(configuredHostValue) : undefined; @@ -480,12 +482,12 @@ async function run( const roots: LocalWorkspaceConfig[] = canonicalWorkerDirectory ? [ { - id: workspaceId, + id: workspaceId, root: canonicalWorkerDirectory, - writable: allowWorkspaceWrites, + writable: allowWorkspaceWrites, name: - option(args, '--workspace-name') ?? - process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? + option(args, '--workspace-name') ?? + process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? (useDefaultWorkspace ? workspaceId : defaultWorkspaceName(workerDirectory!, workspaceId)), @@ -509,8 +511,7 @@ async function run( const separator = value.indexOf('='); if ( separator < 1 || - separator === value.length - 1 || - !canonicalWorkerDirectory || + separator === value.length - 1 || !canonicalWorkerDirectory || commandSandboxMode !== 'native-srt' ) { throw new Error( @@ -520,7 +521,7 @@ async function run( roots.push({ id: value.slice(0, separator), root: await realpath(value.slice(separator + 1)), - writable: allowWorkspaceWrites, + writable: allowWorkspaceWrites, }); } // Aliases and nested grants are not independent execution domains. @@ -554,7 +555,7 @@ async function run( } if ( workspaceLeaseSlots > 1 && - (!allowWorkspaceCommands || commandSandboxMode !== 'native-srt') + ( !allowWorkspaceCommands || commandSandboxMode !== 'native-srt') ) { throw new Error('Concurrent workspace leases require native-srt commands'); } @@ -747,38 +748,38 @@ async function run( statefulWorkspace, }); const nativeOptions: NativeProcessSandboxOptions = { - workspaceRoot: canonicalWorkerDirectory!, + workspaceRoot: canonicalWorkerDirectory!, protectedPaths: [ identityPath, ...rootQuarantinePaths.values(), - github.privateKeyPath, + github.privateKeyPath, ].filter((path): path is string => path != null), - allowedDomains: commandAllowedDomains, - ...(github.provider - ? { - maskedEnvironment: { - variables: [ - { - name: GITHUB_CREDENTIAL_ENV_NAME, - extract: '^(.+)$', - injectHosts: [github.host], - }, - ], - async resolve(signal?: AbortSignal) { - return gitHubCredentialEnvironment( - await github.provider!.getCredential(signal), - ); - }, - wrapCommand(command: string, platform: NodeJS.Platform) { - return wrapGitHubCredentialCommand( - command, - github.host, - platform, - ); - }, - }, - } - : {}), + allowedDomains: commandAllowedDomains, + ...(github.provider + ? { + maskedEnvironment: { + variables: [ + { + name: GITHUB_CREDENTIAL_ENV_NAME, + extract: '^(.+)$', + injectHosts: [github.host], + }, + ], + async resolve(signal?: AbortSignal) { + return gitHubCredentialEnvironment( + await github.provider!.getCredential(signal), + ); + }, + wrapCommand(command: string, platform: NodeJS.Platform) { + return wrapGitHubCredentialCommand( + command, + github.host, + platform, + ); + }, + }, + } + : {}), }; const nativeCommandSandbox = allowWorkspaceCommands && commandSandboxMode === 'native-srt' @@ -786,7 +787,7 @@ async function run( ? new NativeWorkspaceCommandPool( new Map( roots.map((root) => [ - root.id, + root.id, { ...nativeOptions, workspaceRoot: root.root }, ]), ), @@ -859,11 +860,11 @@ async function run( ? { workspaceQuarantines: new Map( roots.map((root) => [ - root.id, + root.id, workspaceMutationGuard( rootQuarantinePaths.get(root.id)!, workerId, - root.id, + root.id, incarnationId, ), ]), @@ -875,41 +876,41 @@ async function run( workspaceLeaseSlots === 1 && roots.length === 1 ? { - async assertAvailable() { - const record = await loadWorkspaceMutationQuarantine( - mutationQuarantinePath, - ); - if (record != null) { - throw new BridgeProtocolError( - `Workspace mutations are quarantined since ${record.quarantinedAt}: ${record.reason}. Inspect or restore the workspace, then run librechat-code clear-workspace-quarantine`, - undefined, - 'WORKER_QUARANTINED', - ); - } - }, - async arm(reason) { - await saveWorkspaceMutationQuarantine(mutationQuarantinePath, { - version: 1, - workerId, - workspaceId, - ownerId: incarnationId, - quarantinedAt: new Date().toISOString(), - reason, - }); - }, - async clear() { - await clearWorkspaceMutationQuarantine( - mutationQuarantinePath, - incarnationId, - ); - }, - async quarantine() { - await assertWorkspaceMutationQuarantineOwner( - mutationQuarantinePath, - incarnationId, + async assertAvailable() { + const record = await loadWorkspaceMutationQuarantine( + mutationQuarantinePath, + ); + if (record != null) { + throw new BridgeProtocolError( + `Workspace mutations are quarantined since ${record.quarantinedAt}: ${record.reason}. Inspect or restore the workspace, then run librechat-code clear-workspace-quarantine`, + undefined, + 'WORKER_QUARANTINED', ); - }, - } + } + }, + async arm(reason) { + await saveWorkspaceMutationQuarantine(mutationQuarantinePath, { + version: 1, + workerId, + workspaceId, + ownerId: incarnationId, + quarantinedAt: new Date().toISOString(), + reason, + }); + }, + async clear() { + await clearWorkspaceMutationQuarantine( + mutationQuarantinePath, + incarnationId, + ); + }, + async quarantine() { + await assertWorkspaceMutationQuarantineOwner( + mutationQuarantinePath, + incarnationId, + ); + }, + } : undefined, onIdentityChange: pairedIdentity && identityPath diff --git a/packages/code/src/native-pool.test.ts b/packages/code/src/native-pool.test.ts new file mode 100644 index 00000000..c37a9b92 --- /dev/null +++ b/packages/code/src/native-pool.test.ts @@ -0,0 +1,112 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { NativeWorkspaceCommandPool } from './native-pool.js'; +import type { WorkspaceExecuteCommandRequest } from './protocol.js'; + +const roots = new Map( + ['a', 'b', 'c'].map((id) => [id, { workspaceRoot: `/fixture/${id}` }]), +); +const request = (workspaceId: string): WorkspaceExecuteCommandRequest => ({ + protocolVersion: 1, + operation: 'execute_command', + workspaceId, + command: 'fixture', +}); +test('native pool reuses roots and evicts only idle processes within its bound', async () => { + const created: string[] = []; + const closed: string[] = []; + const pool = new NativeWorkspaceCommandPool(roots, 2, (options) => { + created.push(options.workspaceRoot); + return { + async prepare() {}, + async close() { + closed.push(options.workspaceRoot); + }, + async execute(req) { + return { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: req.workspaceId, + stdout: '', + stderr: '', + exitCode: 0, + truncated: false, + timedOut: false, + }; + }, + }; + }); + await pool.execute(request('a')); + await pool.execute(request('b')); + await pool.execute(request('a')); + assert.equal(created.length, 2); + await pool.execute(request('c')); + assert.deepEqual(closed, ['/fixture/b']); + await pool.close(); + assert.equal(closed.length, 3); +}); + +test('native pool never evicts an executing root or misclassifies pre-dispatch exhaustion', async () => { + let finish!: () => void; + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + const pending = new Promise((resolve) => { + finish = resolve; + }); + const pool = new NativeWorkspaceCommandPool(roots, 1, () => ({ + async prepare() {}, + async close() {}, + async execute(req) { + entered(); + await pending; + return { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: req.workspaceId, + stdout: '', + stderr: '', + exitCode: 0, + truncated: false, + timedOut: false, + }; + }, + })); + const executing = pool.execute(request('a')); + await started; + await assert.rejects(pool.execute(request('b')), { + code: 'COMMAND_UNAVAILABLE', + mutationMayHaveCommitted: false, + }); + finish(); + await executing; + await pool.close(); +}); + +test('idle eviction failure stays mutation-atomic for the new root', async () => { + const pool = new NativeWorkspaceCommandPool(roots, 1, () => ({ + async prepare() {}, + async close() { + throw new Error('fixture cleanup failure'); + }, + async execute(req) { + return { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: req.workspaceId, + stdout: '', + stderr: '', + exitCode: 0, + truncated: false, + timedOut: false, + }; + }, + })); + await pool.execute(request('a')); + await assert.rejects(pool.execute(request('b')), { + code: 'COMMAND_UNAVAILABLE', + mutationMayHaveCommitted: false, + }); + await assert.rejects(pool.close(), AggregateError); +}); diff --git a/packages/code/src/native-pool.ts b/packages/code/src/native-pool.ts index 82507d65..0cdd1d2e 100644 --- a/packages/code/src/native-pool.ts +++ b/packages/code/src/native-pool.ts @@ -7,7 +7,10 @@ import type { } from './protocol.js'; interface Entry { - sandbox: NativeProcessWorkspaceCommandSandbox; + sandbox: Pick< + NativeProcessWorkspaceCommandSandbox, + 'prepare' | 'execute' | 'close' + >; busy: boolean; } @@ -22,6 +25,10 @@ export class NativeWorkspaceCommandPool { constructor( private readonly roots: ReadonlyMap, private readonly capacity: number, + private readonly createSandbox: ( + options: NativeProcessSandboxOptions, + ) => Entry['sandbox'] = (options) => + new NativeProcessWorkspaceCommandSandbox(options), ) { if ( !Number.isSafeInteger(capacity) || @@ -61,7 +68,7 @@ export class NativeWorkspaceCommandPool { this.entries.delete(idle[0]); } entry = { - sandbox: new NativeProcessWorkspaceCommandSandbox(options), + sandbox: this.createSandbox(options), busy: false, }; } @@ -71,8 +78,17 @@ export class NativeWorkspaceCommandPool { this.entries.set(root, entry); return entry; }); - this.allocation = pending.catch(() => undefined); - return pending; + const checked = pending.catch((error) => { + // Allocation/idle eviction precedes dispatch into the requested root. + // Do not turn a pool resource failure into an uncertain mutation there. + if (error instanceof WorkspaceToolError) throw error; + throw new WorkspaceToolError( + 'Native executor allocation failed', + 'COMMAND_UNAVAILABLE', + ); + }); + this.allocation = checked.catch(() => undefined); + return checked; } async prepare(): Promise { diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 07088f94..44cfdfd0 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -390,7 +390,12 @@ const WORKSPACE_COMMAND_RESULT_KEYS = new Set([ 'truncated', 'timedOut', ]); -const WORKSPACE_SEARCH_MATCH_KEYS = new Set(['path', 'line', 'column', 'text']); +const WORKSPACE_SEARCH_MATCH_KEYS = new Set([ + 'path', + 'line', + 'column', + 'text', +]); export interface BridgeWorkerCapabilities { /** Opt-in protocol: maximum concurrently leased independent workspace roots. */ @@ -551,8 +556,7 @@ export function isWorkspaceToolErrorCode( } export type BridgeSettlement = - | BridgeFulfilledSettlement - | BridgeRejectedSettlement; + BridgeFulfilledSettlement | BridgeRejectedSettlement; export interface BridgeSettlementResponse { protocolVersion: BridgeProtocolVersion; @@ -609,10 +613,7 @@ function normalizePortableRelativePath(value: string): string { } /** Compare path segments in ripgrep's sorted, depth-first traversal order. */ -export function comparePortableRelativePaths( - left: string, - right: string, -): number { +export function comparePortableRelativePaths(left: string, right: string): number { const encoder = new TextEncoder(); const leftSegments = left.split('/'); const rightSegments = right.split('/'); @@ -642,14 +643,9 @@ function isWithinRequestedPath(candidate: string, requested?: string): boolean { ); } -function isValidWorkspaceEditRequest( - request: Record, -): boolean { +function isValidWorkspaceEditRequest(request: Record): boolean { const hasBatch = request.edits !== undefined; - if ( - hasBatch && - (request.oldText !== undefined || request.newText !== undefined) - ) { + if (hasBatch && (request.oldText !== undefined || request.newText !== undefined)) { return false; } const edits = hasBatch @@ -755,8 +751,7 @@ export function isWorkspaceToolRequest( isSafePortableRelativePath(request.path)) && (request.afterPath === undefined || (isSafePortableRelativePath(request.afterPath) && - normalizePortableRelativePath(request.afterPath) === - request.afterPath && + normalizePortableRelativePath(request.afterPath) === request.afterPath && isWithinRequestedPath(request.afterPath, request.path))) && (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && @@ -843,16 +838,11 @@ export function isWorkspaceToolResult( const maxLines = request.maxLines ?? 200; const content = typeof result.content === 'string' ? result.content : null; const reportedLineCount = - Number.isSafeInteger(result.endLine) && - Number(result.endLine) >= startLine - 1 + Number.isSafeInteger(result.endLine) && Number(result.endLine) >= startLine - 1 ? Number(result.endLine) - startLine + 1 : -1; const actualLineCount = - content === null - ? -1 - : content.length === 0 - ? reportedLineCount - : content.split('\n').length; + content === null ? -1 : content.length === 0 ? reportedLineCount : content.split('\n').length; return ( hasOnlyKeys(result, WORKSPACE_READ_RESULT_KEYS) && result.path === request.path && diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts index d31aebed..47498789 100644 --- a/packages/code/src/worker-slots.test.ts +++ b/packages/code/src/worker-slots.test.ts @@ -1,7 +1,10 @@ import test from 'node:test'; import assert from 'node:assert/strict'; import { BridgeWorker } from './worker.js'; -import type { BridgeWorkspaceToolCapabilities } from './protocol.js'; +import type { + BridgeAssignment, + BridgeWorkspaceToolCapabilities, +} from './protocol.js'; const capabilities: BridgeWorkspaceToolCapabilities = { protocolVersion: 1, @@ -71,3 +74,57 @@ for (const receipt of [undefined, 1]) { await assert.rejects(worker.lease(undefined, 0), /negotiated capacity/); }); } + +for (const cancelled of [false, true]) { + test(`local cleanup wait rejects unexecuted work on ${cancelled ? 'cancellation' : 'expiry'}`, async () => { + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'fixture', + runtimes: [], + }, + }); + // Exercise the handoff seam without involving the unrelated HTTP settlement retry loop. + const internals = worker as unknown as { + activeWorkspaceAssignments: Map< + string, + { id: string; done: Promise } + >; + rejectUnexecutedAssignment: () => Promise; + executeOwned: () => Promise; + }; + internals.activeWorkspaceAssignments.set('a', { + id: 'previous', + done: new Promise(() => {}), + }); + let rejected = false; + internals.rejectUnexecutedAssignment = async () => { + rejected = true; + }; + internals.executeOwned = async () => { + assert.fail('must not enter a root still cleaning up'); + }; + const controller = new AbortController(); + if (cancelled) controller.abort(); + await worker.executeAndSettle( + { + assignmentId: 'next', + executionKind: 'workspace_tool', + remainingMs: cancelled ? 60000 : 5, + request: { + protocolVersion: 1, + workspaceId: 'a', + operation: 'read_file', + path: 'test.txt', + }, + } as BridgeAssignment, + controller.signal, + ); + assert.equal(rejected, true); + assert.equal(internals.activeWorkspaceAssignments.get('a')?.id, 'previous'); + }); +} diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 037ecd68..068422fc 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -147,23 +147,21 @@ function workspaceCapabilitiesMatch( advertised.writeFileModes?.length === executor.writeFileModes?.length && (advertised.writeFileModes?.every( (mode, index) => mode === executor.writeFileModes?.[index], - ) ?? - executor.writeFileModes == null) && + ) ?? executor.writeFileModes == null) && advertised.editFileModes?.length === executor.editFileModes?.length && (advertised.editFileModes?.every( (mode, index) => mode === executor.editFileModes?.[index], - ) ?? - executor.editFileModes == null) && - advertised.editFileFeatures?.length === executor.editFileFeatures?.length && + ) ?? executor.editFileModes == null) && + advertised.editFileFeatures?.length === + executor.editFileFeatures?.length && (advertised.editFileFeatures?.every( (feature, index) => feature === executor.editFileFeatures?.[index], - ) ?? - executor.editFileFeatures == null) && - advertised.listFileFeatures?.length === executor.listFileFeatures?.length && + ) ?? executor.editFileFeatures == null) && + advertised.listFileFeatures?.length === + executor.listFileFeatures?.length && (advertised.listFileFeatures?.every( (feature, index) => feature === executor.listFileFeatures?.[index], - ) ?? - executor.listFileFeatures == null) && + ) ?? executor.listFileFeatures == null) && advertised.workspaces.length === executor.workspaces.length && advertised.workspaces.every( (workspace, index) => @@ -175,8 +173,7 @@ function workspaceCapabilitiesMatch( (operation, operationIndex) => operation === executor.workspaces[index]?.operations?.[operationIndex], - ) ?? - executor.workspaces[index]?.operations == null), + ) ?? executor.workspaces[index]?.operations == null), ) ); } @@ -188,7 +185,8 @@ function registrationCompatibleCapabilities( if ( workspaceTools == null || (workspaceTools.operations.every( - (operation) => operation === 'read_file' || operation === 'search_text', + (operation) => + operation === 'read_file' || operation === 'search_text', ) && workspaceTools.workspaces.every( (workspace) => workspace.operations == null, @@ -197,7 +195,8 @@ function registrationCompatibleCapabilities( return capabilities; } const operations = workspaceTools.operations.filter( - (operation) => operation === 'read_file' || operation === 'search_text', + (operation) => + operation === 'read_file' || operation === 'search_text', ); if (operations.length === 0) { const { workspaceTools: _workspaceTools, ...compatible } = capabilities; @@ -206,9 +205,7 @@ function registrationCompatibleCapabilities( const workspaces = workspaceTools.workspaces.flatMap((workspace) => { if ( workspace.operations != null && - !operations.every((operation) => - workspace.operations?.includes(operation), - ) + !operations.every((operation) => workspace.operations?.includes(operation)) ) { return []; } @@ -1090,9 +1087,33 @@ export class BridgeWorker { ); // A settlement can commit remotely before the local durable guard clears. // Keep the next lane out of the root until that cleanup has finished. - await active.done; - if (signal?.aborted) - throw signal.reason ?? new DOMException('aborted', 'AbortError'); + const waitController = new AbortController(); + const onAbort = (): void => waitController.abort(); + signal?.addEventListener('abort', onAbort, { once: true }); + if (signal?.aborted) waitController.abort(); + let finished: boolean; + try { + finished = await Promise.race([ + active.done.then(() => true), + abortableDelay( + Math.max( + 0, + this.assignmentRemainingMs(assignment) - (Date.now() - waitingAt), + ), + waitController.signal, + ).then(() => false), + ]); + } finally { + waitController.abort(); + signal?.removeEventListener('abort', onAbort); + } + if (!finished || signal?.aborted) { + await this.rejectUnexecutedAssignment( + assignment, + 'Workspace cleanup wait ended before execution', + ); + return; + } } let release!: () => void; if (root != null) @@ -1251,7 +1272,8 @@ export class BridgeWorker { } const workspaceRequest = assignment.request; try { - await guard?.assertAvailable(); + if (this.options.workspaceQuarantines != null) + await guard?.assertAvailable(); } catch (error) { throw new BridgeWorkspaceQuarantinedError( 'Workspace is quarantined', diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts index b94aaece..84f41d0f 100644 --- a/service/src/bridge/concurrent-store.test.ts +++ b/service/src/bridge/concurrent-store.test.ts @@ -133,3 +133,42 @@ test('same-root work waits while another root progresses', async () => { await settle(third!); await Promise.all([nextA, b]); }); + +test('queued cancellation never leases and does not block another root', async () => { + await register(); + const a = dispatch('a'); + const first = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 0))!; + const controller = new AbortController(); + const cancelled = dispatch('a', controller.signal); + const queue = `codeapi:bridge:v1:worker:${workerId}:admission`; + for (let i = 0; i < 100 && await redis.zcard(queue) < 2; i++) { + await new Promise(resolve => setTimeout(resolve, 5)); + } + expect(await redis.zcard(queue)).toBe(2); + controller.abort(); + await expect(cancelled).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + const b = dispatch('b'); + const second = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 1))!; + expect(second.request).toMatchObject({ workspaceId: 'b' }); + await settle(first); await settle(second); await Promise.all([a, b]); + expect(await store.lease(workerId, incarnationId, 0, undefined, undefined, 0)).toBeUndefined(); +}); + +test('late quarantine releases its slot after caller cancellation and retains only its root fence', async () => { + await register(); + const controller = new AbortController(); + const a = dispatch('a', controller.signal); + const first = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 0))!; + await store.acknowledgeLease(workerId, incarnationId, first.assignmentId, first.generation, first.leaseToken); + controller.abort(); + await expect(a).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + await store.settle(workerId, first.assignmentId, { protocolVersion: 1, + incarnationId, generation: first.generation, leaseToken: first.leaseToken, + status: 'rejected', error: 'uncertain mutation' }, undefined, undefined, true); + expect(await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`)).toBeNull(); + const b = dispatch('b'); + const next = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 0))!; + expect(next.request).toMatchObject({ workspaceId: 'b' }); + await settle(next); await b; + await expect(dispatch('a')).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); +}); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index b362e1b3..4410eb89 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -94,9 +94,9 @@ function sendStoreError(error: BridgeStoreError, res: Response): void { ? 404 : error.code === 'WORKER_UNAUTHORIZED' ? 403 - : error.code === 'WORKER_BUSY' - ? 503 - : 409; + : error.code === 'WORKER_BUSY' + ? 503 + : 409; res.status(status).json({ error: error.message, code: error.code }); } @@ -143,16 +143,18 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ?.match(/^Bearer\s+(.+)$/i)?.[1] ?.trim() ?? ''; - const adminAuth = (req: Request, res: Response, next: NextFunction): void => { + const adminAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { if (!options.adminToken) { res.status(503).json({ error: 'Code bridge is not configured' }); return; } const token = bearerToken(req); if (!token || !sameToken(token, options.adminToken)) { - res - .status(401) - .json({ error: 'Invalid code bridge administrator token' }); + res.status(401).json({ error: 'Invalid code bridge administrator token' }); return; } next(); @@ -226,89 +228,73 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { const workerAuth = options.authMode === 'paired' ? pairedWorkerAuth : staticWorkerAuth; - router.post( - '/pairings', - adminAuth, - asyncRoute(async (req, res) => { - if (options.authMode !== 'paired') { - res - .status(409) - .json({ error: 'Paired worker authentication is disabled' }); - return; - } - const workerId = isRecord(req.body) ? req.body.workerId : undefined; - if ( - typeof workerId !== 'string' || - !validWorkerId(workerId) || - !configuredWorker(workerId) - ) { - res.status(400).json({ error: 'Invalid bridge worker ID' }); - return; - } - const hasBinding = - isRecord(req.body) && - Object.prototype.hasOwnProperty.call(req.body, 'binding'); - const binding = isRecord(req.body) - ? parseBinding(req.body.binding) - : undefined; - if (hasBinding && binding == null) { - res - .status(400) - .json({ error: 'Invalid bridge worker principal binding' }); - return; - } - if (options.allowDynamicWorkers === true && binding == null) { - res.status(400).json({ - error: 'Dynamic bridge workers require a valid principal binding', - }); - return; - } - const pairing = await options.pairings.issue(workerId, binding); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); - }), - ); + router.post('/pairings', adminAuth, asyncRoute(async (req, res) => { + if (options.authMode !== 'paired') { + res.status(409).json({ error: 'Paired worker authentication is disabled' }); + return; + } + const workerId = isRecord(req.body) ? req.body.workerId : undefined; + if ( + typeof workerId !== 'string' || + !validWorkerId(workerId) || + !configuredWorker(workerId) + ) { + res.status(400).json({ error: 'Invalid bridge worker ID' }); + return; + } + const hasBinding = isRecord(req.body) && + Object.prototype.hasOwnProperty.call(req.body, 'binding'); + const binding = isRecord(req.body) ? parseBinding(req.body.binding) : undefined; + if (hasBinding && binding == null) { + res.status(400).json({ error: 'Invalid bridge worker principal binding' }); + return; + } + if (options.allowDynamicWorkers === true && binding == null) { + res.status(400).json({ + error: 'Dynamic bridge workers require a valid principal binding', + }); + return; + } + const pairing = await options.pairings.issue(workerId, binding); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); + })); - router.post( - '/pairings/redeem', - asyncRoute(async (req, res) => { - if (options.authMode !== 'paired') { - res - .status(409) - .json({ error: 'Paired worker authentication is disabled' }); - return; - } - const redemption = req.body as unknown; - if ( - !isRecord(redemption) || - redemption.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - typeof redemption.workerId !== 'string' || - !validWorkerId(redemption.workerId) || - !configuredWorker(redemption.workerId) || - typeof redemption.code !== 'string' || - redemption.code.length < 16 || - typeof redemption.publicKey !== 'string' || - redemption.publicKey.length > 4096 - ) { - res.status(400).json({ error: 'Invalid bridge pairing redemption' }); + router.post('/pairings/redeem', asyncRoute(async (req, res) => { + if (options.authMode !== 'paired') { + res.status(409).json({ error: 'Paired worker authentication is disabled' }); + return; + } + const redemption = req.body as unknown; + if ( + !isRecord(redemption) || + redemption.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof redemption.workerId !== 'string' || + !validWorkerId(redemption.workerId) || + !configuredWorker(redemption.workerId) || + typeof redemption.code !== 'string' || + redemption.code.length < 16 || + typeof redemption.publicKey !== 'string' || + redemption.publicKey.length > 4096 + ) { + res.status(400).json({ error: 'Invalid bridge pairing redemption' }); + return; + } + try { + const credential = await options.pairings.redeem({ + workerId: redemption.workerId, + code: redemption.code, + publicKey: redemption.publicKey, + }); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); + } catch (error) { + if (error instanceof BridgePairingError) { + const status = error.code === 'PUBLIC_KEY_INVALID' ? 400 : 401; + res.status(status).json({ error: error.message, code: error.code }); return; } - try { - const credential = await options.pairings.redeem({ - workerId: redemption.workerId, - code: redemption.code, - publicKey: redemption.publicKey, - }); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); - } catch (error) { - if (error instanceof BridgePairingError) { - const status = error.code === 'PUBLIC_KEY_INVALID' ? 400 : 401; - res.status(status).json({ error: error.message, code: error.code }); - return; - } - throw error; - } - }), - ); + throw error; + } + })); router.post( '/workers/:workerId/revoke', @@ -345,84 +331,87 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ); router.post( - '/workers/register', + '/workers/register', workerAuth, asyncRoute(async (req, res) => { - const registration = req.body as unknown; - if ( - !isRecord(registration) || - registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - typeof registration.workerId !== 'string' || - !validWorkerId(registration.workerId) || - !validIncarnationId(registration.incarnationId) || - !isValidBridgeWorkerCapabilities(registration.capabilities) - ) { - res.status(400).json({ error: 'Invalid bridge worker registration' }); - return; - } - if (!configuredWorker(registration.workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - const authorization = - options.authMode === 'paired' - ? (res.locals.bridgeWorkerAuthorization as { - identityId: string; - pairingGeneration: number; - credentialId: string; - activeCredentialId: string; - binding?: BridgeWorkerBinding; - }) - : undefined; - const trustedRegistration: BridgeWorkerRegistration & { - binding?: BridgeWorkerBinding; - credentialId?: string; - identityId?: string; - } = { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - workerId: registration.workerId, - incarnationId: registration.incarnationId, - capabilities: registration.capabilities, - ...(authorization?.credentialId != null - ? { credentialId: authorization.credentialId } - : {}), - ...(authorization?.identityId != null - ? { identityId: authorization.identityId } - : {}), - ...(authorization?.binding != null - ? { binding: authorization.binding } - : {}), - }; + const registration = req.body as unknown; + if ( + !isRecord(registration) || + registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof registration.workerId !== 'string' || + !validWorkerId(registration.workerId) || + !validIncarnationId(registration.incarnationId) || + !isValidBridgeWorkerCapabilities(registration.capabilities) + ) { + res.status(400).json({ error: 'Invalid bridge worker registration' }); + return; + } + if ( + !configuredWorker(registration.workerId) + ) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + const authorization = options.authMode === 'paired' + ? ( + res.locals.bridgeWorkerAuthorization as { + identityId: string; + pairingGeneration: number; + credentialId: string; + activeCredentialId: string; + binding?: BridgeWorkerBinding; + } + ) + : undefined; + const trustedRegistration: BridgeWorkerRegistration & { + binding?: BridgeWorkerBinding; + credentialId?: string; + identityId?: string; + } = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + capabilities: registration.capabilities, + ...(authorization?.credentialId != null + ? { credentialId: authorization.credentialId } + : {}), + ...(authorization?.identityId != null + ? { identityId: authorization.identityId } + : {}), + ...(authorization?.binding != null + ? { binding: authorization.binding } + : {}), + }; try { - const registrationGeneration = await options.store.register( - trustedRegistration, - authorization, - ); + const registrationGeneration = await options.store.register( + trustedRegistration, + authorization, + ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: registration.workerId, incarnationId: registration.incarnationId, - registrationGeneration, - registeredAt: new Date().toISOString(), - leaseTtlMs: 60_000, + registrationGeneration, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, workspaceLeaseSlots: options.store.workspaceLeaseCapacity( registration.capabilities.workspaceLeaseSlots, ), - supportedWorkspaceToolOperations: [ - 'read_file', - 'search_text', - 'list_files', - 'write_file', - 'preview_edit', - 'edit_file', - 'execute_command', - ], - supportedWorkspaceWriteFileModes: ['replace', 'create'], - supportedWorkspaceEditFileModes: ['single', 'batch'], - supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], - supportedWorkspaceListFileFeatures: ['after_path'], + supportedWorkspaceToolOperations: [ + 'read_file', + 'search_text', + 'list_files', + 'write_file', + 'preview_edit', + 'edit_file', + 'execute_command', + ], + supportedWorkspaceWriteFileModes: ['replace', 'create'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], + supportedWorkspaceListFileFeatures: ['after_path'], }); } catch (error) { if (error instanceof BridgeStoreError) { @@ -434,124 +423,124 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }), ); - router.post( - '/workers/:workerId/ready', - workerAuth, - asyncRoute(async (req, res) => { - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; - if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isSafeInteger(body.registrationGeneration) || - Number(body.registrationGeneration) < 1 - ) { - res.status(400).json({ - error: 'Invalid bridge worker readiness confirmation', - }); - return; - } - if (!configuredWorker(workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); +router.post( + '/workers/:workerId/ready', + workerAuth, + asyncRoute(async (req, res) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isSafeInteger(body.registrationGeneration) || + Number(body.registrationGeneration) < 1 + ) { + res.status(400).json({ + error: 'Invalid bridge worker readiness confirmation', + }); + return; + } + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { + await options.store.confirmReady( + workerId, + body.incarnationId, + Number(body.registrationGeneration), + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ready: true }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); return; } + throw error; + } + }), +); + +router.post( + '/workers/:workerId/workspaces/reset', + workerAuth, + asyncRoute(async (req, res) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + if ( + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + typeof body.runtimeSessionId !== 'string' || + body.runtimeSessionId.trim().length === 0 || + body.runtimeSessionId.length > 512 || + body.confirmDiscarded !== true + ) { + res.status(400).json({ + error: 'Workspace reset requires confirmation of local discard', + }); + return; + } + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { + const resetController = new AbortController(); + const abortReset = (): void => resetController.abort(); + req.once('aborted', abortReset); + res.once('close', abortReset); try { - await options.store.confirmReady( + await options.store.resetWorkspace( workerId, body.incarnationId, - Number(body.registrationGeneration), + body.runtimeSessionId, + resetController.signal, ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ready: true }); - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); - return; + if (!resetController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); } - throw error; - } - }), - ); - - router.post( - '/workers/:workerId/workspaces/reset', - workerAuth, - asyncRoute(async (req, res) => { - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; - if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - typeof body.runtimeSessionId !== 'string' || - body.runtimeSessionId.trim().length === 0 || - body.runtimeSessionId.length > 512 || - body.confirmDiscarded !== true - ) { - res.status(400).json({ - error: 'Workspace reset requires confirmation of local discard', - }); - return; + } finally { + req.off('aborted', abortReset); + res.off('close', abortReset); } - if (!configuredWorker(workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); return; } - try { - const resetController = new AbortController(); - const abortReset = (): void => resetController.abort(); - req.once('aborted', abortReset); - res.once('close', abortReset); - try { - await options.store.resetWorkspace( - workerId, - body.incarnationId, - body.runtimeSessionId, - resetController.signal, - ); - if (!resetController.signal.aborted) { - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, reset: true }); - } - } finally { - req.off('aborted', abortReset); - res.off('close', abortReset); - } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); - return; - } - throw error; - } - }), - ); + throw error; + } + }), +); router.post( - '/workers/:workerId/lease', + '/workers/:workerId/lease', workerAuth, asyncRoute(async (req, res) => { - const requestStartedAtMs = Date.now(); + const requestStartedAtMs = Date.now(); const workerId = req.params.workerId; const body = isRecord(req.body) ? req.body : {}; - const requestedWait = Number(body.waitMs ?? 25_000); + const requestedWait = Number(body.waitMs ?? 25_000); if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isFinite(requestedWait) || - requestedWait < 0 || + !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isFinite(requestedWait) || + requestedWait < 0 || (body.workspaceLeaseSlot !== undefined && (!Number.isSafeInteger(body.workspaceLeaseSlot) || Number(body.workspaceLeaseSlot) < 0 || Number(body.workspaceLeaseSlot) >= 8)) ) { - res.status(400).json({ error: 'Invalid bridge lease request' }); - return; - } + res.status(400).json({ error: 'Invalid bridge lease request' }); + return; + } if (!configuredWorker(workerId)) { res.status(403).json({ error: 'Worker is not authorized for this Code API deployment', @@ -559,16 +548,16 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { return; } try { - const leaseController = new AbortController(); - const abortLease = (): void => leaseController.abort(); - req.once('aborted', abortLease); - res.once('close', abortLease); - let assignment: CodeBridgeAssignment | undefined; + const leaseController = new AbortController(); + const abortLease = (): void => leaseController.abort(); + req.once('aborted', abortLease); + res.once('close', abortLease); + let assignment: CodeBridgeAssignment | undefined; try { assignment = await options.store.lease( workerId, body.incarnationId, - Math.min(requestedWait, MAX_LEASE_WAIT_MS), + Math.min(requestedWait, MAX_LEASE_WAIT_MS), leaseController.signal, ( res.locals.bridgeWorkerAuthorization as @@ -579,19 +568,19 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ? undefined : Number(body.workspaceLeaseSlot), ); - if (leaseController.signal.aborted) { - if (assignment != null) await options.store.returnLease(assignment); - return; - } - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), - assignment, - }); - } finally { - req.off('aborted', abortLease); - res.off('close', abortLease); + if (leaseController.signal.aborted) { + if (assignment != null) await options.store.returnLease(assignment); + return; } + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), + assignment, + }); + } finally { + req.off('aborted', abortLease); + res.off('close', abortLease); + } } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); @@ -602,74 +591,71 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }), ); - router.post( - '/workers/:workerId/assignments/:assignmentId/ack', - workerAuth, - asyncRoute(async (req, res) => { - const body = isRecord(req.body) ? req.body : {}; - if ( - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isSafeInteger(body.generation) || - Number(body.generation) < 1 || - typeof body.leaseToken !== 'string' || - body.leaseToken.length < 32 - ) { - res.status(400).json({ error: 'Invalid bridge lease acknowledgement' }); - return; - } +router.post( + '/workers/:workerId/assignments/:assignmentId/ack', + workerAuth, + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || + !Number.isSafeInteger(body.generation) || + Number(body.generation) < 1 || + typeof body.leaseToken !== 'string' || + body.leaseToken.length < 32 + ) { + res.status(400).json({ error: 'Invalid bridge lease acknowledgement' }); + return; + } + try { + const acknowledgementController = new AbortController(); + const abortAcknowledgement = (): void => + acknowledgementController.abort(); + req.once('aborted', abortAcknowledgement); + res.once('close', abortAcknowledgement); try { - const acknowledgementController = new AbortController(); - const abortAcknowledgement = (): void => - acknowledgementController.abort(); - req.once('aborted', abortAcknowledgement); - res.once('close', abortAcknowledgement); - try { - await options.store.acknowledgeLease( - req.params.workerId, - body.incarnationId, - req.params.assignmentId, - Number(body.generation), - body.leaseToken, - acknowledgementController.signal, - ); - if (!acknowledgementController.signal.aborted) { - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - accepted: true, - }); - } - } finally { - req.off('aborted', abortAcknowledgement); - res.off('close', abortAcknowledgement); - } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); - return; + await options.store.acknowledgeLease( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + Number(body.generation), + body.leaseToken, + acknowledgementController.signal, + ); + if (!acknowledgementController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, accepted: true }); } - throw error; + } finally { + req.off('aborted', abortAcknowledgement); + res.off('close', abortAcknowledgement); } - }), - ); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }), +); router.post( [ - '/workers/:workerId/assignments/:assignmentId/settle', + '/workers/:workerId/assignments/:assignmentId/settle', '/workers/:workerId/assignments/:assignmentId/quarantine', ], workerAuth, asyncRoute(async (req, res) => { - const settlement = req.body as unknown; - if (!isSettlement(settlement)) { - res.status(400).json({ error: 'Invalid bridge settlement' }); - return; - } + const settlement = req.body as unknown; + if (!isSettlement(settlement)) { + res.status(400).json({ error: 'Invalid bridge settlement' }); + return; + } try { - const settlementController = new AbortController(); - const abortSettlement = (): void => settlementController.abort(); - req.once('aborted', abortSettlement); - res.once('close', abortSettlement); + const settlementController = new AbortController(); + const abortSettlement = (): void => settlementController.abort(); + req.once('aborted', abortSettlement); + res.once('close', abortSettlement); try { await options.store.settle( req.params.workerId, @@ -683,16 +669,16 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { )?.identityId, req.path.endsWith('/quarantine'), ); - if (!settlementController.signal.aborted) { - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - accepted: true, - }); - } - } finally { - req.off('aborted', abortSettlement); - res.off('close', abortSettlement); + if (!settlementController.signal.aborted) { + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + accepted: true, + }); } + } finally { + req.off('aborted', abortSettlement); + res.off('close', abortSettlement); + } } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); @@ -703,38 +689,38 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }), ); - router.post( - '/workers/:workerId/assignments/:assignmentId/cancellation', - workerAuth, - asyncRoute(async (req, res) => { - const body = isRecord(req.body) ? req.body : {}; - if ( - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) - ) { - res.status(400).json({ error: 'Invalid bridge cancellation request' }); - return; - } - const cancellationController = new AbortController(); - const abortCancellation = (): void => cancellationController.abort(); - req.once('aborted', abortCancellation); - res.once('close', abortCancellation); - try { - const cancelled = await options.store.cancelled( - req.params.workerId, - body.incarnationId, - req.params.assignmentId, - cancellationController.signal, - ); - if (!cancellationController.signal.aborted) { - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); - } - } finally { - req.off('aborted', abortCancellation); - res.off('close', abortCancellation); +router.post( + '/workers/:workerId/assignments/:assignmentId/cancellation', + workerAuth, + asyncRoute(async (req, res) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) + ) { + res.status(400).json({ error: 'Invalid bridge cancellation request' }); + return; + } + const cancellationController = new AbortController(); + const abortCancellation = (): void => cancellationController.abort(); + req.once('aborted', abortCancellation); + res.once('close', abortCancellation); + try { + const cancelled = await options.store.cancelled( + req.params.workerId, + body.incarnationId, + req.params.assignmentId, + cancellationController.signal, + ); + if (!cancellationController.signal.aborted) { + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); } - }), - ); + } finally { + req.off('aborted', abortCancellation); + res.off('close', abortCancellation); + } + }), +); router.post( '/workers/:workerId/credentials/refresh', diff --git a/service/src/bridge/slots.test.ts b/service/src/bridge/slots.test.ts index 8694810c..dfe89495 100644 --- a/service/src/bridge/slots.test.ts +++ b/service/src/bridge/slots.test.ts @@ -81,3 +81,13 @@ test('slots reject invalid capacity and replaced incarnation', async () => { await redis.set(`${prefix}:incarnation`, 'replacement'); await expect(slots.reserve(a)).rejects.toThrow('replaced'); }); + +test('releasing a long slot shortens the aggregate expiry to remaining work', async () => { + const a = await enqueue('a', 'root-a'); + const b = { ...await enqueue('b', 'root-b'), expiresAtMs: Date.now() + 3000 }; + await slots.reserve(a); + await slots.reserve(b); + expect(await redis.pttl(`${prefix}:lock`)).toBeGreaterThan(8000); + await slots.release(workerId, incarnationId, 'a'); + expect(await redis.pttl(`${prefix}:lock`)).toBeLessThanOrEqual(3000); +}); diff --git a/service/src/bridge/slots.ts b/service/src/bridge/slots.ts index 2a2b128d..9c9e7d3b 100644 --- a/service/src/bridge/slots.ts +++ b/service/src/bridge/slots.ts @@ -113,14 +113,20 @@ export class BridgeWorkspaceSlots { ): Promise { await this.redis.eval( [ + 'local latest = 0', `for slot = 0, ${MAX_WORKSPACE_LEASE_SLOTS - 1} do`, - " local entry = redis.call('HMGET', KEYS[1], 'a:' .. slot, 'i:' .. slot)", + " local entry = redis.call('HMGET', KEYS[1], 'a:' .. slot, 'i:' .. slot, 'e:' .. slot)", ' if entry[2] == ARGV[1] and entry[1] == ARGV[2] then', " redis.call('HDEL', KEYS[1], 'a:' .. slot, 'i:' .. slot, 'w:' .. slot, 'e:' .. slot)", + ' elseif type(entry[1]) == "string" then latest = math.max(latest, tonumber(entry[3]))', ' end', 'end', "if redis.call('HLEN', KEYS[1]) == 0 and redis.call('GET', KEYS[2]) == 'workspace-slots:' .. ARGV[1] then", " redis.call('DEL', KEYS[1], KEYS[2], KEYS[3])", + "elseif latest > 0 and redis.call('GET', KEYS[2]) == 'workspace-slots:' .. ARGV[1] then", + ' for _, key in ipairs(KEYS) do', + " redis.call('PEXPIREAT', key, latest)", + ' end', 'end', 'return 1', ].join('\n'), diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 22747de6..f7404934 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -34,8 +34,7 @@ export type CodeBridgeSettlement = BridgeSettlement< run?: t.ExecuteResponse['run']; } >; -export type CodeBridgeWorkspaceSettlement = - BridgeSettlement; +export type CodeBridgeWorkspaceSettlement = BridgeSettlement; type AnyCodeBridgeSettlement = | CodeBridgeSettlement | CodeBridgeWorkspaceSettlement; @@ -119,8 +118,7 @@ function supportsWorkspaceTool( ) { const mode = request.edits === undefined ? 'single' : 'batch'; const modes = capabilities?.editFileModes; - const supportsMode = - modes == null ? mode === 'single' : modes.includes(mode); + const supportsMode = modes == null ? mode === 'single' : modes.includes(mode); if (request.operation === 'preview_edit') return supportsMode; return ( supportsMode && @@ -362,7 +360,7 @@ export class RedisBridgeStore { this.redis.eval( [ "local registration = redis.call('GET', KEYS[1])", - 'if not registration then return { false, false, false, -2 } end', + "if not registration then return { false, false, false, -2 } end", 'return {', ' registration,', " redis.call('GET', KEYS[2]) or false,", @@ -378,17 +376,8 @@ export class RedisBridgeStore { this.redisCommandTimeoutMs, 'Bridge worker status', )) as [string | null, string | null, string | null, number]; - const [ - rawRegistration, - readyToken, - registrationGeneration, - leaseExpiresInMs, - ] = snapshot; - if ( - rawRegistration == null || - rawRegistration === '' || - leaseExpiresInMs <= 0 - ) { + const [rawRegistration, readyToken, registrationGeneration, leaseExpiresInMs] = snapshot; + if (rawRegistration == null || rawRegistration === '' || leaseExpiresInMs <= 0) { return { online: false, ready: false }; } @@ -408,16 +397,11 @@ export class RedisBridgeStore { return { online: false, ready: false }; } - const requiresConfirmation = - registration.capabilities.requiresReadyConfirmation === true; + const requiresConfirmation = registration.capabilities.requiresReadyConfirmation === true; const ready = !requiresConfirmation || (registrationGeneration != null && - readyToken === - workerReadyToken( - registration.incarnationId, - Number(registrationGeneration), - )); + readyToken === workerReadyToken(registration.incarnationId, Number(registrationGeneration))); return { online: true, ready, @@ -428,18 +412,15 @@ export class RedisBridgeStore { async register( registration: RegisteredBridgeWorker, - authorization?: - | string - | { - identityId?: string; - pairingGeneration?: number; - activeCredentialId?: string; - }, + authorization?: string | { + identityId?: string; + pairingGeneration?: number; + activeCredentialId?: string; + }, ): Promise { if (registration.capabilities.workspaceLeaseSlots !== undefined) { if ( - !isValidBridgeWorkerCapabilities(registration.capabilities) || - registration.capabilities.requiresReadyConfirmation !== true + !isValidBridgeWorkerCapabilities(registration.capabilities) ||registration.capabilities.requiresReadyConfirmation !== true ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -467,12 +448,12 @@ export class RedisBridgeStore { ' local pairingGeneration = redis.call(\'GET\', KEYS[7]) or "0"', ' if pairingGeneration ~= ARGV[5] then return -5 end', ' if ARGV[6] ~= "" then', - " if redis.call('GET', KEYS[8]) ~= ARGV[6] then return -5 end", + ' if redis.call(\'GET\', KEYS[8]) ~= ARGV[6] then return -5 end', ' elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -5', ' end', 'end', 'if ARGV[8] ~= "" then', - " local stableIdentity = redis.call('GET', KEYS[8])", + ' local stableIdentity = redis.call(\'GET\', KEYS[8])', ' if stableIdentity and stableIdentity ~= ARGV[8] then return -4 end', ' if not stableIdentity then', ' if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4 end', @@ -480,26 +461,26 @@ export class RedisBridgeStore { ' end', 'elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4', 'end', - "if redis.call('EXISTS', KEYS[3]) == 1 then return -2 end", - "if redis.call('EXISTS', KEYS[2]) == 1 then return -1 end", - "local current = redis.call('GET', KEYS[4])", + 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', + 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', + 'local current = redis.call(\'GET\', KEYS[4])', "if current == ARGV[1] and redis.call('EXISTS', KEYS[5]) == 1 and (redis.call('GET', KEYS[13]) or \"1\") ~= ARGV[10] then return -3 end", - "if not current and redis.call('EXISTS', KEYS[5]) == 1 then", - " local owner = redis.call('GET', KEYS[6])", + 'if not current and redis.call(\'EXISTS\', KEYS[5]) == 1 then', + ' local owner = redis.call(\'GET\', KEYS[6])', ' if owner ~= ARGV[1] then return -3 end', 'end', 'if current then', ' if current ~= ARGV[1] then', - " if redis.call('EXISTS', KEYS[5]) == 1 then return -3 end", - " redis.call('SET', ARGV[4] .. current .. ':fenced', \"1\")", + ' if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -3 end', + ' redis.call(\'SET\', ARGV[4] .. current .. \':fenced\', \"1\")', ' end', 'end', 'local registrationGeneration = tonumber(redis.call(\'GET\', KEYS[10]) or \"0\")', - "local registrationGenerationIncarnation = redis.call('GET', KEYS[11])", + 'local registrationGenerationIncarnation = redis.call(\'GET\', KEYS[11])', 'local registrationGenerationChanged = false', 'if registrationGeneration < 1 or registrationGenerationIncarnation ~= ARGV[1] then', - " registrationGeneration = redis.call('INCR', KEYS[10])", - " redis.call('SET', KEYS[11], ARGV[1])", + ' registrationGeneration = redis.call(\'INCR\', KEYS[10])', + ' redis.call(\'SET\', KEYS[11], ARGV[1])', ' registrationGenerationChanged = true', 'end', 'redis.call(\'SET\', KEYS[1], ARGV[2], \"EX\", ARGV[3])', @@ -514,10 +495,7 @@ export class RedisBridgeStore { script, 13, workerKey(registration.workerId), - incarnationFenceKey( - registration.workerId, - registration.incarnationId, - ), + incarnationFenceKey(registration.workerId, registration.incarnationId), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), lockKey(registration.workerId), @@ -539,9 +517,7 @@ export class RedisBridgeStore { authorizationObject?.identityId ?? '', expectedActiveCredentialId ?? '', registration.identityId ?? '', - registration.capabilities.requiresReadyConfirmation === true - ? '1' - : '0', + registration.capabilities.requiresReadyConfirmation === true ? '1' : '0', String(registration.capabilities.workspaceLeaseSlots ?? 1), ), this.redisCommandTimeoutMs, @@ -579,9 +555,7 @@ export class RedisBridgeStore { ); } if (!Number.isSafeInteger(result) || result < 1) { - throw new Error( - 'Bridge worker registration returned an invalid generation', - ); + throw new Error('Bridge worker registration returned an invalid generation'); } return result; } @@ -595,12 +569,12 @@ export class RedisBridgeStore { await boundedCommand( this.redis.eval( [ - "if redis.call('EXISTS', KEYS[1]) == 0 then return -1 end", - "if redis.call('GET', KEYS[2]) ~= ARGV[1] then return -2 end", - "if redis.call('GET', KEYS[3]) ~= ARGV[2] then return -2 end", - "if redis.call('GET', KEYS[4]) ~= ARGV[1] then return -2 end", - "if redis.call('EXISTS', KEYS[5]) == 1 then return -2 end", - "if redis.call('EXISTS', KEYS[6]) == 1 then return -3 end", + 'if redis.call(\'EXISTS\', KEYS[1]) == 0 then return -1 end', + 'if redis.call(\'GET\', KEYS[2]) ~= ARGV[1] then return -2 end', + 'if redis.call(\'GET\', KEYS[3]) ~= ARGV[2] then return -2 end', + 'if redis.call(\'GET\', KEYS[4]) ~= ARGV[1] then return -2 end', + 'if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -2 end', + 'if redis.call(\'EXISTS\', KEYS[6]) == 1 then return -3 end', 'redis.call(\'SET\', KEYS[7], ARGV[3], "EX", ARGV[4])', 'return 1', ].join('\n'), @@ -704,17 +678,12 @@ export class RedisBridgeStore { registration: RegisteredBridgeWorker, ) => Promise; }): Promise { - if ( - args.executionTimeoutMs !== undefined && - (args.workspaceRequest == null || - !Number.isSafeInteger(args.executionTimeoutMs) || - args.executionTimeoutMs < 1 || - args.executionTimeoutMs > 305_000) - ) { - throw new BridgeStoreError( - 'ASSIGNMENT_INVALID', - 'Invalid workspace execution budget', - ); + if (args.executionTimeoutMs !== undefined && ( + args.workspaceRequest == null || + !Number.isSafeInteger(args.executionTimeoutMs) || + args.executionTimeoutMs < 1 || args.executionTimeoutMs > 305_000 + )) { + throw new BridgeStoreError('ASSIGNMENT_INVALID', 'Invalid workspace execution budget'); } this.assertDispatchActive(args.signal, args.deadlineAtMs); const dispatchable = await this.dispatchCommand( @@ -779,9 +748,7 @@ export class RedisBridgeStore { const assignmentId = randomBytes(18).toString('base64url'); const leaseToken = randomBytes(32).toString('base64url'); // The lock is acquired before admission finishes; it must outlive the later execution deadline. - const ttlSeconds = assignmentTtlSeconds( - args.deadlineAtMs + (args.executionTimeoutMs ?? 0), - ); + const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs + (args.executionTimeoutMs ?? 0)); const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; let workspaceLeaseSlot: number | undefined; @@ -791,16 +758,14 @@ export class RedisBridgeStore { ? new BridgeWorkspaceSlots(this.redis) : undefined; let resultCommitted = false; - const admission = - args.workspaceRequest == null - ? undefined - : new BridgeAdmissionQueue(this.redis); + const admission = args.workspaceRequest == null + ? undefined + : new BridgeAdmissionQueue(this.redis); try { if ( admission != null && !(await this.dispatchCommand( - () => - admission.enter( + () => admission.enter( args.workerId, assignmentId, args.deadlineAtMs, @@ -809,39 +774,31 @@ export class RedisBridgeStore { : args.workspaceRequest?.workspaceId, ), args, - 'Bridge admission enqueue', + 'Bridge admission enqueue', )) ) { - throw new BridgeStoreError( - 'WORKER_QUEUE_FULL', - 'Bridge worker pending request limit reached', - ); + throw new BridgeStoreError('WORKER_QUEUE_FULL', 'Bridge worker pending request limit reached'); } let locked = false; do { if ( admission != null && - workspaceSlots == null && - !(await this.dispatchCommand( - () => admission.isHead(args.workerId, assignmentId), - args, - 'Bridge admission position', - )) + workspaceSlots == null && !(await this.dispatchCommand( + () => admission.isHead(args.workerId, assignmentId), + args, + 'Bridge admission position', + )) ) { - await delay( - Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), - args.signal, - ); + await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); continue; } if (workspaceSlots != null) { workspaceLeaseSlot = await this.dispatchCommand( () => workspaceSlots.reserve({ - workerId: args.workerId, + workerId: args.workerId, incarnationId: lockIncarnationId, - assignmentId, - workspaceId: args.workspaceRequest!.workspaceId, + assignmentId,workspaceId: args.workspaceRequest!.workspaceId, capacity: registration.capabilities.workspaceLeaseSlots!, expiresAtMs: Date.now() + ttlSeconds * 1000, }), @@ -850,23 +807,20 @@ export class RedisBridgeStore { ); locked = workspaceLeaseSlot !== undefined; } else { - locked = await this.dispatchCommand( - () => - this.acquireLock( - args.workerId, - assignmentId, - lockIncarnationId, - ttlSeconds, - ), - args, - 'Bridge assignment lock acquisition', - ); + locked = await this.dispatchCommand( + () => + this.acquireLock( + args.workerId, + assignmentId, + lockIncarnationId, + ttlSeconds, + ), + args, + 'Bridge assignment lock acquisition', + ); } if (!locked && admission != null) { - await delay( - Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), - args.signal, - ); + await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); } } while (!locked && admission != null); if (!locked) { @@ -887,21 +841,12 @@ export class RedisBridgeStore { current == null || current.registration.incarnationId !== registration.incarnationId || current.registration.identityId !== registration.identityId || - current.registration.binding?.tenantId !== - registration.binding?.tenantId + current.registration.binding?.tenantId !== registration.binding?.tenantId ) { - throw new BridgeStoreError( - 'WORKER_OFFLINE', - 'Bridge worker changed while the request was waiting', - ); + throw new BridgeStoreError('WORKER_OFFLINE', 'Bridge worker changed while the request was waiting'); } - if ( - !supportsWorkspaceTool(current.registration, args.workspaceRequest!) - ) { - throw new BridgeStoreError( - 'WORKER_MISMATCH', - 'Bridge worker capabilities changed while the request was waiting', - ); + if (!supportsWorkspaceTool(current.registration, args.workspaceRequest!)) { + throw new BridgeStoreError('WORKER_MISMATCH', 'Bridge worker capabilities changed while the request was waiting'); } } this.assertDispatchActive(args.signal, args.deadlineAtMs); @@ -925,7 +870,7 @@ export class RedisBridgeStore { ? {} : { workspaceLeaseSlot, - workspaceFence: `native-workspace:${args.workspaceRequest!.workspaceId}`, + workspaceFence: `native-workspace:${ args.workspaceRequest!.workspaceId}`, }), ...(registration.identityId != null ? { workerIdentityId: registration.identityId } @@ -987,10 +932,7 @@ export class RedisBridgeStore { } if ( args.workspaceRequest != null && - !supportsWorkspaceTool( - replacement.registration, - args.workspaceRequest, - ) + !supportsWorkspaceTool(replacement.registration, args.workspaceRequest) ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -1088,7 +1030,9 @@ export class RedisBridgeStore { } const deadline = Date.now() + waitMs; let firstPoll = true; - while (!signalAborted(signal) && (firstPoll || Date.now() < deadline)) { + while ( + !signalAborted(signal) && + (firstPoll || Date.now() < deadline)) { firstPoll = false; let assignmentId: string | null; try { @@ -1404,8 +1348,7 @@ export class RedisBridgeStore { ); if (existingSettlement === serializedSettlement && !quarantineWorkspace) return; - if ( - existingSettlement != null && + if (existingSettlement != null && existingSettlement !== serializedSettlement ) { throw new BridgeStoreError( @@ -1419,7 +1362,7 @@ export class RedisBridgeStore { 'Bridge settlement assignment read', ); if (assignment == null) { - if (existingSettlement === serializedSettlement) return; + if (existingSettlement === serializedSettlement) return; throw new BridgeStoreError( 'ASSIGNMENT_NOT_FOUND', 'Bridge assignment was not found', @@ -1435,7 +1378,7 @@ export class RedisBridgeStore { quarantineWorkspace && (assignment.workspaceFence == null || assignment.workspaceLeaseSlot === undefined || - settlement.status !== 'rejected') + settlement.status !== 'rejected') ) { throw new BridgeStoreError( 'ASSIGNMENT_INVALID', @@ -1495,7 +1438,7 @@ export class RedisBridgeStore { workerIncarnationKey(workerId), ); const script = [ - "local existing = redis.call('GET', KEYS[2])", + 'local existing = redis.call(\'GET\', KEYS[2])', 'if existing then', ' if existing == ARGV[1] then return 2 end', ' return -1', @@ -1505,12 +1448,12 @@ export class RedisBridgeStore { 'if ARGV[4] ~= "rejected" and redis.call(\'EXISTS\', KEYS[5]) == 0 then return -3 end', 'local stableIdentityKey = KEYS[#KEYS - 1]', 'if ARGV[5] ~= "" then', - " if redis.call('GET', stableIdentityKey) ~= ARGV[5] then return -4 end", - "elseif redis.call('EXISTS', stableIdentityKey) == 1 then return -4", + ' if redis.call(\'GET\', stableIdentityKey) ~= ARGV[5] then return -4 end', + 'elseif redis.call(\'EXISTS\', stableIdentityKey) == 1 then return -4', 'end', - "if redis.call('GET', KEYS[#KEYS]) ~= ARGV[7] then return -4 end", + 'if redis.call(\'GET\', KEYS[#KEYS]) ~= ARGV[7] then return -4 end', 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', - "if redis.call('GET', KEYS[3]) == ARGV[3] then redis.call('DEL', KEYS[3], KEYS[4]) end", + 'if redis.call(\'GET\', KEYS[3]) == ARGV[3] then redis.call(\'DEL\', KEYS[3], KEYS[4]) end', 'if ARGV[6] == "1" and ARGV[4] == "rejected" then', ' if ARGV[8] == "1" then redis.call(\'SET\', KEYS[6], "quarantined:" .. ARGV[3])', " else redis.call('DEL', KEYS[6]) end", @@ -1617,9 +1560,9 @@ export class RedisBridgeStore { const script = [ 'redis.call(\'SET\', KEYS[2], \"1\")', 'if #KEYS == 4 then redis.call(\'SET\', KEYS[4], \"1\") end', - "local current = redis.call('GET', KEYS[3])", + 'local current = redis.call(\'GET\', KEYS[3])', 'if current == ARGV[1] then', - " return redis.call('DEL', KEYS[1], KEYS[3])", + ' return redis.call(\'DEL\', KEYS[1], KEYS[3])', 'end', 'return 0', ].join('\n'); @@ -1632,7 +1575,12 @@ export class RedisBridgeStore { keys.push(workspaceQuarantineKey(workerId, runtimeSessionId)); } await boundedCommand( - this.redis.eval(script, keys.length, ...keys, incarnationId), + this.redis.eval( + script, + keys.length, + ...keys, + incarnationId, + ), this.redisCommandTimeoutMs, 'Bridge worker quarantine', ); @@ -1681,23 +1629,21 @@ export class RedisBridgeStore { workerId: string, ): Promise { const raw = await this.redis.get(workerKey(workerId)); - return raw == null - ? undefined - : (JSON.parse(raw) as RegisteredBridgeWorker); + return raw == null ? undefined : (JSON.parse(raw) as RegisteredBridgeWorker); } private async dispatchableRegistration( workerId: string, ): Promise< - { registration: RegisteredBridgeWorker; readyToken?: string } | undefined + | { registration: RegisteredBridgeWorker; readyToken?: string } + | undefined > { - const [raw, ready, generation, generationIncarnation] = - await this.redis.mget( - workerKey(workerId), - workerReadyKey(workerId), - workerRegistrationGenerationKey(workerId), - workerRegistrationGenerationIncarnationKey(workerId), - ); + const [raw, ready, generation, generationIncarnation] = await this.redis.mget( + workerKey(workerId), + workerReadyKey(workerId), + workerRegistrationGenerationKey(workerId), + workerRegistrationGenerationIncarnationKey(workerId), + ); if (raw == null) return undefined; const registration = JSON.parse(raw) as RegisteredBridgeWorker; if (registration.capabilities.requiresReadyConfirmation !== true) { @@ -1708,8 +1654,7 @@ export class RedisBridgeStore { !Number.isSafeInteger(registrationGeneration) || registrationGeneration < 1 || generationIncarnation !== registration.incarnationId || - ready !== - workerReadyToken(registration.incarnationId, registrationGeneration) + ready !== workerReadyToken(registration.incarnationId, registrationGeneration) ) { return undefined; } @@ -1777,11 +1722,11 @@ export class RedisBridgeStore { // Keep acknowledged assignment metadata for late clean rejection recovery, // but atomically revoke fulfillment when no settlement has won yet. const closeScript = [ - "local settlement = redis.call('GET', KEYS[2])", + 'local settlement = redis.call(\'GET\', KEYS[2])', 'if settlement then return settlement end', - "redis.call('DEL', KEYS[3])", - "if #KEYS == 4 and redis.call('GET', KEYS[4]) == ARGV[1] then return nil end", - "redis.call('DEL', KEYS[1])", + 'redis.call(\'DEL\', KEYS[3])', + 'if #KEYS == 4 and redis.call(\'GET\', KEYS[4]) == ARGV[1] then return nil end', + 'redis.call(\'DEL\', KEYS[1])', 'return nil', ].join('\n'); const finalSettlement = await boundedCommand( @@ -1815,7 +1760,12 @@ export class RedisBridgeStore { ? 30 : assignmentTtlSeconds(Date.parse(assignment.expiresAt)); await boundedCommand( - this.redis.set(cancellationKey(assignmentId), '1', 'EX', ttlSeconds), + this.redis.set( + cancellationKey(assignmentId), + '1', + 'EX', + ttlSeconds, + ), this.redisCommandTimeoutMs, 'Bridge assignment cancellation', ); @@ -1829,15 +1779,16 @@ export class RedisBridgeStore { const script = [ "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", 'if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[6]) ~= ARGV[7] then return 0 end', - "if #KEYS == 7 and redis.call('EXISTS', KEYS[7]) == 1 then return -1 end", + 'if #KEYS == 7 and redis.call(\'EXISTS\', KEYS[7]) == 1 then return -1 end', 'redis.call(\'SET\', KEYS[2], ARGV[2], \"EX\", ARGV[3])', - "redis.call('RPUSH', KEYS[3], ARGV[4])", - "redis.call('EXPIRE', KEYS[3], ARGV[3])", + 'redis.call(\'RPUSH\', KEYS[3], ARGV[4])', + 'redis.call(\'EXPIRE\', KEYS[3], ARGV[3])', ...(assignment.workspaceLeaseSlot === undefined - ? ['redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])'] + ? [ + 'redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])'] : []), 'redis.call(\'SET\', KEYS[5], "1", \"PXAT\", ARGV[6])', - "if #KEYS == 7 then redis.call('SET', KEYS[7], ARGV[4]) end", + 'if #KEYS == 7 then redis.call(\'SET\', KEYS[7], ARGV[4]) end', 'return 1', ].join('\n'); const keys = [ @@ -1888,7 +1839,7 @@ export class RedisBridgeStore { ttlSeconds: number, ): Promise { const script = [ - "if redis.call('EXISTS', KEYS[1]) == 1 then return 0 end", + 'if redis.call(\'EXISTS\', KEYS[1]) == 1 then return 0 end', 'redis.call(\'SET\', KEYS[1], ARGV[1], \"PX\", ARGV[3])', 'redis.call(\'SET\', KEYS[2], ARGV[2], \"PX\", ARGV[3])', 'return 1', @@ -1934,8 +1885,8 @@ export class RedisBridgeStore { } const runtimeSessionId = assignmentWorkspace(assignment)!; const script = [ - "if redis.call('GET', KEYS[1]) == ARGV[1] then", - " return redis.call('DEL', KEYS[1])", + 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1])', 'end', 'return 0', ].join('\n'); @@ -1944,7 +1895,10 @@ export class RedisBridgeStore { this.redis.eval( script, 1, - workspaceQuarantineKey(assignment.workerId, runtimeSessionId), + workspaceQuarantineKey( + assignment.workerId, + runtimeSessionId, + ), assignment.assignmentId, ), // Once settlement wins, caller cancellation must not prevent its @@ -2039,7 +1993,8 @@ export class RedisBridgeStore { if (cleanupResult !== -1) { await boundedCommand( assignment.workspaceLeaseSlot === undefined - ? this.releaseLock(assignment.workerId, assignment.assignmentId) + ? + this.releaseLock(assignment.workerId, assignment.assignmentId) : new BridgeWorkspaceSlots(this.redis).release( assignment.workerId, assignment.incarnationId, @@ -2056,8 +2011,8 @@ export class RedisBridgeStore { assignmentId: string, ): Promise { const script = [ - "if redis.call('GET', KEYS[1]) == ARGV[1] then", - " return redis.call('DEL', KEYS[1], KEYS[2])", + 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', + ' return redis.call(\'DEL\', KEYS[1], KEYS[2])', 'end', 'return 0', ].join('\n'); diff --git a/service/src/config.ts b/service/src/config.ts index d9f363f7..67a3561b 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -8,35 +8,12 @@ import { resolveExecutionProfileSource, } from './execution-profile'; -export const languageConfig: Record< - Languages | string, - t.LanguageConfig | undefined -> = { - [Languages.bash]: { - language: 'bash', - version: '5.2.0', - fileName: 'script.sh', - }, - [Languages.js]: { - language: 'bun-js', - version: '1.3.14', - fileName: 'index.js', - }, - [Languages.node]: { - language: 'node', - version: '24.15.0', - fileName: 'index.js', - }, - [Languages.py]: { - language: 'python', - version: '3.14.4', - fileName: 'main.py', - }, - [Languages.ts]: { - language: 'bun-ts', - version: '1.3.14', - fileName: 'main.ts', - }, +export const languageConfig: Record = { + [Languages.bash]: { language: 'bash', version: '5.2.0', fileName: 'script.sh' }, + [Languages.js]: { language: 'bun-js', version: '1.3.14', fileName: 'index.js' }, + [Languages.node]: { language: 'node', version: '24.15.0', fileName: 'index.js' }, + [Languages.py]: { language: 'python', version: '3.14.4', fileName: 'main.py' }, + [Languages.ts]: { language: 'bun-ts', version: '1.3.14', fileName: 'main.ts' }, }; const languageAliases: Record = { @@ -72,12 +49,8 @@ export function resolveLanguage(lang: string): Languages | undefined { } const defaultJobTimeoutMs = Number(process.env.JOB_TIMEOUT) || 300000; -const defaultMaxFileSize = - Number(process.env.MAX_FILE_SIZE) || 25 * 1024 * 1024; -const defaultExecutionManifestTtlSeconds = Math.min( - Math.ceil((defaultJobTimeoutMs + 60000) / 1000), - 600, -); +const defaultMaxFileSize = Number(process.env.MAX_FILE_SIZE) || 25 * 1024 * 1024; +const defaultExecutionManifestTtlSeconds = Math.min(Math.ceil((defaultJobTimeoutMs + 60000) / 1000), 600); const EGRESS_GRANT_GRACE_MS = 10 * 60 * 1000; /** Object-store listing and marker writes are metadata operations, not * checkpoint transfers. Bound each tightly so the post-exec checkpoint @@ -108,12 +81,11 @@ export function checkpointPipelineBudgetMs( checkpointTimeoutMs, CHECKPOINT_METADATA_TIMEOUT_CAP_MS, ); - return ( - launchTimeoutMs + - 2 * checkpointTimeoutMs + - 2 * metadataTimeoutMs + - POST_EXEC_CHECKPOINT_REGISTRY_COMMANDS * - RUNTIME_SESSION_REDIS_COMMAND_TIMEOUT_MS + return ( launchTimeoutMs + + 2 * checkpointTimeoutMs + + 2 * metadataTimeoutMs + + POST_EXEC_CHECKPOINT_REGISTRY_COMMANDS + * RUNTIME_SESSION_REDIS_COMMAND_TIMEOUT_MS ); } @@ -141,20 +113,16 @@ export function jobCompletionWaitTimeoutMs( backendCleanupTimeoutMs: number, egressRevokeTimeoutMs: number, ): number { - return ( - jobTimeoutMs + - backendCleanupTimeoutMs + - egressRevokeTimeoutMs + - WORKER_COMPLETION_OVERHEAD_MS + return ( jobTimeoutMs + + backendCleanupTimeoutMs + + egressRevokeTimeoutMs + + WORKER_COMPLETION_OVERHEAD_MS ); } export function parseArnList(raw: string | undefined): string[] | undefined { if (raw == null) return undefined; - const entries = raw - .split(',') - .map((entry) => entry.trim()) - .filter((entry) => entry.length > 0); + const entries = raw.split(',').map((entry) => entry.trim()).filter((entry) => entry.length > 0); return entries.length > 0 ? entries : undefined; } @@ -177,10 +145,7 @@ interface IntegerRange { max?: number; } -const lambdaMicrovmNumericRanges: Record< - keyof LambdaMicrovmNumericConfig, - IntegerRange -> = { +const lambdaMicrovmNumericRanges: Record = { LAMBDA_MICROVM_PORT: { min: 1, max: 65_535 }, LAMBDA_MICROVM_MAX_DURATION_SECONDS: { min: 1, max: 28_800 }, LAMBDA_MICROVM_IDLE_SECONDS: { min: 60, max: 28_800 }, @@ -217,20 +182,14 @@ export function resolveLambdaMicrovmNumericConfig( ): LambdaMicrovmNumericConfig { const read = (name: keyof LambdaMicrovmNumericConfig): number => { const raw = source[name]; - return raw == null || raw.trim() === '' - ? lambdaMicrovmNumericDefaults[name] - : Number(raw); + return raw == null || raw.trim() === '' ? lambdaMicrovmNumericDefaults[name] : Number(raw); }; return { LAMBDA_MICROVM_PORT: read('LAMBDA_MICROVM_PORT'), - LAMBDA_MICROVM_MAX_DURATION_SECONDS: read( - 'LAMBDA_MICROVM_MAX_DURATION_SECONDS', - ), + LAMBDA_MICROVM_MAX_DURATION_SECONDS: read('LAMBDA_MICROVM_MAX_DURATION_SECONDS'), LAMBDA_MICROVM_IDLE_SECONDS: read('LAMBDA_MICROVM_IDLE_SECONDS'), LAMBDA_MICROVM_SUSPEND_SECONDS: read('LAMBDA_MICROVM_SUSPEND_SECONDS'), - LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: read( - 'LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS', - ), + LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS: read('LAMBDA_MICROVM_AUTH_TOKEN_TTL_SECONDS'), LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: read('LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS'), LAMBDA_MICROVM_HEALTH_TIMEOUT_MS: read('LAMBDA_MICROVM_HEALTH_TIMEOUT_MS'), LAMBDA_MICROVM_LAUNCH_TPS: read('LAMBDA_MICROVM_LAUNCH_TPS'), @@ -242,28 +201,18 @@ export function resolveLambdaMicrovmNumericConfig( export function lambdaMicrovmNumericConfigError( config: LambdaMicrovmNumericConfig, ): string | undefined { - for (const name of Object.keys(lambdaMicrovmNumericRanges) as Array< - keyof LambdaMicrovmNumericConfig - >) { + for (const name of Object.keys(lambdaMicrovmNumericRanges) as Array) { const value = config[name]; const { min, max } = lambdaMicrovmNumericRanges[name]; - if ( - !Number.isSafeInteger(value) || - value < min || - (max != null && value > max) - ) { - const range = - max == null ? `at least ${min}` : `between ${min} and ${max}`; + if (!Number.isSafeInteger(value) || value < min || (max != null && value > max)) { + const range = max == null ? `at least ${min}` : `between ${min} and ${max}`; return `${name} must be a whole number ${range}`; } } return undefined; } -export function resolvePositiveIntEnv( - raw: string | undefined, - defaultValue: number, -): number { +export function resolvePositiveIntEnv(raw: string | undefined, defaultValue: number): number { if (raw == null || raw.trim() === '') { return defaultValue; } @@ -274,14 +223,8 @@ export function resolvePositiveIntEnv( return parsed; } -export function resolveEgressGrantTtlSeconds( - rawTtlSeconds: string | undefined, - jobTimeoutMs: number, -): number { - const defaultTtlSeconds = Math.max( - 1, - Math.ceil((jobTimeoutMs + EGRESS_GRANT_GRACE_MS) / 1000), - ); +export function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, jobTimeoutMs: number): number { + const defaultTtlSeconds = Math.max(1, Math.ceil((jobTimeoutMs + EGRESS_GRANT_GRACE_MS) / 1000)); if (rawTtlSeconds == null || rawTtlSeconds.trim() === '') { return defaultTtlSeconds; } @@ -294,16 +237,12 @@ export function resolveEgressGrantTtlSeconds( return Math.max(1, Math.ceil(configuredTtlSeconds)); } -const lambdaMicrovmNumericConfig = resolveLambdaMicrovmNumericConfig( - process.env, -); +const lambdaMicrovmNumericConfig = resolveLambdaMicrovmNumericConfig(process.env); export function hostedAppOperationTimeoutMs(): number { - return ( - env.CHECKPOINT_TIMEOUT_MS * 9 + - env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS * 7 + - env.HOSTED_APP_START_TIMEOUT_MS + - 30_000 + return ( env.CHECKPOINT_TIMEOUT_MS * 9 + + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS * 7 + + env.HOSTED_APP_START_TIMEOUT_MS + 30_000 ); } @@ -325,41 +264,39 @@ function configuredChoice( export function resolveSandboxBackend( raw: string | undefined, ): 'http' | 'lambda-microvm' | 'remote-bridge' { - return configuredChoice(raw, 'CODEAPI_SANDBOX_BACKEND', 'http', [ + return configuredChoice( + raw, + 'CODEAPI_SANDBOX_BACKEND', 'http', - 'lambda-microvm', - 'remote-bridge', - ]); + ['http', 'lambda-microvm', 'remote-bridge'], + ); } export function resolveRuntimeSessionMode( raw: string | undefined, ): 'stateless' | 'affinity' | 'strict' { - return configuredChoice(raw, 'CODEAPI_RUNTIME_SESSION_MODE', 'stateless', [ + return configuredChoice( + raw, + 'CODEAPI_RUNTIME_SESSION_MODE', 'stateless', - 'affinity', - 'strict', - ]); + ['stateless', 'affinity', 'strict'], + ); } export function resolveBridgeAuthMode( raw: string | undefined, ): 'static' | 'paired' { - return configuredChoice(raw, 'CODEAPI_BRIDGE_AUTH_MODE', 'static', [ + return configuredChoice( + raw, + 'CODEAPI_BRIDGE_AUTH_MODE', 'static', - 'paired', - ]); + ['static', 'paired'], + ); } -const sandboxBackend = resolveSandboxBackend( - process.env.CODEAPI_SANDBOX_BACKEND, -); -const runtimeSessionMode = resolveRuntimeSessionMode( - process.env.CODEAPI_RUNTIME_SESSION_MODE, -); -const bridgeAuthMode = resolveBridgeAuthMode( - process.env.CODEAPI_BRIDGE_AUTH_MODE, -); +const sandboxBackend = resolveSandboxBackend(process.env.CODEAPI_SANDBOX_BACKEND); +const runtimeSessionMode = resolveRuntimeSessionMode(process.env.CODEAPI_RUNTIME_SESSION_MODE); +const bridgeAuthMode = resolveBridgeAuthMode(process.env.CODEAPI_BRIDGE_AUTH_MODE); export const env = { PORT: process.env.SERVICE_PORT ?? 3112, @@ -367,61 +304,26 @@ export const env = { HARDENED_SANDBOX_MODE: process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', INSTANCE_ID: process.env.INSTANCE_ID ?? nanoid(), HTTP_JSON_LIMIT: process.env.CODEAPI_HTTP_JSON_LIMIT ?? '50mb', - SANDBOX_ENDPOINT: - process.env.SANDBOX_ENDPOINT ?? 'http://localhost:2000/api/v2', + SANDBOX_ENDPOINT: process.env.SANDBOX_ENDPOINT ?? 'http://localhost:2000/api/v2', EGRESS_GATEWAY_URL: process.env.EGRESS_GATEWAY_URL ?? '', FILE_SERVER_URL: process.env.FILE_SERVER_URL ?? 'http://localhost:3000', - TOOL_CALL_SERVER_URL: - process.env.TOOL_CALL_SERVER_URL ?? 'http://localhost:3033', + TOOL_CALL_SERVER_URL: process.env.TOOL_CALL_SERVER_URL ?? 'http://localhost:3033', EGRESS_GATEWAY_PORT: Number(process.env.EGRESS_GATEWAY_PORT) || 3190, - EGRESS_GATEWAY_FILE_SERVER_URL: - process.env.EGRESS_GATEWAY_FILE_SERVER_URL ?? - process.env.FILE_SERVER_URL ?? - 'http://localhost:3000', - EGRESS_GATEWAY_TOOL_CALL_SERVER_URL: - process.env.EGRESS_GATEWAY_TOOL_CALL_SERVER_URL ?? - process.env.TOOL_CALL_SERVER_URL ?? - 'http://localhost:3033', - EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES: - Number(process.env.EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES) || 1024 * 1024, + EGRESS_GATEWAY_FILE_SERVER_URL: process.env.EGRESS_GATEWAY_FILE_SERVER_URL ?? process.env.FILE_SERVER_URL ?? 'http://localhost:3000', + EGRESS_GATEWAY_TOOL_CALL_SERVER_URL: process.env.EGRESS_GATEWAY_TOOL_CALL_SERVER_URL ?? process.env.TOOL_CALL_SERVER_URL ?? 'http://localhost:3033', + EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES: Number(process.env.EGRESS_GATEWAY_MAX_TOOL_CALL_BYTES) || 1024 * 1024, // Per-entry / aggregate caps for PTC tool results persisted in `tool_history:` (see replay-state.ts). - PTC_MAX_TOOL_RESULT_BYTES: resolvePositiveIntEnv( - process.env.PTC_MAX_TOOL_RESULT_BYTES, - 5_000_000, - ), - PTC_MAX_TOOL_HISTORY_TOTAL_BYTES: resolvePositiveIntEnv( - process.env.PTC_MAX_TOOL_HISTORY_TOTAL_BYTES, - 40_000_000, - ), - EGRESS_GATEWAY_MAX_FILE_BYTES: - Number( - process.env.EGRESS_GATEWAY_MAX_FILE_BYTES ?? - process.env.SANDBOX_MAX_FILE_SIZE, - ) || 10_000_000, - EGRESS_GATEWAY_MAX_PATH_LENGTH: - Number( - process.env.EGRESS_GATEWAY_MAX_PATH_LENGTH ?? - process.env.SANDBOX_MAX_PATH_LENGTH, - ) || 256, - EGRESS_GATEWAY_MAX_NESTING_DEPTH: - Number( - process.env.EGRESS_GATEWAY_MAX_NESTING_DEPTH ?? - process.env.SANDBOX_MAX_NESTING_DEPTH, - ) || 10, - EGRESS_GATEWAY_REQUEST_TIMEOUT_MS: - Number(process.env.EGRESS_GATEWAY_REQUEST_TIMEOUT_MS) || 30_000, - EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: - Number(process.env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS) || 5_000, - EGRESS_LEDGER_REQUIRED: - process.env.CODEAPI_EGRESS_LEDGER_REQUIRED === 'true' || - process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', - EGRESS_LEDGER_TTL_GRACE_SECONDS: - Number(process.env.CODEAPI_EGRESS_LEDGER_TTL_GRACE_SECONDS) || 300, + PTC_MAX_TOOL_RESULT_BYTES: resolvePositiveIntEnv(process.env.PTC_MAX_TOOL_RESULT_BYTES, 5_000_000), + PTC_MAX_TOOL_HISTORY_TOTAL_BYTES: resolvePositiveIntEnv(process.env.PTC_MAX_TOOL_HISTORY_TOTAL_BYTES, 40_000_000), + EGRESS_GATEWAY_MAX_FILE_BYTES: Number(process.env.EGRESS_GATEWAY_MAX_FILE_BYTES ?? process.env.SANDBOX_MAX_FILE_SIZE) || 10_000_000, + EGRESS_GATEWAY_MAX_PATH_LENGTH: Number(process.env.EGRESS_GATEWAY_MAX_PATH_LENGTH ?? process.env.SANDBOX_MAX_PATH_LENGTH) || 256, + EGRESS_GATEWAY_MAX_NESTING_DEPTH: Number(process.env.EGRESS_GATEWAY_MAX_NESTING_DEPTH ?? process.env.SANDBOX_MAX_NESTING_DEPTH) || 10, + EGRESS_GATEWAY_REQUEST_TIMEOUT_MS: Number(process.env.EGRESS_GATEWAY_REQUEST_TIMEOUT_MS) || 30_000, + EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: Number(process.env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS) || 5_000, + EGRESS_LEDGER_REQUIRED: process.env.CODEAPI_EGRESS_LEDGER_REQUIRED === 'true' || process.env.CODEAPI_HARDENED_SANDBOX_MODE === 'true', + EGRESS_LEDGER_TTL_GRACE_SECONDS: Number(process.env.CODEAPI_EGRESS_LEDGER_TTL_GRACE_SECONDS) || 300, EGRESS_GRANT_SECRET: process.env.CODEAPI_EGRESS_GRANT_SECRET ?? '', - EGRESS_GRANT_TTL_SECONDS: resolveEgressGrantTtlSeconds( - process.env.EGRESS_GRANT_TTL_SECONDS, - defaultJobTimeoutMs, - ), + EGRESS_GRANT_TTL_SECONDS: resolveEgressGrantTtlSeconds(process.env.EGRESS_GRANT_TTL_SECONDS, defaultJobTimeoutMs), PYTHON_CONCURRENCY: Number(process.env.PYTHON_CONCURRENCY) || 1, OTHER_CONCURRENCY: Number(process.env.OTHER_CONCURRENCY) || 8, JOB_WINDOW: Number(process.env.JOB_WINDOW) || 1000, @@ -460,32 +362,22 @@ export const env = { * working; multi-tenant deploys MUST set this to `true` before any tenant * is multi-homed, otherwise a missing tenantId would silently bucket * cross-tenant requests under the same `'legacy'` prefix. */ - TENANT_ISOLATION_STRICT: - process.env.CODEAPI_TENANT_ISOLATION_STRICT === 'true', + TENANT_ISOLATION_STRICT: process.env.CODEAPI_TENANT_ISOLATION_STRICT === 'true', // Signed execution manifests. Prefer private/public key mode for split-runner // deployments so sandbox-runner receives only a verifier, not a signing secret. - EXECUTION_MANIFEST_PRIVATE_KEY: - process.env.CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY ?? '', - EXECUTION_MANIFEST_PUBLIC_KEY: - process.env.CODEAPI_EXECUTION_MANIFEST_PUBLIC_KEY ?? '', + EXECUTION_MANIFEST_PRIVATE_KEY: process.env.CODEAPI_EXECUTION_MANIFEST_PRIVATE_KEY ?? '', + EXECUTION_MANIFEST_PUBLIC_KEY: process.env.CODEAPI_EXECUTION_MANIFEST_PUBLIC_KEY ?? '', // Legacy HMAC fallback for non-split deployments. Do not mount into sandbox-runner. - EXECUTION_MANIFEST_SECRET: - process.env.CODEAPI_EXECUTION_MANIFEST_SECRET ?? '', + EXECUTION_MANIFEST_SECRET: process.env.CODEAPI_EXECUTION_MANIFEST_SECRET ?? '', EXECUTION_MANIFEST_TTL_SECONDS: Math.min( - Number(process.env.EXECUTION_MANIFEST_TTL_SECONDS) || - defaultExecutionManifestTtlSeconds, + Number(process.env.EXECUTION_MANIFEST_TTL_SECONDS) || defaultExecutionManifestTtlSeconds, 600, ), - EXECUTION_MANIFEST_MAX_UPLOAD_BYTES: - Number(process.env.EXECUTION_MANIFEST_MAX_UPLOAD_BYTES) || - defaultMaxFileSize, - EXECUTION_MANIFEST_MAX_OUTPUT_FILES: - Number(process.env.EXECUTION_MANIFEST_MAX_OUTPUT_FILES) || 50, - EXECUTION_MANIFEST_MAX_REQUESTS: - Number(process.env.EXECUTION_MANIFEST_MAX_REQUESTS) || 1000, + EXECUTION_MANIFEST_MAX_UPLOAD_BYTES: Number(process.env.EXECUTION_MANIFEST_MAX_UPLOAD_BYTES) || defaultMaxFileSize, + EXECUTION_MANIFEST_MAX_OUTPUT_FILES: Number(process.env.EXECUTION_MANIFEST_MAX_OUTPUT_FILES) || 50, + EXECUTION_MANIFEST_MAX_REQUESTS: Number(process.env.EXECUTION_MANIFEST_MAX_REQUESTS) || 1000, // Redis - Alternative DNS Lookup for AWS ElastiCache TLS connections - REDIS_USE_ALTERNATIVE_DNS_LOOKUP: - process.env.REDIS_USE_ALTERNATIVE_DNS_LOOKUP === 'true', + REDIS_USE_ALTERNATIVE_DNS_LOOKUP: process.env.REDIS_USE_ALTERNATIVE_DNS_LOOKUP === 'true', /** * Programmatic Tool Calling execution model. * - `replay` (default): Temporal-style replay. Sandbox exits between round-trips; @@ -496,9 +388,7 @@ export const env = { * via a long-polling HTTP callback through the Tool Call Server. Retained as * an explicit opt-in during rollout; scheduled for removal in a follow-up. */ - PTC_MODE: (process.env.PTC_MODE === 'blocking' ? 'blocking' : 'replay') as - | 'replay' - | 'blocking', + PTC_MODE: (process.env.PTC_MODE === 'blocking' ? 'blocking' : 'replay') as 'replay' | 'blocking', PTC_DEBUG: process.env.PTC_DEBUG === 'true', /** * Sandbox execution backend. @@ -548,20 +438,14 @@ export const env = { ), // Lambda MicroVM backend. Connector lists are comma-separated ARNs. LAMBDA_MICROVM_IMAGE_ARN: process.env.LAMBDA_MICROVM_IMAGE_ARN ?? '', - LAMBDA_MICROVM_IMAGE_VERSION: - process.env.LAMBDA_MICROVM_IMAGE_VERSION || undefined, - LAMBDA_MICROVM_EXECUTION_ROLE_ARN: - process.env.LAMBDA_MICROVM_EXECUTION_ROLE_ARN || undefined, + LAMBDA_MICROVM_IMAGE_VERSION: process.env.LAMBDA_MICROVM_IMAGE_VERSION || undefined, + LAMBDA_MICROVM_EXECUTION_ROLE_ARN: process.env.LAMBDA_MICROVM_EXECUTION_ROLE_ARN || undefined, /* Runtime VM stdout reaches CloudWatch only when RunMicrovm sends a logging * config AND an executionRoleArn is set — pairs with the role above. */ LAMBDA_MICROVM_LOG_GROUP: process.env.LAMBDA_MICROVM_LOG_GROUP || undefined, LAMBDA_MICROVM_REGION: process.env.LAMBDA_MICROVM_REGION || undefined, - LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS: parseArnList( - process.env.LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS, - ), - LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS: parseArnList( - process.env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS, - ), + LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS: parseArnList(process.env.LAMBDA_MICROVM_INGRESS_CONNECTOR_ARNS), + LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS: parseArnList(process.env.LAMBDA_MICROVM_EGRESS_CONNECTOR_ARNS), ...lambdaMicrovmNumericConfig, /* CreateMicrovmAuthToken is minted per execute + per checkpoint; share a * fleet-wide budget so concurrent warm-session executes queue instead of @@ -575,19 +459,14 @@ export const env = { process.env.CODEAPI_CHECKPOINT_MAX_BYTES, 512 * 1024 * 1024, ), - CHECKPOINT_TIMEOUT_MS: configuredNumber( - process.env.CODEAPI_CHECKPOINT_TIMEOUT_MS, - 60_000, - ), - CHECKPOINT_PREFIX: - process.env.CODEAPI_CHECKPOINT_PREFIX ?? 'rtsx-checkpoints/', + CHECKPOINT_TIMEOUT_MS: configuredNumber(process.env.CODEAPI_CHECKPOINT_TIMEOUT_MS, 60_000), + CHECKPOINT_PREFIX: process.env.CODEAPI_CHECKPOINT_PREFIX ?? 'rtsx-checkpoints/', /** Dedicated Lambda MicroVM resident-server fleet. This remains an explicit * stateful-stack capability; the ordinary/default HTTP profile never starts * or preserves application processes. */ HOSTED_APPS_ENABLED: process.env.CODEAPI_HOSTED_APPS_ENABLED === 'true', HOSTED_APP_IMAGE_ARN: process.env.LAMBDA_MICROVM_APP_IMAGE_ARN ?? '', - HOSTED_APP_IMAGE_VERSION: - process.env.LAMBDA_MICROVM_APP_IMAGE_VERSION || undefined, + HOSTED_APP_IMAGE_VERSION: process.env.LAMBDA_MICROVM_APP_IMAGE_VERSION || undefined, /* These values are part of the pinned app-host image contract. RunMicrovm * cannot inject environment variables into the image, so exposing overrides * here would only make the control plane call ports the runner never opened. */ @@ -606,10 +485,8 @@ export const env = { 900, ), HOSTED_APP_START_TIMEOUT_MS: 30_000 as number, - HOSTED_APP_CREDENTIAL_KEY: - process.env.CODEAPI_HOSTED_APP_CREDENTIAL_KEY ?? '', - HOSTED_APP_PREVIEW_ORIGIN: - process.env.CODEAPI_HOSTED_APP_PREVIEW_ORIGIN ?? '', + HOSTED_APP_CREDENTIAL_KEY: process.env.CODEAPI_HOSTED_APP_CREDENTIAL_KEY ?? '', + HOSTED_APP_PREVIEW_ORIGIN: process.env.CODEAPI_HOSTED_APP_PREVIEW_ORIGIN ?? '', HOSTED_APP_PREVIEW_SIGNING_KEY: process.env.CODEAPI_HOSTED_APP_PREVIEW_SIGNING_KEY ?? '', }; @@ -632,9 +509,7 @@ type PlanLimits = { * JSON object keyed by the `plan_id` JWT claim. Unknown or absent plan ids * fall back to the default tier, which is the only entry defined in code. */ -export function parsePlanLimits( - raw: string | undefined, -): Record { +export function parsePlanLimits(raw: string | undefined): Record { if (raw == null || raw.trim() === '') { return {}; } @@ -642,14 +517,10 @@ export function parsePlanLimits( try { parsed = JSON.parse(raw); } catch (error) { - throw new Error( - `CODEAPI_PLAN_LIMITS is not valid JSON: ${(error as Error).message}`, - ); + throw new Error(`CODEAPI_PLAN_LIMITS is not valid JSON: ${(error as Error).message}`); } if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) { - throw new Error( - 'CODEAPI_PLAN_LIMITS must be a JSON object keyed by plan id', - ); + throw new Error('CODEAPI_PLAN_LIMITS must be a JSON object keyed by plan id'); } return parsed as Record; } @@ -657,8 +528,7 @@ export function parsePlanLimits( export const planLimits: PlanLimits = { ...parsePlanLimits(process.env.CODEAPI_PLAN_LIMITS), default: { - run_memory_limit: - Number(process.env.SANDBOX_RUN_MEMORY_LIMIT) || default_run_memory_limit, + run_memory_limit: Number(process.env.SANDBOX_RUN_MEMORY_LIMIT) || default_run_memory_limit, max_file_size: env.MAX_FILE_SIZE, }, }; From e3348dc2816917da0847f4950502de254bd4adb8 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 8 Sep 2026 22:51:38 -0400 Subject: [PATCH 3/7] fix: Enforce Replica Slot Ceilings at Dispatch --- packages/code/src/cli.ts | 145 ++++++------- packages/code/src/worker-slots.test.ts | 52 +++++ packages/code/src/worker.ts | 7 +- service/src/bridge/concurrent-store.test.ts | 133 ++++++++++-- service/src/bridge/router.ts | 181 ++++++++-------- service/src/bridge/store.ts | 224 ++++++++++++-------- service/src/config.ts | 15 +- 7 files changed, 475 insertions(+), 282 deletions(-) diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index a9f59556..f88784d9 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -482,12 +482,12 @@ async function run( const roots: LocalWorkspaceConfig[] = canonicalWorkerDirectory ? [ { - id: workspaceId, + id: workspaceId, root: canonicalWorkerDirectory, - writable: allowWorkspaceWrites, + writable: allowWorkspaceWrites, name: - option(args, '--workspace-name') ?? - process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? + option(args, '--workspace-name') ?? + process.env.LIBRECHAT_CODE_WORKSPACE_NAME?.trim() ?? (useDefaultWorkspace ? workspaceId : defaultWorkspaceName(workerDirectory!, workspaceId)), @@ -511,7 +511,8 @@ async function run( const separator = value.indexOf('='); if ( separator < 1 || - separator === value.length - 1 || !canonicalWorkerDirectory || + separator === value.length - 1 || + !canonicalWorkerDirectory || commandSandboxMode !== 'native-srt' ) { throw new Error( @@ -521,7 +522,7 @@ async function run( roots.push({ id: value.slice(0, separator), root: await realpath(value.slice(separator + 1)), - writable: allowWorkspaceWrites, + writable: allowWorkspaceWrites, }); } // Aliases and nested grants are not independent execution domains. @@ -555,7 +556,7 @@ async function run( } if ( workspaceLeaseSlots > 1 && - ( !allowWorkspaceCommands || commandSandboxMode !== 'native-srt') + (!allowWorkspaceCommands || commandSandboxMode !== 'native-srt') ) { throw new Error('Concurrent workspace leases require native-srt commands'); } @@ -748,38 +749,38 @@ async function run( statefulWorkspace, }); const nativeOptions: NativeProcessSandboxOptions = { - workspaceRoot: canonicalWorkerDirectory!, + workspaceRoot: canonicalWorkerDirectory!, protectedPaths: [ identityPath, ...rootQuarantinePaths.values(), - github.privateKeyPath, + github.privateKeyPath, ].filter((path): path is string => path != null), - allowedDomains: commandAllowedDomains, - ...(github.provider - ? { - maskedEnvironment: { - variables: [ - { - name: GITHUB_CREDENTIAL_ENV_NAME, - extract: '^(.+)$', - injectHosts: [github.host], - }, - ], - async resolve(signal?: AbortSignal) { - return gitHubCredentialEnvironment( - await github.provider!.getCredential(signal), - ); - }, - wrapCommand(command: string, platform: NodeJS.Platform) { - return wrapGitHubCredentialCommand( - command, - github.host, - platform, - ); - }, - }, - } - : {}), + allowedDomains: commandAllowedDomains, + ...(github.provider + ? { + maskedEnvironment: { + variables: [ + { + name: GITHUB_CREDENTIAL_ENV_NAME, + extract: '^(.+)$', + injectHosts: [github.host], + }, + ], + async resolve(signal?: AbortSignal) { + return gitHubCredentialEnvironment( + await github.provider!.getCredential(signal), + ); + }, + wrapCommand(command: string, platform: NodeJS.Platform) { + return wrapGitHubCredentialCommand( + command, + github.host, + platform, + ); + }, + }, + } + : {}), }; const nativeCommandSandbox = allowWorkspaceCommands && commandSandboxMode === 'native-srt' @@ -787,7 +788,7 @@ async function run( ? new NativeWorkspaceCommandPool( new Map( roots.map((root) => [ - root.id, + root.id, { ...nativeOptions, workspaceRoot: root.root }, ]), ), @@ -860,11 +861,11 @@ async function run( ? { workspaceQuarantines: new Map( roots.map((root) => [ - root.id, + root.id, workspaceMutationGuard( rootQuarantinePaths.get(root.id)!, workerId, - root.id, + root.id, incarnationId, ), ]), @@ -876,41 +877,41 @@ async function run( workspaceLeaseSlots === 1 && roots.length === 1 ? { - async assertAvailable() { - const record = await loadWorkspaceMutationQuarantine( - mutationQuarantinePath, - ); - if (record != null) { - throw new BridgeProtocolError( - `Workspace mutations are quarantined since ${record.quarantinedAt}: ${record.reason}. Inspect or restore the workspace, then run librechat-code clear-workspace-quarantine`, - undefined, - 'WORKER_QUARANTINED', + async assertAvailable() { + const record = await loadWorkspaceMutationQuarantine( + mutationQuarantinePath, ); - } - }, - async arm(reason) { - await saveWorkspaceMutationQuarantine(mutationQuarantinePath, { - version: 1, - workerId, - workspaceId, - ownerId: incarnationId, - quarantinedAt: new Date().toISOString(), - reason, - }); - }, - async clear() { - await clearWorkspaceMutationQuarantine( - mutationQuarantinePath, - incarnationId, - ); - }, - async quarantine() { - await assertWorkspaceMutationQuarantineOwner( - mutationQuarantinePath, - incarnationId, - ); - }, - } + if (record != null) { + throw new BridgeProtocolError( + `Workspace mutations are quarantined since ${record.quarantinedAt}: ${record.reason}. Inspect or restore the workspace, then run librechat-code clear-workspace-quarantine`, + undefined, + 'WORKER_QUARANTINED', + ); + } + }, + async arm(reason) { + await saveWorkspaceMutationQuarantine(mutationQuarantinePath, { + version: 1, + workerId, + workspaceId, + ownerId: incarnationId, + quarantinedAt: new Date().toISOString(), + reason, + }); + }, + async clear() { + await clearWorkspaceMutationQuarantine( + mutationQuarantinePath, + incarnationId, + ); + }, + async quarantine() { + await assertWorkspaceMutationQuarantineOwner( + mutationQuarantinePath, + incarnationId, + ); + }, + } : undefined, onIdentityChange: pairedIdentity && identityPath diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts index 47498789..4ddb6423 100644 --- a/packages/code/src/worker-slots.test.ts +++ b/packages/code/src/worker-slots.test.ts @@ -128,3 +128,55 @@ for (const cancelled of [false, true]) { assert.equal(internals.activeWorkspaceAssignments.get('a')?.id, 'previous'); }); } + +test('a local cleanup handoff preserves the new assignment owner and remaining budget', async () => { + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'fixture', + runtimes: [], + }, + }); + const internals = worker as unknown as { + activeWorkspaceAssignments: Map< + string, + { id: string; done: Promise } + >; + executeOwned: (assignment: BridgeAssignment) => Promise; + }; + let release!: () => void; + internals.activeWorkspaceAssignments.set('a', { + id: 'previous', + done: new Promise((resolve) => { + release = resolve; + }), + }); + let executed = false; + internals.executeOwned = async (assignment) => { + executed = true; + assert.equal(internals.activeWorkspaceAssignments.get('a')?.id, 'next'); + assert.ok(assignment.remainingMs! < 1000 && assignment.remainingMs! > 0); + }; + const pending = worker.executeAndSettle({ + assignmentId: 'next', + executionKind: 'workspace_tool', + remainingMs: 1000, + request: { + protocolVersion: 1, + workspaceId: 'a', + operation: 'read_file', + path: 'test.txt', + }, + } as BridgeAssignment); + await new Promise((resolve) => setTimeout(resolve, 5)); + assert.equal(executed, false); + internals.activeWorkspaceAssignments.delete('a'); + release(); + await pending; + assert.equal(executed, true); + assert.equal(internals.activeWorkspaceAssignments.size, 0); +}); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 068422fc..e37c3b88 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -1054,7 +1054,8 @@ export class BridgeWorker { error instanceof BridgeProtocolError && (error.status === 401 || error.status === 403); const credentialRemainingMs = - Date.parse(identity.expiresAt) - (Date.now() + serverClockOffsetMs); + Date.parse(identity.expiresAt) - + (Date.now() + serverClockOffsetMs); if (terminal || credentialRemainingMs <= 0) throw error; await abortableDelay( Math.min( @@ -1630,9 +1631,7 @@ export class BridgeWorker { return await lease.execute({ body, headers, signal }); } if (lease.endpoint == null) { - throw new BridgeProtocolError( - 'Runtime lease does not provide an execution transport', - ); + throw new BridgeProtocolError('Runtime lease does not provide an execution transport'); } const endpoint = lease.endpoint.replace(/\/+$/, ''); const response = await this.fetchImpl(`${endpoint}/execute`, { diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts index 84f41d0f..177b4f74 100644 --- a/service/src/bridge/concurrent-store.test.ts +++ b/service/src/bridge/concurrent-store.test.ts @@ -12,7 +12,7 @@ const incarnationId = 'concurrent-incarnation'; afterEach(async () => { await redis.flushall(); }); -async function register() { +async function register(workspaceLeaseSlots = 2) { const generation = await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId, @@ -22,7 +22,7 @@ async function register() { runtimes: [], sandboxProfile: 'native-srt', requiresReadyConfirmation: true, - workspaceLeaseSlots: 2, + workspaceLeaseSlots, workspaceTools: { protocolVersion: BRIDGE_PROTOCOL_VERSION, operations: ['read_file'], @@ -96,6 +96,50 @@ test('store routes simultaneous roots through separate acknowledged slots', asyn ).toBeNull(); }); +test('a replica never dispatches above its own configured ceiling', async () => { + await register(); + const serialReplica = new RedisBridgeStore(redis, 60, 1000, 1); + await expect( + serialReplica.dispatchWorkspaceTool({ + workerId, + signal: new AbortController().signal, + deadlineAtMs: Date.now() + 1000, + request: { + protocolVersion: 1, + operation: 'read_file', + workspaceId: 'a', + path: 'file.txt', + }, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + await expect( + serialReplica.lease(workerId, incarnationId, 0, undefined, undefined, 1), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect( + await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`), + ).toBeNull(); +}); + +test('capacity changes require the active slots to drain', async () => { + await register(); + const pending = dispatch('a'); + const assignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await expect(register(1)).rejects.toMatchObject({ code: 'WORKER_BUSY' }); + await settle(assignment); + await pending; + await register(1); + expect( + (await store.workerStatus(workerId)).capabilities?.workspaceLeaseSlots, + ).toBe(1); +}); + test('same-root work waits while another root progresses', async () => { await register(); const a = dispatch('a'); @@ -137,38 +181,93 @@ test('same-root work waits while another root progresses', async () => { test('queued cancellation never leases and does not block another root', async () => { await register(); const a = dispatch('a'); - const first = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 0))!; + const first = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; const controller = new AbortController(); const cancelled = dispatch('a', controller.signal); const queue = `codeapi:bridge:v1:worker:${workerId}:admission`; - for (let i = 0; i < 100 && await redis.zcard(queue) < 2; i++) { - await new Promise(resolve => setTimeout(resolve, 5)); + for (let i = 0; i < 100 && (await redis.zcard(queue)) < 2; i++) { + await new Promise((resolve) => setTimeout(resolve, 5)); } expect(await redis.zcard(queue)).toBe(2); controller.abort(); await expect(cancelled).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); const b = dispatch('b'); - const second = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 1))!; + const second = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 1, + ))!; expect(second.request).toMatchObject({ workspaceId: 'b' }); - await settle(first); await settle(second); await Promise.all([a, b]); - expect(await store.lease(workerId, incarnationId, 0, undefined, undefined, 0)).toBeUndefined(); + await settle(first); + await settle(second); + await Promise.all([a, b]); + expect( + await store.lease(workerId, incarnationId, 0, undefined, undefined, 0), + ).toBeUndefined(); }); test('late quarantine releases its slot after caller cancellation and retains only its root fence', async () => { await register(); const controller = new AbortController(); const a = dispatch('a', controller.signal); - const first = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 0))!; - await store.acknowledgeLease(workerId, incarnationId, first.assignmentId, first.generation, first.leaseToken); + const first = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await store.acknowledgeLease( + workerId, + incarnationId, + first.assignmentId, + first.generation, + first.leaseToken, + ); controller.abort(); await expect(a).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); - await store.settle(workerId, first.assignmentId, { protocolVersion: 1, - incarnationId, generation: first.generation, leaseToken: first.leaseToken, - status: 'rejected', error: 'uncertain mutation' }, undefined, undefined, true); - expect(await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`)).toBeNull(); + await store.settle( + workerId, + first.assignmentId, + { + protocolVersion: 1, + incarnationId, + generation: first.generation, + leaseToken: first.leaseToken, + status: 'rejected', + error: 'uncertain mutation', + }, + undefined, + undefined, + true, + ); + expect( + await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`), + ).toBeNull(); const b = dispatch('b'); - const next = (await store.lease(workerId, incarnationId, 1000, undefined, undefined, 0))!; + const next = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; expect(next.request).toMatchObject({ workspaceId: 'b' }); - await settle(next); await b; - await expect(dispatch('a')).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + await settle(next); + await b; + await expect(dispatch('a')).rejects.toMatchObject({ + code: 'WORKSPACE_QUARANTINED', + }); }); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 4410eb89..90ded8a8 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -330,10 +330,10 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }), ); - router.post( +router.post( '/workers/register', - workerAuth, - asyncRoute(async (req, res) => { + workerAuth, + asyncRoute(async (req, res) => { const registration = req.body as unknown; if ( !isRecord(registration) || @@ -384,21 +384,21 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ? { binding: authorization.binding } : {}), }; - try { + try { const registrationGeneration = await options.store.register( trustedRegistration, authorization, ); - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - workerId: registration.workerId, - incarnationId: registration.incarnationId, + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, registrationGeneration, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, - workspaceLeaseSlots: options.store.workspaceLeaseCapacity( - registration.capabilities.workspaceLeaseSlots, - ), + workspaceLeaseSlots: options.store.workspaceLeaseCapacity( + registration.capabilities.workspaceLeaseSlots, + ), supportedWorkspaceToolOperations: [ 'read_file', 'search_text', @@ -412,16 +412,16 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { supportedWorkspaceEditFileModes: ['single', 'batch'], supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], supportedWorkspaceListFileFeatures: ['after_path'], - }); - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); - return; - } - throw error; + }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; } - }), - ); + throw error; + } + }), +); router.post( '/workers/:workerId/ready', @@ -519,55 +519,53 @@ router.post( }), ); - router.post( +router.post( '/workers/:workerId/lease', - workerAuth, - asyncRoute(async (req, res) => { + workerAuth, + asyncRoute(async (req, res) => { const requestStartedAtMs = Date.now(); - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; const requestedWait = Number(body.waitMs ?? 25_000); - if ( + if ( !validWorkerId(workerId) || body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || !validIncarnationId(body.incarnationId) || !Number.isFinite(requestedWait) || requestedWait < 0 || - (body.workspaceLeaseSlot !== undefined && - (!Number.isSafeInteger(body.workspaceLeaseSlot) || - Number(body.workspaceLeaseSlot) < 0 || - Number(body.workspaceLeaseSlot) >= 8)) - ) { + (body.workspaceLeaseSlot !== undefined && + (!Number.isSafeInteger(body.workspaceLeaseSlot) || + Number(body.workspaceLeaseSlot) < 0 || + Number(body.workspaceLeaseSlot) >= 8)) + ) { res.status(400).json({ error: 'Invalid bridge lease request' }); return; } - if (!configuredWorker(workerId)) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - try { + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + try { const leaseController = new AbortController(); const abortLease = (): void => leaseController.abort(); req.once('aborted', abortLease); res.once('close', abortLease); let assignment: CodeBridgeAssignment | undefined; - try { - assignment = await options.store.lease( - workerId, - body.incarnationId, + try { + assignment = await options.store.lease( + workerId, + body.incarnationId, Math.min(requestedWait, MAX_LEASE_WAIT_MS), - leaseController.signal, - ( - res.locals.bridgeWorkerAuthorization as - | { identityId: string } - | undefined - )?.identityId, - body.workspaceLeaseSlot === undefined - ? undefined - : Number(body.workspaceLeaseSlot), - ); + leaseController.signal, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, + body.workspaceLeaseSlot === undefined ? undefined : Number(body.workspaceLeaseSlot), + ); if (leaseController.signal.aborted) { if (assignment != null) await options.store.returnLease(assignment); return; @@ -577,19 +575,19 @@ router.post( serverElapsedMs: Math.max(0, Date.now() - requestStartedAtMs), assignment, }); - } finally { + } finally { req.off('aborted', abortLease); res.off('close', abortLease); } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); - return; - } - throw error; + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; } - }), - ); + throw error; + } + }), +); router.post( '/workers/:workerId/assignments/:assignmentId/ack', @@ -639,55 +637,55 @@ router.post( }), ); - router.post( - [ - '/workers/:workerId/assignments/:assignmentId/settle', - '/workers/:workerId/assignments/:assignmentId/quarantine', - ], - workerAuth, - asyncRoute(async (req, res) => { +router.post( + [ + '/workers/:workerId/assignments/:assignmentId/settle', + '/workers/:workerId/assignments/:assignmentId/quarantine', + ], + workerAuth, + asyncRoute(async (req, res) => { const settlement = req.body as unknown; if (!isSettlement(settlement)) { res.status(400).json({ error: 'Invalid bridge settlement' }); return; } - try { + try { const settlementController = new AbortController(); const abortSettlement = (): void => settlementController.abort(); req.once('aborted', abortSettlement); res.once('close', abortSettlement); - try { - await options.store.settle( - req.params.workerId, - req.params.assignmentId, - settlement, - settlementController.signal, - ( - res.locals.bridgeWorkerAuthorization as - | { identityId: string } - | undefined - )?.identityId, - req.path.endsWith('/quarantine'), - ); + try { + await options.store.settle( + req.params.workerId, + req.params.assignmentId, + settlement, + settlementController.signal, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, + req.path.endsWith('/quarantine'), + ); if (!settlementController.signal.aborted) { res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, accepted: true, }); } - } finally { + } finally { req.off('aborted', abortSettlement); res.off('close', abortSettlement); } - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); - return; - } - throw error; + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; } - }), - ); + throw error; + } + }), +); router.post( '/workers/:workerId/assignments/:assignmentId/cancellation', @@ -746,5 +744,6 @@ router.post( }), ); + return router; } diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index f7404934..c03048a1 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -412,15 +412,18 @@ export class RedisBridgeStore { async register( registration: RegisteredBridgeWorker, - authorization?: string | { - identityId?: string; - pairingGeneration?: number; - activeCredentialId?: string; - }, + authorization?: + | string + | { + identityId?: string; + pairingGeneration?: number; + activeCredentialId?: string; + }, ): Promise { if (registration.capabilities.workspaceLeaseSlots !== undefined) { if ( - !isValidBridgeWorkerCapabilities(registration.capabilities) ||registration.capabilities.requiresReadyConfirmation !== true + !isValidBridgeWorkerCapabilities(registration.capabilities) || + registration.capabilities.requiresReadyConfirmation !== true ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -448,12 +451,12 @@ export class RedisBridgeStore { ' local pairingGeneration = redis.call(\'GET\', KEYS[7]) or "0"', ' if pairingGeneration ~= ARGV[5] then return -5 end', ' if ARGV[6] ~= "" then', - ' if redis.call(\'GET\', KEYS[8]) ~= ARGV[6] then return -5 end', + " if redis.call('GET', KEYS[8]) ~= ARGV[6] then return -5 end", ' elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -5', ' end', 'end', 'if ARGV[8] ~= "" then', - ' local stableIdentity = redis.call(\'GET\', KEYS[8])', + " local stableIdentity = redis.call('GET', KEYS[8])", ' if stableIdentity and stableIdentity ~= ARGV[8] then return -4 end', ' if not stableIdentity then', ' if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4 end', @@ -461,26 +464,26 @@ export class RedisBridgeStore { ' end', 'elseif ARGV[7] ~= "" and redis.call(\'GET\', KEYS[9]) ~= ARGV[7] then return -4', 'end', - 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', - 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', - 'local current = redis.call(\'GET\', KEYS[4])', + "if redis.call('EXISTS', KEYS[3]) == 1 then return -2 end", + "if redis.call('EXISTS', KEYS[2]) == 1 then return -1 end", + "local current = redis.call('GET', KEYS[4])", "if current == ARGV[1] and redis.call('EXISTS', KEYS[5]) == 1 and (redis.call('GET', KEYS[13]) or \"1\") ~= ARGV[10] then return -3 end", - 'if not current and redis.call(\'EXISTS\', KEYS[5]) == 1 then', - ' local owner = redis.call(\'GET\', KEYS[6])', + "if not current and redis.call('EXISTS', KEYS[5]) == 1 then", + " local owner = redis.call('GET', KEYS[6])", ' if owner ~= ARGV[1] then return -3 end', 'end', 'if current then', ' if current ~= ARGV[1] then', - ' if redis.call(\'EXISTS\', KEYS[5]) == 1 then return -3 end', - ' redis.call(\'SET\', ARGV[4] .. current .. \':fenced\', \"1\")', + " if redis.call('EXISTS', KEYS[5]) == 1 then return -3 end", + " redis.call('SET', ARGV[4] .. current .. ':fenced', \"1\")", ' end', 'end', 'local registrationGeneration = tonumber(redis.call(\'GET\', KEYS[10]) or \"0\")', - 'local registrationGenerationIncarnation = redis.call(\'GET\', KEYS[11])', + "local registrationGenerationIncarnation = redis.call('GET', KEYS[11])", 'local registrationGenerationChanged = false', 'if registrationGeneration < 1 or registrationGenerationIncarnation ~= ARGV[1] then', - ' registrationGeneration = redis.call(\'INCR\', KEYS[10])', - ' redis.call(\'SET\', KEYS[11], ARGV[1])', + " registrationGeneration = redis.call('INCR', KEYS[10])", + " redis.call('SET', KEYS[11], ARGV[1])", ' registrationGenerationChanged = true', 'end', 'redis.call(\'SET\', KEYS[1], ARGV[2], \"EX\", ARGV[3])', @@ -495,7 +498,10 @@ export class RedisBridgeStore { script, 13, workerKey(registration.workerId), - incarnationFenceKey(registration.workerId, registration.incarnationId), + incarnationFenceKey( + registration.workerId, + registration.incarnationId, + ), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), lockKey(registration.workerId), @@ -517,7 +523,9 @@ export class RedisBridgeStore { authorizationObject?.identityId ?? '', expectedActiveCredentialId ?? '', registration.identityId ?? '', - registration.capabilities.requiresReadyConfirmation === true ? '1' : '0', + registration.capabilities.requiresReadyConfirmation === true + ? '1' + : '0', String(registration.capabilities.workspaceLeaseSlots ?? 1), ), this.redisCommandTimeoutMs, @@ -555,7 +563,9 @@ export class RedisBridgeStore { ); } if (!Number.isSafeInteger(result) || result < 1) { - throw new Error('Bridge worker registration returned an invalid generation'); + throw new Error( + 'Bridge worker registration returned an invalid generation', + ); } return result; } @@ -678,12 +688,17 @@ export class RedisBridgeStore { registration: RegisteredBridgeWorker, ) => Promise; }): Promise { - if (args.executionTimeoutMs !== undefined && ( - args.workspaceRequest == null || - !Number.isSafeInteger(args.executionTimeoutMs) || - args.executionTimeoutMs < 1 || args.executionTimeoutMs > 305_000 - )) { - throw new BridgeStoreError('ASSIGNMENT_INVALID', 'Invalid workspace execution budget'); + if ( + args.executionTimeoutMs !== undefined && + (args.workspaceRequest == null || + !Number.isSafeInteger(args.executionTimeoutMs) || + args.executionTimeoutMs < 1 || + args.executionTimeoutMs > 305_000) + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_INVALID', + 'Invalid workspace execution budget', + ); } this.assertDispatchActive(args.signal, args.deadlineAtMs); const dispatchable = await this.dispatchCommand( @@ -698,6 +713,15 @@ export class RedisBridgeStore { ); } let { registration, readyToken } = dispatchable; + if ( + (registration.capabilities.workspaceLeaseSlots ?? 1) > + this.maxWorkspaceLeaseSlots + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Worker slot negotiation exceeds this Code API replica ceiling; use consistent replica configuration', + ); + } if ( (args.requireTenantBinding === true && registration.binding == null) || (registration.binding != null && @@ -748,7 +772,9 @@ export class RedisBridgeStore { const assignmentId = randomBytes(18).toString('base64url'); const leaseToken = randomBytes(32).toString('base64url'); // The lock is acquired before admission finishes; it must outlive the later execution deadline. - const ttlSeconds = assignmentTtlSeconds(args.deadlineAtMs + (args.executionTimeoutMs ?? 0)); + const ttlSeconds = assignmentTtlSeconds( + args.deadlineAtMs + (args.executionTimeoutMs ?? 0), + ); const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; let workspaceLeaseSlot: number | undefined; @@ -758,14 +784,16 @@ export class RedisBridgeStore { ? new BridgeWorkspaceSlots(this.redis) : undefined; let resultCommitted = false; - const admission = args.workspaceRequest == null - ? undefined - : new BridgeAdmissionQueue(this.redis); + const admission = + args.workspaceRequest == null + ? undefined + : new BridgeAdmissionQueue(this.redis); try { if ( admission != null && !(await this.dispatchCommand( - () => admission.enter( + () => + admission.enter( args.workerId, assignmentId, args.deadlineAtMs, @@ -774,31 +802,39 @@ export class RedisBridgeStore { : args.workspaceRequest?.workspaceId, ), args, - 'Bridge admission enqueue', + 'Bridge admission enqueue', )) ) { - throw new BridgeStoreError('WORKER_QUEUE_FULL', 'Bridge worker pending request limit reached'); + throw new BridgeStoreError( + 'WORKER_QUEUE_FULL', + 'Bridge worker pending request limit reached', + ); } let locked = false; do { if ( admission != null && - workspaceSlots == null && !(await this.dispatchCommand( - () => admission.isHead(args.workerId, assignmentId), - args, - 'Bridge admission position', - )) + workspaceSlots == null && + !(await this.dispatchCommand( + () => admission.isHead(args.workerId, assignmentId), + args, + 'Bridge admission position', + )) ) { - await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); + await delay( + Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), + args.signal, + ); continue; } if (workspaceSlots != null) { workspaceLeaseSlot = await this.dispatchCommand( () => workspaceSlots.reserve({ - workerId: args.workerId, + workerId: args.workerId, incarnationId: lockIncarnationId, - assignmentId,workspaceId: args.workspaceRequest!.workspaceId, + assignmentId, + workspaceId: args.workspaceRequest!.workspaceId, capacity: registration.capabilities.workspaceLeaseSlots!, expiresAtMs: Date.now() + ttlSeconds * 1000, }), @@ -807,20 +843,23 @@ export class RedisBridgeStore { ); locked = workspaceLeaseSlot !== undefined; } else { - locked = await this.dispatchCommand( - () => - this.acquireLock( - args.workerId, - assignmentId, - lockIncarnationId, - ttlSeconds, - ), - args, - 'Bridge assignment lock acquisition', - ); + locked = await this.dispatchCommand( + () => + this.acquireLock( + args.workerId, + assignmentId, + lockIncarnationId, + ttlSeconds, + ), + args, + 'Bridge assignment lock acquisition', + ); } if (!locked && admission != null) { - await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); + await delay( + Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), + args.signal, + ); } } while (!locked && admission != null); if (!locked) { @@ -841,12 +880,21 @@ export class RedisBridgeStore { current == null || current.registration.incarnationId !== registration.incarnationId || current.registration.identityId !== registration.identityId || - current.registration.binding?.tenantId !== registration.binding?.tenantId + current.registration.binding?.tenantId !== + registration.binding?.tenantId ) { - throw new BridgeStoreError('WORKER_OFFLINE', 'Bridge worker changed while the request was waiting'); + throw new BridgeStoreError( + 'WORKER_OFFLINE', + 'Bridge worker changed while the request was waiting', + ); } - if (!supportsWorkspaceTool(current.registration, args.workspaceRequest!)) { - throw new BridgeStoreError('WORKER_MISMATCH', 'Bridge worker capabilities changed while the request was waiting'); + if ( + !supportsWorkspaceTool(current.registration, args.workspaceRequest!) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + 'Bridge worker capabilities changed while the request was waiting', + ); } } this.assertDispatchActive(args.signal, args.deadlineAtMs); @@ -870,7 +918,7 @@ export class RedisBridgeStore { ? {} : { workspaceLeaseSlot, - workspaceFence: `native-workspace:${ args.workspaceRequest!.workspaceId}`, + workspaceFence: `native-workspace:${args.workspaceRequest!.workspaceId}`, }), ...(registration.identityId != null ? { workerIdentityId: registration.identityId } @@ -932,7 +980,10 @@ export class RedisBridgeStore { } if ( args.workspaceRequest != null && - !supportsWorkspaceTool(replacement.registration, args.workspaceRequest) + !supportsWorkspaceTool( + replacement.registration, + args.workspaceRequest, + ) ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -1018,6 +1069,7 @@ export class RedisBridgeStore { if ( !Number.isSafeInteger(slot) || slot < 0 || + slot >= this.maxWorkspaceLeaseSlots || slot >= (registration?.capabilities.workspaceLeaseSlots ?? 1) || (registration?.capabilities.workspaceLeaseSlots ?? 1) <= 1 || registration?.incarnationId !== incarnationId @@ -1030,9 +1082,7 @@ export class RedisBridgeStore { } const deadline = Date.now() + waitMs; let firstPoll = true; - while ( - !signalAborted(signal) && - (firstPoll || Date.now() < deadline)) { + while (!signalAborted(signal) && (firstPoll || Date.now() < deadline)) { firstPoll = false; let assignmentId: string | null; try { @@ -1348,7 +1398,8 @@ export class RedisBridgeStore { ); if (existingSettlement === serializedSettlement && !quarantineWorkspace) return; - if (existingSettlement != null && + if ( + existingSettlement != null && existingSettlement !== serializedSettlement ) { throw new BridgeStoreError( @@ -1362,7 +1413,7 @@ export class RedisBridgeStore { 'Bridge settlement assignment read', ); if (assignment == null) { - if (existingSettlement === serializedSettlement) return; + if (existingSettlement === serializedSettlement) return; throw new BridgeStoreError( 'ASSIGNMENT_NOT_FOUND', 'Bridge assignment was not found', @@ -1378,7 +1429,7 @@ export class RedisBridgeStore { quarantineWorkspace && (assignment.workspaceFence == null || assignment.workspaceLeaseSlot === undefined || - settlement.status !== 'rejected') + settlement.status !== 'rejected') ) { throw new BridgeStoreError( 'ASSIGNMENT_INVALID', @@ -1438,7 +1489,7 @@ export class RedisBridgeStore { workerIncarnationKey(workerId), ); const script = [ - 'local existing = redis.call(\'GET\', KEYS[2])', + "local existing = redis.call('GET', KEYS[2])", 'if existing then', ' if existing == ARGV[1] then return 2 end', ' return -1', @@ -1448,12 +1499,12 @@ export class RedisBridgeStore { 'if ARGV[4] ~= "rejected" and redis.call(\'EXISTS\', KEYS[5]) == 0 then return -3 end', 'local stableIdentityKey = KEYS[#KEYS - 1]', 'if ARGV[5] ~= "" then', - ' if redis.call(\'GET\', stableIdentityKey) ~= ARGV[5] then return -4 end', - 'elseif redis.call(\'EXISTS\', stableIdentityKey) == 1 then return -4', + " if redis.call('GET', stableIdentityKey) ~= ARGV[5] then return -4 end", + "elseif redis.call('EXISTS', stableIdentityKey) == 1 then return -4", 'end', - 'if redis.call(\'GET\', KEYS[#KEYS]) ~= ARGV[7] then return -4 end', + "if redis.call('GET', KEYS[#KEYS]) ~= ARGV[7] then return -4 end", 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', - 'if redis.call(\'GET\', KEYS[3]) == ARGV[3] then redis.call(\'DEL\', KEYS[3], KEYS[4]) end', + "if redis.call('GET', KEYS[3]) == ARGV[3] then redis.call('DEL', KEYS[3], KEYS[4]) end", 'if ARGV[6] == "1" and ARGV[4] == "rejected" then', ' if ARGV[8] == "1" then redis.call(\'SET\', KEYS[6], "quarantined:" .. ARGV[3])', " else redis.call('DEL', KEYS[6]) end", @@ -1722,11 +1773,11 @@ export class RedisBridgeStore { // Keep acknowledged assignment metadata for late clean rejection recovery, // but atomically revoke fulfillment when no settlement has won yet. const closeScript = [ - 'local settlement = redis.call(\'GET\', KEYS[2])', + "local settlement = redis.call('GET', KEYS[2])", 'if settlement then return settlement end', - 'redis.call(\'DEL\', KEYS[3])', - 'if #KEYS == 4 and redis.call(\'GET\', KEYS[4]) == ARGV[1] then return nil end', - 'redis.call(\'DEL\', KEYS[1])', + "redis.call('DEL', KEYS[3])", + "if #KEYS == 4 and redis.call('GET', KEYS[4]) == ARGV[1] then return nil end", + "redis.call('DEL', KEYS[1])", 'return nil', ].join('\n'); const finalSettlement = await boundedCommand( @@ -1779,16 +1830,15 @@ export class RedisBridgeStore { const script = [ "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", 'if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[6]) ~= ARGV[7] then return 0 end', - 'if #KEYS == 7 and redis.call(\'EXISTS\', KEYS[7]) == 1 then return -1 end', + "if #KEYS == 7 and redis.call('EXISTS', KEYS[7]) == 1 then return -1 end", 'redis.call(\'SET\', KEYS[2], ARGV[2], \"EX\", ARGV[3])', - 'redis.call(\'RPUSH\', KEYS[3], ARGV[4])', - 'redis.call(\'EXPIRE\', KEYS[3], ARGV[3])', + "redis.call('RPUSH', KEYS[3], ARGV[4])", + "redis.call('EXPIRE', KEYS[3], ARGV[3])", ...(assignment.workspaceLeaseSlot === undefined - ? [ - 'redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])'] + ? ['redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])'] : []), 'redis.call(\'SET\', KEYS[5], "1", \"PXAT\", ARGV[6])', - 'if #KEYS == 7 then redis.call(\'SET\', KEYS[7], ARGV[4]) end', + "if #KEYS == 7 then redis.call('SET', KEYS[7], ARGV[4]) end", 'return 1', ].join('\n'); const keys = [ @@ -1885,8 +1935,8 @@ export class RedisBridgeStore { } const runtimeSessionId = assignmentWorkspace(assignment)!; const script = [ - 'if redis.call(\'GET\', KEYS[1]) == ARGV[1] then', - ' return redis.call(\'DEL\', KEYS[1])', + "if redis.call('GET', KEYS[1]) == ARGV[1] then", + " return redis.call('DEL', KEYS[1])", 'end', 'return 0', ].join('\n'); @@ -1895,10 +1945,7 @@ export class RedisBridgeStore { this.redis.eval( script, 1, - workspaceQuarantineKey( - assignment.workerId, - runtimeSessionId, - ), + workspaceQuarantineKey(assignment.workerId, runtimeSessionId), assignment.assignmentId, ), // Once settlement wins, caller cancellation must not prevent its @@ -1993,8 +2040,7 @@ export class RedisBridgeStore { if (cleanupResult !== -1) { await boundedCommand( assignment.workspaceLeaseSlot === undefined - ? - this.releaseLock(assignment.workerId, assignment.assignmentId) + ? this.releaseLock(assignment.workerId, assignment.assignmentId) : new BridgeWorkspaceSlots(this.redis).release( assignment.workerId, assignment.incarnationId, diff --git a/service/src/config.ts b/service/src/config.ts index 67a3561b..94daf5a3 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -81,12 +81,11 @@ export function checkpointPipelineBudgetMs( checkpointTimeoutMs, CHECKPOINT_METADATA_TIMEOUT_CAP_MS, ); - return ( launchTimeoutMs + return launchTimeoutMs + 2 * checkpointTimeoutMs + 2 * metadataTimeoutMs + POST_EXEC_CHECKPOINT_REGISTRY_COMMANDS - * RUNTIME_SESSION_REDIS_COMMAND_TIMEOUT_MS - ); + * RUNTIME_SESSION_REDIS_COMMAND_TIMEOUT_MS; } /** BullMQ's `timestamp` is the enqueue time. Anchor the worker deadline to it @@ -113,11 +112,10 @@ export function jobCompletionWaitTimeoutMs( backendCleanupTimeoutMs: number, egressRevokeTimeoutMs: number, ): number { - return ( jobTimeoutMs + return jobTimeoutMs + backendCleanupTimeoutMs + egressRevokeTimeoutMs - + WORKER_COMPLETION_OVERHEAD_MS - ); + + WORKER_COMPLETION_OVERHEAD_MS; } export function parseArnList(raw: string | undefined): string[] | undefined { @@ -240,10 +238,9 @@ export function resolveEgressGrantTtlSeconds(rawTtlSeconds: string | undefined, const lambdaMicrovmNumericConfig = resolveLambdaMicrovmNumericConfig(process.env); export function hostedAppOperationTimeoutMs(): number { - return ( env.CHECKPOINT_TIMEOUT_MS * 9 + return env.CHECKPOINT_TIMEOUT_MS * 9 + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS * 7 - + env.HOSTED_APP_START_TIMEOUT_MS + 30_000 - ); + + env.HOSTED_APP_START_TIMEOUT_MS + 30_000; } function configuredNumber(raw: string | undefined, fallback: number): number { From 4882def2606db6b7312f93b03ccf208276106a03 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 8 Sep 2026 23:04:09 -0400 Subject: [PATCH 4/7] fix: Fence BYOM Workspaces Independently of Settlement --- packages/code/README.md | 7 +- packages/code/src/cli-slots.test.ts | 82 +++- packages/code/src/cli.ts | 12 +- packages/code/src/worker.ts | 12 + service/src/bridge/concurrent-store.test.ts | 92 +++++ service/src/bridge/concurrent-worker.test.ts | 379 +++++++++++-------- service/src/bridge/store.ts | 189 +++++++-- 7 files changed, 571 insertions(+), 202 deletions(-) diff --git a/packages/code/README.md b/packages/code/README.md index 8ccfc5de..568698cf 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -541,6 +541,8 @@ The primary root keeps its configured workspace ID (default `primary`). Repeat `--workspace id=path` to add named roots, up to the protocol's 32-root limit. Roots must already exist and must not overlap or alias one another. Commands retain the selected root's sandbox boundary, not a shared parent-directory grant. +The `LIBRECHAT_CODE_WORKSPACE_QUARANTINE_FILE` single-file override is rejected +when multiple roots are configured; unset it to use separate root-derived markers. `LIBRECHAT_CODE_WORKSPACE_LEASE_SLOTS` is the equivalent worker setting. Both ceilings must be integers from 1 to 8; the lower ceiling wins. An older Code API @@ -559,7 +561,10 @@ enable concurrent Docker/NsJail sessions or bypass any approval/network policy. An uncertain mutation or executor failure leaves an assignment-owned local guard and a server-side fence for that root. Healthy roots can continue. The worker -does not replay the failed command. To recover a quarantined native root: +does not replay the failed command. A guard-cleanup failure after settlement fences +the root independently without replacing the committed result. Expiring ownership +receipts exclude command payloads; explicit reset invalidates old fence requests. +To recover a quarantined native root: 1. Stop the worker and inspect or restore the affected directory. 2. Run `librechat-code clear-workspace-quarantine --worker-dir /projects/second --workspace-id second` using the same deployment/identity configuration. diff --git a/packages/code/src/cli-slots.test.ts b/packages/code/src/cli-slots.test.ts index cc69ddfc..484f0dc0 100644 --- a/packages/code/src/cli-slots.test.ts +++ b/packages/code/src/cli-slots.test.ts @@ -1,7 +1,7 @@ import test from 'node:test'; import assert from 'node:assert/strict'; import { spawnSync } from 'node:child_process'; -import { mkdtemp, mkdir, rm } from 'node:fs/promises'; +import { mkdtemp, mkdir, rm, stat } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; @@ -61,3 +61,83 @@ test('CLI bounds requested workspace slots before connecting', () => { assert.notEqual(result.status, 0); assert.match(result.stderr, /cannot exceed 8/); }); + +test('CLI rejects one quarantine-file override for multiple roots', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'byom-cli-markers-')); + t.after(() => rm(root, { recursive: true, force: true })); + await mkdir(join(root, 'a')); + await mkdir(join(root, 'b')); + const result = spawnSync( + process.execPath, + [ + fileURLToPath(new URL('./cli.js', import.meta.url)), + 'run', + '--worker-dir', + join(root, 'a'), + '--workspace', + `second=${join(root, 'b')}`, + ], + { + encoding: 'utf8', + timeout: 3000, + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'http://127.0.0.1:1', + LIBRECHAT_CODE_WORKER_TOKEN: 'fixture', + LIBRECHAT_CODE_WORKER_ID: 'fixture-worker', + LIBRECHAT_CODE_COMMAND_SANDBOX: 'native-srt', + LIBRECHAT_CODE_WORKSPACE_QUARANTINE_FILE: join(root, 'shared.json'), + }, + }, + ); + assert.notEqual(result.status, 0); + assert.match(result.stderr, /single-root override/); +}); + +test('CLI preserves distinct case-sensitive roots on a non-Linux platform', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'byom-cli-case-')); + t.after(() => rm(root, { recursive: true, force: true })); + await mkdir(join(root, 'Foo')); + await mkdir(join(root, 'foo'), { recursive: true }); + if ( + (await stat(join(root, 'Foo'))).ino === (await stat(join(root, 'foo'))).ino + ) { + t.skip('requires a case-sensitive test filesystem'); + return; + } + const argv = [ + 'fixture', + 'run', + '--worker-dir', + join(root, 'Foo'), + '--workspace', + `second=${join(root, 'foo')}`, + '--workspace-lease-slots', + '2', + ]; + const result = spawnSync( + process.execPath, + [ + '--input-type=module', + '-e', + `Object.defineProperty(process, 'platform', {value:'darwin'}); process.argv=[process.execPath,...${JSON.stringify(argv)}]; await import(${JSON.stringify(new URL('./cli.js', import.meta.url).href)});`, + ], + { + encoding: 'utf8', + timeout: 3000, + env: { + ...process.env, + LIBRECHAT_CODE_URL: 'http://127.0.0.1:1', + LIBRECHAT_CODE_WORKER_TOKEN: 'fixture', + LIBRECHAT_CODE_WORKER_ID: 'fixture-worker', + LIBRECHAT_CODE_COMMAND_SANDBOX: 'native-srt', + LIBRECHAT_CODE_WORKSPACE_QUARANTINE_FILE: '', + }, + }, + ); + assert.match( + result.stderr, + /Concurrent workspace leases require native-srt commands/, + ); + assert.doesNotMatch(result.stderr, /overlap or alias/); +}); diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index f88784d9..59735f75 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -531,9 +531,7 @@ async function run( const rootIdentities = await Promise.all( roots.map((root) => stat(root.root)), ); - const normalized = roots.map((root) => - process.platform === 'linux' ? root.root : root.root.toLowerCase(), - ); + const normalized = roots.map((root) => root.root); for (let i = 0; i < roots.length; i++) for (let j = 0; j < i; j++) { const inside = (a: string, b: string): boolean => { @@ -560,6 +558,14 @@ async function run( ) { throw new Error('Concurrent workspace leases require native-srt commands'); } + if ( + roots.length > 1 && + process.env.LIBRECHAT_CODE_WORKSPACE_QUARANTINE_FILE?.trim() + ) { + throw new Error( + 'LIBRECHAT_CODE_WORKSPACE_QUARANTINE_FILE is a single-root override; unset it for multiple workspace roots', + ); + } const rootQuarantinePaths = new Map( roots.map((root) => [ root.id, diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index e37c3b88..92cf8747 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -350,6 +350,7 @@ export class BridgeWorker { private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; private lastRegisteredAtMs = 0; private mutationGuardArmed = false; + private readonly quarantinedWorkspaces = new Set(); private readonly activeWorkspaceAssignments = new Map< string, { id: string; done: Promise } @@ -673,6 +674,7 @@ export class BridgeWorker { DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, signal, ); + this.quarantinedWorkspaces.delete(workspaceId); } async lease( @@ -1136,6 +1138,13 @@ export class BridgeWorker { }; try { await this.executeOwned(adjusted, signal); + } catch (error) { + if (root != null && error instanceof BridgeWorkspaceQuarantinedError) { + // A failed unlink/fsync may have removed the durable marker already. + // Fence locally before releasing the handoff to the next root assignment. + this.quarantinedWorkspaces.add(root); + } + throw error; } finally { if (root != null) { this.activeWorkspaceAssignments.delete(root); @@ -1273,6 +1282,9 @@ export class BridgeWorker { } const workspaceRequest = assignment.request; try { + if (this.quarantinedWorkspaces.has(workspaceRequest.workspaceId)) { + throw new Error('Workspace requires an explicit quarantine reset'); + } if (this.options.workspaceQuarantines != null) await guard?.assertAvailable(); } catch (error) { diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts index 177b4f74..8bc53f50 100644 --- a/service/src/bridge/concurrent-store.test.ts +++ b/service/src/bridge/concurrent-store.test.ts @@ -271,3 +271,95 @@ test('late quarantine releases its slot after caller cancellation and retains on code: 'WORKSPACE_QUARANTINED', }); }); + +test('post-settlement fences are authenticated, idempotent, and invalidated by reset', async () => { + await register(); + const pending = dispatch('a'); + const assignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await settle(assignment); + await pending; // Normal dispatch cleanup has removed the full assignment. + const receiptKey = `codeapi:bridge:v1:assignment:${assignment.assignmentId}:workspace-fence-owner`; + expect(await redis.ttl(receiptKey)).toBeGreaterThan(0); + expect( + JSON.parse((await redis.hget(receiptKey, 'metadata'))!), + ).not.toHaveProperty('request'); + const resultKey = `codeapi:bridge:v1:assignment:${assignment.assignmentId}:settlement`; + const originalResult = await redis.get(resultKey); + const intent = { + protocolVersion: 1 as const, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected' as const, + error: 'local cleanup failed after commit', + }; + await expect( + store.settle( + workerId, + assignment.assignmentId, + { ...intent, leaseToken: 'forged' }, + undefined, + undefined, + true, + ), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + await expect( + store.settle( + workerId, + assignment.assignmentId, + intent, + undefined, + 'different-principal', + true, + ), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + await store.settle( + workerId, + assignment.assignmentId, + intent, + undefined, + undefined, + true, + ); + await store.settle( + workerId, + assignment.assignmentId, + intent, + undefined, + undefined, + true, + ); + expect(await redis.get(resultKey)).toBe(originalResult); + await expect(dispatch('a')).rejects.toMatchObject({ + code: 'WORKSPACE_QUARANTINED', + }); + await store.resetWorkspace(workerId, incarnationId, 'native-workspace:a'); + await expect( + store.settle( + workerId, + assignment.assignmentId, + intent, + undefined, + undefined, + true, + ), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + const next = dispatch('a'); + const nextAssignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await settle(nextAssignment); + await next; +}); diff --git a/service/src/bridge/concurrent-worker.test.ts b/service/src/bridge/concurrent-worker.test.ts index 3a359f07..8d9e7962 100644 --- a/service/src/bridge/concurrent-worker.test.ts +++ b/service/src/bridge/concurrent-worker.test.ts @@ -7,195 +7,240 @@ import { WorkspaceToolError } from '../../../packages/code/src/workspace'; import type { WorkspaceMutationQuarantine } from '../../../packages/code/src/worker'; import type { BridgeWorkspaceToolCapabilities } from '../../../packages/code/src/protocol'; -test('concurrent worker quarantines one root while another settles', async () => { - const redis = new RedisMock() as unknown as Redis; - const store = new RedisBridgeStore(redis, 60, 1000, 2); - const controller = new AbortController(); - const workerId = 'worker-concurrency'; - const incarnationId = 'incarnation-concurrency'; - const guards = new Map(); - const pending = new Set(); - for (const root of ['a', 'b']) - guards.set(root, { - async assertAvailable() { - if (pending.has(root)) throw new Error('quarantined'); - }, - async arm() { - pending.add(root); - }, - async clear() { - pending.delete(root); +for (const failure of ['execution', 'cleanup', 'post-unlink']) { + const cleanupFailure = failure !== 'execution'; + test(`concurrent worker isolates ${failure} failure`, async () => { + const redis = new RedisMock() as unknown as Redis; + const store = new RedisBridgeStore(redis, 60, 1000, 2); + const controller = new AbortController(); + const workerId = 'worker-concurrency'; + const incarnationId = 'incarnation-concurrency'; + const guards = new Map(); + const pending = new Set(); + for (const root of ['a', 'b']) + guards.set(root, { + async assertAvailable() { + if (pending.has(root)) throw new Error('quarantined'); + }, + async arm() { + pending.add(root); + }, + async clear() { + if (failure === 'post-unlink') pending.delete(root); + if (cleanupFailure && root === 'a') + throw new Error('injected guard cleanup failure'); + pending.delete(root); + }, + async quarantine() { + expect(pending.has(root)).toBe(true); + }, + }); + const capabilities: BridgeWorkspaceToolCapabilities = { + protocolVersion: 1, + operations: ['execute_command'], + workspaces: [{ id: 'a' }, { id: 'b' }], + }; + let startBoth!: () => void; + const bothStarted = new Promise((resolve) => { + startBoth = resolve; + }); + const started = new Set(); + const errors: unknown[] = []; + let registered!: () => void; + const ready = new Promise((resolve) => { + registered = resolve; + }); + const worker = new BridgeWorker({ + codeApiUrl: 'http://fixture.invalid', + token: 'fixture', + workerId, + incarnationId, + sandboxEndpoint: 'http://sandbox.invalid', + leaseWaitMs: 50, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + requiresReadyConfirmation: true, + workspaceLeaseSlots: 2, + workspaceTools: capabilities, }, - async quarantine() { - expect(pending.has(root)).toBe(true); + workspaceQuarantines: guards, + onError: (error) => { + errors.push(error); }, - }); - const capabilities: BridgeWorkspaceToolCapabilities = { - protocolVersion: 1, - operations: ['execute_command'], - workspaces: [{ id: 'a' }, { id: 'b' }], - }; - let startBoth!: () => void; - const bothStarted = new Promise((resolve) => { - startBoth = resolve; - }); - const started = new Set(); - const errors: unknown[] = []; - let registered!: () => void; - const ready = new Promise((resolve) => { - registered = resolve; - }); - const worker = new BridgeWorker({ - codeApiUrl: 'http://fixture.invalid', - token: 'fixture', - workerId, - incarnationId, - sandboxEndpoint: 'http://sandbox.invalid', - leaseWaitMs: 50, - capabilities: { - statefulWorkspace: false, - sandboxProfile: 'native-srt', - runtimes: [], - requiresReadyConfirmation: true, - workspaceLeaseSlots: 2, - workspaceTools: capabilities, - }, - workspaceQuarantines: guards, - onError: (error) => { - errors.push(error); - }, - workspaceTools: { - capabilities, - mutationFailuresAreAtomic: true, - async execute(request) { - started.add(request.workspaceId); - if (started.size === 2) startBoth(); - await bothStarted; - if (request.workspaceId === 'a') - throw new WorkspaceToolError( - 'uncertain command', - 'COMMAND_UNAVAILABLE', - true, - ); - return { - protocolVersion: 1, - operation: 'execute_command', - workspaceId: 'b', - stdout: 'completed', - stderr: '', - exitCode: 0, - truncated: false, - timedOut: false, - }; + workspaceTools: { + capabilities, + mutationFailuresAreAtomic: true, + async execute(request) { + started.add(request.workspaceId); + if (started.size === 2) startBoth(); + await bothStarted; + if (request.workspaceId === 'a' && !cleanupFailure) + throw new WorkspaceToolError( + 'uncertain command', + 'COMMAND_UNAVAILABLE', + true, + ); + return { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: request.workspaceId, + stdout: 'completed', + stderr: '', + exitCode: 0, + truncated: false, + timedOut: false, + }; + }, }, - }, - fetchImpl: (async (url, init) => { - const path = new URL(String(url)).pathname; - const body = JSON.parse(String(init?.body)); - const signal = init?.signal ?? undefined; - let result: object; - if (path.endsWith('/register')) { - const generation = await store.register(body); - result = { - protocolVersion: 1, - workerId, - incarnationId, - registrationGeneration: generation, - registeredAt: new Date().toISOString(), - leaseTtlMs: 60000, - workspaceLeaseSlots: 2, - supportedWorkspaceToolOperations: ['execute_command'], - }; - } else if (path.endsWith('/ready')) { - await store.confirmReady( - workerId, - incarnationId, - body.registrationGeneration, - ); - registered(); - result = { protocolVersion: 1, ready: true }; - } else if (path.endsWith('/lease')) { - result = { - protocolVersion: 1, - serverElapsedMs: 0, - assignment: await store.lease( + fetchImpl: (async (url, init) => { + const path = new URL(String(url)).pathname; + const body = JSON.parse(String(init?.body)); + const signal = init?.signal ?? undefined; + let result: object; + if (path.endsWith('/register')) { + const generation = await store.register(body); + result = { + protocolVersion: 1, workerId, incarnationId, - body.waitMs, - signal, - undefined, - body.workspaceLeaseSlot, - ), - }; - } else { - const id = path.split('/').at(-2)!; - if (path.endsWith('/ack')) { - await store.acknowledgeLease( + registrationGeneration: generation, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60000, + workspaceLeaseSlots: 2, + supportedWorkspaceToolOperations: ['execute_command'], + }; + } else if (path.endsWith('/ready')) { + await store.confirmReady( workerId, incarnationId, - id, - body.generation, - body.leaseToken, - signal, + body.registrationGeneration, ); - result = { protocolVersion: 1, accepted: true }; - } else if (path.endsWith('/cancellation')) { + registered(); + result = { protocolVersion: 1, ready: true }; + } else if (path.endsWith('/lease')) { result = { protocolVersion: 1, - cancelled: await store.cancelled( + serverElapsedMs: 0, + assignment: await store.lease( workerId, incarnationId, - id, + body.waitMs, signal, + undefined, + body.workspaceLeaseSlot, ), }; } else { - await store.settle( - workerId, - id, - body, - signal, - undefined, - path.endsWith('/quarantine'), - ); - result = { protocolVersion: 1, accepted: true }; + const id = path.split('/').at(-2)!; + if (path.endsWith('/ack')) { + await store.acknowledgeLease( + workerId, + incarnationId, + id, + body.generation, + body.leaseToken, + signal, + ); + result = { protocolVersion: 1, accepted: true }; + } else if (path.endsWith('/cancellation')) { + result = { + protocolVersion: 1, + cancelled: await store.cancelled( + workerId, + incarnationId, + id, + signal, + ), + }; + } else { + await store.settle( + workerId, + id, + body, + signal, + undefined, + path.endsWith('/quarantine'), + ); + result = { protocolVersion: 1, accepted: true }; + } } - } - return Response.json(result); - }) as typeof fetch, - }); - const running = worker.run(controller.signal); - void running.catch(() => undefined); - try { - await ready; - const results = await Promise.all( - ['a', 'b'].map((workspaceId) => + return Response.json(result); + }) as typeof fetch, + }); + const running = worker.run(controller.signal); + void running.catch(() => undefined); + try { + await ready; + const results = await Promise.allSettled( + ['a', 'b'].map((workspaceId) => + store.dispatchWorkspaceTool({ + workerId, + signal: controller.signal, + deadlineAtMs: Date.now() + 3000, + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId, + command: 'fixture', + }, + }), + ), + ); + if (!cleanupFailure) + expect(results[0]).toMatchObject({ + status: 'fulfilled', + value: { status: 'rejected' }, + }); + else if (results[0].status === 'fulfilled') + expect(results[0].value).toMatchObject({ status: 'fulfilled', result: { stdout: 'completed' } }); + else expect(results[0].reason).toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + expect(results[1]).toMatchObject({ + status: 'fulfilled', + value: { status: 'fulfilled' }, + }); + for (let i = 0; i < 100 && errors.length === 0; i++) + await new Promise((resolve) => setTimeout(resolve, 5)); + expect(errors.length).toBe(1); + await expect( store.dispatchWorkspaceTool({ workerId, signal: controller.signal, - deadlineAtMs: Date.now() + 3000, + deadlineAtMs: Date.now() + 1000, request: { protocolVersion: 1, operation: 'execute_command', - workspaceId, - command: 'fixture', + workspaceId: 'a', + command: 'must not execute', }, }), - ), - ); - expect(results[0]).toMatchObject({ status: 'rejected' }); - expect(results[1]).toMatchObject({ status: 'fulfilled' }); - expect([...pending]).toEqual(['a']); - expect(started.size).toBe(2); - } catch (error) { - throw new AggregateError( - [error, ...errors], - `Started roots: ${[...started].join(',')}`, - ); - } finally { - controller.abort(); - await running; - await redis.flushall(); - redis.disconnect(); - } -}); + ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + await expect( + store.dispatchWorkspaceTool({ + workerId, + signal: controller.signal, + deadlineAtMs: Date.now() + 1000, + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: 'b', + command: 'still healthy', + }, + }), + ).resolves.toMatchObject({ status: 'fulfilled' }); + expect([...pending]).toEqual(failure === 'post-unlink' ? [] : ['a']); + expect(started.size).toBe(2); + } catch (error) { + throw new AggregateError( + [error, ...errors], + `Started roots: ${[...started].join(',')}`, + ); + } finally { + controller.abort(); + await running; + await redis.flushall(); + redis.disconnect(); + } + }); +} diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index c03048a1..c15cf425 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -68,7 +68,27 @@ interface StoredAssignment extends CodeBridgeAssignment { workspaceFence?: string; } -function assignmentWorkspace(assignment: StoredAssignment): string | undefined { +type AssignmentOwnership = Pick< + StoredAssignment, + | 'assignmentId' + | 'workerId' + | 'incarnationId' + | 'workspaceFence' + | 'workspaceLeaseSlot' + | 'generation' + | 'leaseTokenHash' + | 'workerIdentityId' + | 'expiresAt' + | 'runtimeSessionId' +>; + +function workspaceFenceReceiptKey(assignmentId: string): string { + return `${assignmentKey(assignmentId)}:workspace-fence-owner`; +} + +function assignmentWorkspace( + assignment: AssignmentOwnership, +): string | undefined { return assignment.workspaceFence ?? assignment.runtimeSessionId; } @@ -1390,14 +1410,23 @@ export class RedisBridgeStore { identityId?: string, quarantineWorkspace = false, ): Promise { + if (quarantineWorkspace) { + await this.quarantineSettledWorkspace( + workerId, + assignmentId, + settlement, + signal, + identityId, + ); + return; + } const serializedSettlement = JSON.stringify(settlement); const existingSettlement = await this.leaseCommand( this.redis.get(settlementKey(assignmentId)), signal, 'Bridge settlement existing read', ); - if (existingSettlement === serializedSettlement && !quarantineWorkspace) - return; + if (existingSettlement === serializedSettlement) return; if ( existingSettlement != null && existingSettlement !== serializedSettlement @@ -1425,17 +1454,6 @@ export class RedisBridgeStore { 'Bridge assignment belongs to another worker', ); } - if ( - quarantineWorkspace && - (assignment.workspaceFence == null || - assignment.workspaceLeaseSlot === undefined || - settlement.status !== 'rejected') - ) { - throw new BridgeStoreError( - 'ASSIGNMENT_INVALID', - 'Quarantine requires a concurrent workspace assignment', - ); - } const registration = await this.leaseCommand( this.registration(workerId), signal, @@ -1506,8 +1524,7 @@ export class RedisBridgeStore { 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', "if redis.call('GET', KEYS[3]) == ARGV[3] then redis.call('DEL', KEYS[3], KEYS[4]) end", 'if ARGV[6] == "1" and ARGV[4] == "rejected" then', - ' if ARGV[8] == "1" then redis.call(\'SET\', KEYS[6], "quarantined:" .. ARGV[3])', - " else redis.call('DEL', KEYS[6]) end", + " redis.call('DEL', KEYS[6])", 'end', 'return 1', ].join('\n'); @@ -1524,7 +1541,6 @@ export class RedisBridgeStore { identityId ?? '', hasWorkspace ? '1' : '0', settlement.incarnationId, - quarantineWorkspace ? '1' : '0', ), signal, 'Bridge settlement commit', @@ -1637,6 +1653,97 @@ export class RedisBridgeStore { ); } + private async quarantineSettledWorkspace( + workerId: string, + assignmentId: string, + settlement: AnyCodeBridgeSettlement, + signal?: AbortSignal, + identityId?: string, + ): Promise { + // Keep a small, expiring ownership receipt separate from assignment cleanup. + // A local guard can fail to clear after the result has already committed. + const receiptKey = workspaceFenceReceiptKey(assignmentId); + const raw = await this.leaseCommand( + this.redis.hgetall(receiptKey), + signal, + 'Workspace fence ownership read', + ); + const receipt = + raw.metadata == null + ? undefined + : (JSON.parse(raw.metadata) as AssignmentOwnership); + if ( + receipt == null || + receipt.workerId !== workerId || + receipt.assignmentId !== assignmentId || + receipt.workspaceFence == null || + receipt.workspaceLeaseSlot === undefined || + settlement.status !== 'rejected' || + receipt.incarnationId !== settlement.incarnationId || + receipt.generation !== settlement.generation || + receipt.leaseTokenHash !== tokenHash(settlement.leaseToken) || + receipt.workerIdentityId !== identityId + ) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Workspace quarantine ownership is stale', + ); + } + const fence = workspaceQuarantineKey(workerId, receipt.workspaceFence); + const accepted = Number( + await this.leaseCommand( + this.redis.eval( + [ + "if redis.call('HGET', KEYS[1], 'metadata') ~= ARGV[1] then return 0 end", + "if redis.call('HGET', KEYS[1], 'epoch') ~= ARGV[4] then return 0 end", + "if redis.call('GET', KEYS[2]) ~= ARGV[2] then return 0 end", + "if (redis.call('GET', KEYS[3]) or '') ~= ARGV[3] then return 0 end", + "if (redis.call('GET', KEYS[4]) or '0') ~= ARGV[4] then return 0 end", + "redis.call('SET', KEYS[5], 'quarantined:' .. ARGV[5])", + // Never replace a committed result. Before settlement, terminate the waiter. + "redis.call('SET', KEYS[6], ARGV[6], 'EX', ARGV[7], 'NX')", + "if redis.call('GET', KEYS[7]) == ARGV[5] then redis.call('DEL', KEYS[7]) end", + "if redis.call('GET', KEYS[8]) == ARGV[5] then redis.call('DEL', KEYS[8]) end", + 'return 1', + ].join('\n'), + 8, + receiptKey, + workerIncarnationKey(workerId), + workerStableIdentityKey(workerId), + `${fence}:epoch`, + fence, + settlementKey(assignmentId), + leaseClaimKey( + workerId, + receipt.incarnationId, + receipt.workspaceLeaseSlot, + ), + leaseAckKey( + workerId, + receipt.incarnationId, + receipt.workspaceLeaseSlot, + ), + raw.metadata, + receipt.incarnationId, + identityId ?? '', + raw.epoch, + assignmentId, + JSON.stringify(settlement), + assignmentTtlSeconds(Date.parse(receipt.expiresAt)), + ), + signal, + 'Workspace quarantine fence commit', + ), + ); + if (accepted !== 1) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Workspace quarantine ownership changed or was reset', + ); + } + await this.cleanupWithRetry(workerId, assignmentId, receipt); + } + async resetWorkspace( workerId: string, incarnationId: string, @@ -1650,12 +1757,14 @@ export class RedisBridgeStore { "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return -1 end", "if redis.call('EXISTS', KEYS[2]) == 1 then return -2 end", "redis.call('DEL', KEYS[3])", + "if redis.call('EXISTS', KEYS[4]) == 1 then redis.call('INCR', KEYS[4]) end", 'return 1', ].join('\n'), - 3, + 4, workerIncarnationKey(workerId), lockKey(workerId), workspaceQuarantineKey(workerId, runtimeSessionId), + `${workspaceQuarantineKey(workerId, runtimeSessionId)}:epoch`, incarnationId, ), signal, @@ -1804,19 +1913,14 @@ export class RedisBridgeStore { private async cancel( assignmentId: string, - assignment?: StoredAssignment, + assignment?: AssignmentOwnership, ): Promise { const ttlSeconds = assignment == null ? 30 : assignmentTtlSeconds(Date.parse(assignment.expiresAt)); await boundedCommand( - this.redis.set( - cancellationKey(assignmentId), - '1', - 'EX', - ttlSeconds, - ), + this.redis.set(cancellationKey(assignmentId), '1', 'EX', ttlSeconds), this.redisCommandTimeoutMs, 'Bridge assignment cancellation', ); @@ -1830,7 +1934,7 @@ export class RedisBridgeStore { const script = [ "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", 'if ARGV[7] ~= "" and redis.call(\'GET\', KEYS[6]) ~= ARGV[7] then return 0 end', - "if #KEYS == 7 and redis.call('EXISTS', KEYS[7]) == 1 then return -1 end", + "if #KEYS >= 7 and redis.call('EXISTS', KEYS[7]) == 1 then return -1 end", 'redis.call(\'SET\', KEYS[2], ARGV[2], \"EX\", ARGV[3])', "redis.call('RPUSH', KEYS[3], ARGV[4])", "redis.call('EXPIRE', KEYS[3], ARGV[3])", @@ -1838,7 +1942,14 @@ export class RedisBridgeStore { ? ['redis.call(\'SET\', KEYS[4], ARGV[1], \"PX\", ARGV[5])'] : []), 'redis.call(\'SET\', KEYS[5], "1", \"PXAT\", ARGV[6])', - "if #KEYS == 7 then redis.call('SET', KEYS[7], ARGV[4]) end", + "if #KEYS >= 7 then redis.call('SET', KEYS[7], ARGV[4]) end", + 'if #KEYS == 9 then', + " local epoch = redis.call('GET', KEYS[9])", + " if type(epoch) ~= 'string' then epoch = '0'; redis.call('SET', KEYS[9], epoch, 'EX', ARGV[3]) end", + " if redis.call('PTTL', KEYS[9]) < tonumber(ARGV[3]) * 1000 then redis.call('EXPIRE', KEYS[9], ARGV[3]) end", + " redis.call('HSET', KEYS[8], 'metadata', ARGV[8], 'epoch', epoch)", + " redis.call('EXPIRE', KEYS[8], ARGV[3])", + 'end', 'return 1', ].join('\n'); const keys = [ @@ -1861,6 +1972,23 @@ export class RedisBridgeStore { ), ); } + const receipt: AssignmentOwnership = { + assignmentId: assignment.assignmentId, + workerId: assignment.workerId, + incarnationId: assignment.incarnationId, + workspaceFence: assignment.workspaceFence, + workspaceLeaseSlot: assignment.workspaceLeaseSlot, + generation: assignment.generation, + leaseTokenHash: assignment.leaseTokenHash, + workerIdentityId: assignment.workerIdentityId, + expiresAt: assignment.expiresAt, + }; + if (assignment.workspaceLeaseSlot !== undefined) { + keys.push( + workspaceFenceReceiptKey(assignment.assignmentId), + `${workspaceQuarantineKey(assignment.workerId, assignment.workspaceFence!)}:epoch`, + ); + } const result = await this.redis.eval( script, keys.length, @@ -1872,6 +2000,7 @@ export class RedisBridgeStore { String(ttlSeconds * 1000), String(Date.parse(assignment.expiresAt)), readyToken ?? '', + JSON.stringify(receipt), ); if (Number(result) === -1) { throw new BridgeStoreError( @@ -1909,7 +2038,7 @@ export class RedisBridgeStore { private async cleanupDispatch( workerId: string, assignmentId: string, - assignment: StoredAssignment | undefined, + assignment: AssignmentOwnership | undefined, ): Promise { await Promise.all([ this.cancel(assignmentId, assignment), @@ -1965,7 +2094,7 @@ export class RedisBridgeStore { private async cleanupWithRetry( workerId: string, assignmentId: string, - assignment: StoredAssignment | undefined, + assignment: AssignmentOwnership | undefined, ): Promise { let lastError: unknown; for (let attempt = 0; attempt < 3; attempt += 1) { @@ -1980,7 +2109,7 @@ export class RedisBridgeStore { throw lastError; } - private async cleanup(assignment: StoredAssignment): Promise { + private async cleanup(assignment: AssignmentOwnership): Promise { const keys = [ assignmentKey(assignment.assignmentId), queueKey( From 760ab7f3ff152b4512e69a1297834e617e4812c6 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 8 Sep 2026 23:50:23 -0400 Subject: [PATCH 5/7] fix: Retain Workspace Ownership Through Cleanup --- packages/code/README.md | 6 + packages/code/src/cli.ts | 2 +- packages/code/src/worker-slots.test.ts | 35 +++++ packages/code/src/worker.ts | 149 ++++++++++++++++--- service/src/bridge/concurrent-store.test.ts | 109 +++++++++++++- service/src/bridge/concurrent-worker.test.ts | 65 +++++++- service/src/bridge/router.ts | 7 +- service/src/bridge/store.ts | 68 ++++++++- 8 files changed, 408 insertions(+), 33 deletions(-) diff --git a/packages/code/README.md b/packages/code/README.md index 568698cf..3573a696 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -564,6 +564,12 @@ and a server-side fence for that root. Healthy roots can continue. The worker does not replay the failed command. A guard-cleanup failure after settlement fences the root independently without replacing the committed result. Expiring ownership receipts exclude command payloads; explicit reset invalidates old fence requests. +The server releases a root only after result finalization **and** explicit local +cleanup confirmation. Local guard cleanup has a five-second bound; an expired +receipt never implies a clean root. Control receipt delivery retries three times. +If delivery remains unavailable, that lease lane stops while healthy lanes keep +running; inspect/reset the affected root and restart the worker to restore capacity. +Reset-only registration stays unready and cannot attract new assignments. To recover a quarantined native root: 1. Stop the worker and inspect or restore the affected directory. diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 59735f75..27a8d2d9 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -964,7 +964,7 @@ async function run( const resetNativeRoot = option(args, '--reset-workspace-quarantine'); if (resetNativeRoot != null) { await worker.refreshCredential(controller.signal); - await worker.register(controller.signal); + await worker.registerForMaintenance(controller.signal); await worker.resetNativeWorkspace(resetNativeRoot, controller.signal); process.stdout.write( `librechat-code: reset acknowledged for native workspace ${resetNativeRoot}\n`, diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts index 4ddb6423..ba6f37eb 100644 --- a/packages/code/src/worker-slots.test.ts +++ b/packages/code/src/worker-slots.test.ts @@ -11,6 +11,41 @@ const capabilities: BridgeWorkspaceToolCapabilities = { operations: ['read_file'], workspaces: [{ id: 'a' }, { id: 'b' }], }; +test('maintenance registration never advertises readiness or starts leasing', async () => { + const paths: string[] = []; + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + incarnationId: 'incarnation-maintenance', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + }, + fetchImpl: async (url, init) => { + const path = new URL(String(url)).pathname; + paths.push(path); + assert.equal( + JSON.parse(String(init?.body)).capabilities.requiresReadyConfirmation, + true, + ); + return Response.json({ + protocolVersion: 1, + workerId: 'worker', + incarnationId: 'incarnation-maintenance', + registrationGeneration: 1, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60000, + }); + }, + }); + await worker.registerForMaintenance(); + assert.equal(paths.length, 1); + assert.ok(paths[0].endsWith('/register')); + await assert.rejects(worker.run(), /maintenance/i); +}); for (const receipt of [undefined, 1]) { test(`worker keeps serial lease wire format for receipt ${receipt}`, async () => { const controller = new AbortController(); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 92cf8747..be2badbf 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -41,6 +41,7 @@ export interface BridgeWorkerOptions { leaseTransportGraceMs?: number; registrationTransportTimeoutMs?: number; leaseAckTransportTimeoutMs?: number; + workspaceCleanupTimeoutMs?: number; resetTransportTimeoutMs?: number; cancellationPollIntervalMs?: number; cancellationTransportTimeoutMs?: number; @@ -349,6 +350,7 @@ export class BridgeWorker { private activeCapabilities: BridgeWorkerCapabilities; private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; private lastRegisteredAtMs = 0; + private maintenanceOnly = false; private mutationGuardArmed = false; private readonly quarantinedWorkspaces = new Set(); private readonly activeWorkspaceAssignments = new Map< @@ -453,6 +455,22 @@ export class BridgeWorker { return await this.registerWithPolicy(signal, false); } + async registerForMaintenance( + signal?: AbortSignal, + ): Promise { + if ( + this.lastRegisteredAtMs !== 0 || + this.registrationInFlight != null || + this.concurrentRunning + ) { + throw new BridgeProtocolError( + 'Maintenance registration requires a fresh worker', + ); + } + this.maintenanceOnly = true; + return this.register(signal); + } + private async registerWithPolicy( signal: AbortSignal | undefined, allowActiveMutation: boolean, @@ -514,7 +532,9 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: this.options.workerId, incarnationId: this.incarnationId, - capabilities, + capabilities: this.maintenanceOnly + ? { ...capabilities, requiresReadyConfirmation: true } + : capabilities, }, registrationController.signal, ); @@ -582,7 +602,10 @@ export class BridgeWorker { this.registrationTtlMs = registration.leaseTtlMs; this.activeCapabilities = this.registrationCapabilities; await this.options.onRegistered?.(registration); - if (this.options.capabilities.requiresReadyConfirmation === true) { + if ( + !this.maintenanceOnly && + this.options.capabilities.requiresReadyConfirmation === true + ) { await this.confirmReady(registration, signal); } this.lastRegisteredAtMs = registrationStartedAtMs; @@ -820,6 +843,10 @@ export class BridgeWorker { } async run(signal?: AbortSignal): Promise { + if (this.maintenanceOnly) + throw new BridgeProtocolError( + 'Maintenance workers cannot execute assignments', + ); let reconnectAttempt = 0; while (!signal?.aborted) { try { @@ -889,24 +916,24 @@ export class BridgeWorker { // The durable local guard is retained. A distinct receipt tells // Code API to release this slot without declaring the root clean. try { - await this.timedRequest( - this.assignmentUrl(assignment, 'quarantine'), - { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - incarnationId: this.incarnationId, - generation: assignment.generation, - leaseToken: assignment.leaseToken, - status: 'rejected', - error: - 'Workspace quarantined after an uncertain execution or settlement; inspect it before resetting.', - }, - this.options.leaseAckTransportTimeoutMs ?? - DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + await this.reportWorkspaceOwnership( + assignment, + 'quarantine', + controller.signal, ); this.options.onError?.(error); continue; } catch (quarantineError) { - fail(quarantineError); + // The durable server fence remains until explicit cleanup/reset. + // An unavailable control receipt must not cancel healthy roots. + if ( + quarantineError instanceof BridgeProtocolError && + (quarantineError.status === 401 || + quarantineError.status === 403 || + quarantineError.code === 'WORKER_FENCED') + ) + fail(quarantineError); + else this.options.onError?.(quarantineError); return; } } @@ -1597,7 +1624,31 @@ export class BridgeWorker { } if (workspaceMutationArmed) { try { - await guard!.clear(assignment.assignmentId); + if (assignment.workspaceLeaseSlot === undefined) { + await guard!.clear(assignment.assignmentId); + } else { + let timer!: ReturnType; + try { + await Promise.race([ + guard!.clear(assignment.assignmentId), + new Promise((_resolve, reject) => { + timer = setTimeout( + () => + reject(new Error('Workspace guard cleanup timed out')), + Math.min( + 5000, + Math.max( + 1, + this.options.workspaceCleanupTimeoutMs ?? 5000, + ), + ), + ); + }), + ]); + } finally { + clearTimeout(timer); + } + } if ( assignment.executionKind === 'workspace_tool' && isWorkspaceToolRequest(assignment.request) @@ -1612,6 +1663,20 @@ export class BridgeWorker { ); } } + if (assignment.workspaceLeaseSlot !== undefined) { + try { + await this.reportWorkspaceOwnership( + assignment, + 'workspace-cleanup', + signal, + ); + } catch (error) { + throw new BridgeWorkspaceQuarantinedError( + 'Workspace cleanup acknowledgement could not be confirmed', + error, + ); + } + } } finally { heartbeatController.abort(); try { @@ -1623,6 +1688,53 @@ export class BridgeWorker { } } + private async reportWorkspaceOwnership( + assignment: BridgeAssignment, + operation: 'quarantine' | 'workspace-cleanup', + signal?: AbortSignal, + ): Promise { + let failure: unknown; + for (let attempt = 0; attempt < 3 && !signal?.aborted; attempt++) { + try { + await this.timedRequest( + this.assignmentUrl(assignment, operation), + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected', + error: + operation === 'quarantine' + ? 'Workspace requires inspection before reset.' + : 'Local workspace cleanup confirmed.', + }, + this.options.leaseAckTransportTimeoutMs ?? + DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS, + signal, + ); + return; + } catch (error) { + failure = error; + if ( + error instanceof BridgeProtocolError && + error.status != null && + error.status >= 400 && + error.status < 500 && + error.status !== 408 && + error.status !== 429 + ) + throw error; + await abortableDelay(100 * (attempt + 1), signal); + } + } + throw ( + failure ?? + signal?.reason ?? + new Error('Workspace ownership reporting aborted') + ); + } + private assignmentRemainingMs(assignment: BridgeAssignment): number { if ( Number.isSafeInteger(assignment.remainingMs) && @@ -1768,6 +1880,9 @@ export class BridgeWorker { this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, ), ); + if (assignment.workspaceLeaseSlot !== undefined) { + await this.reportWorkspaceOwnership(assignment, 'workspace-cleanup'); + } } finally { heartbeatController.abort(); await heartbeat; diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts index 8bc53f50..43715b37 100644 --- a/service/src/bridge/concurrent-store.test.ts +++ b/service/src/bridge/concurrent-store.test.ts @@ -47,7 +47,7 @@ function dispatch(workspaceId: string, signal = new AbortController().signal) { void promise.catch(() => undefined); return promise; } -async function settle(assignment: CodeBridgeAssignment) { +async function settle(assignment: CodeBridgeAssignment, cleanup = true) { await store.acknowledgeLease( workerId, incarnationId, @@ -63,7 +63,110 @@ async function settle(assignment: CodeBridgeAssignment) { status: 'rejected', error: 'fixture clean rejection', }); + if (cleanup) + await store.confirmWorkspaceCleanup(workerId, assignment.assignmentId, { + protocolVersion: 1, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected', + error: 'local cleanup confirmed', + }); } +test('committed results retain the root fence until cleanup, including receipt expiry', async () => { + await register(); + const pending = dispatch('a'); + const assignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await store.acknowledgeLease( + workerId, + incarnationId, + assignment.assignmentId, + assignment.generation, + assignment.leaseToken, + ); + const intent = { + protocolVersion: 1 as const, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected' as const, + error: 'fixture', + }; + await store.settle(workerId, assignment.assignmentId, intent); + await pending; + const fenceKey = ( + await redis.keys( + `codeapi:bridge:v1:worker:${workerId}:workspace:*:quarantined`, + ) + )[0]; + expect(await redis.get(fenceKey)).toBe(assignment.assignmentId); + await redis.del( + `codeapi:bridge:v1:assignment:${assignment.assignmentId}:workspace-fence-owner`, + ); + await expect( + store.confirmWorkspaceCleanup(workerId, assignment.assignmentId, intent), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + expect(await redis.get(fenceKey)).toBe(assignment.assignmentId); + expect(await redis.ttl(fenceKey)).toBe(-1); + const healthy = dispatch('b'); + const next = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 1, + ))!; + await settle(next); + await healthy; +}); +test('late quarantine after confirmed cleanup cannot fence a newer assignment', async () => { + await register(); + const first = dispatch('a'); + const assignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await settle(assignment); + await first; + const second = dispatch('a'); + const next = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await store.settle( + workerId, + assignment.assignmentId, + { + protocolVersion: 1, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected', + error: 'lost cleanup response', + }, + undefined, + undefined, + true, + ); + await settle(next); + await expect(second).resolves.toMatchObject({ status: 'rejected' }); +}); test('store routes simultaneous roots through separate acknowledged slots', async () => { await register(); const a = dispatch('a'); @@ -283,8 +386,8 @@ test('post-settlement fences are authenticated, idempotent, and invalidated by r undefined, 0, ))!; - await settle(assignment); - await pending; // Normal dispatch cleanup has removed the full assignment. + await settle(assignment, false); + await pending; // Result is committed, but local cleanup remains outstanding. const receiptKey = `codeapi:bridge:v1:assignment:${assignment.assignmentId}:workspace-fence-owner`; expect(await redis.ttl(receiptKey)).toBeGreaterThan(0); expect( diff --git a/service/src/bridge/concurrent-worker.test.ts b/service/src/bridge/concurrent-worker.test.ts index 8d9e7962..16fcc7ab 100644 --- a/service/src/bridge/concurrent-worker.test.ts +++ b/service/src/bridge/concurrent-worker.test.ts @@ -7,8 +7,17 @@ import { WorkspaceToolError } from '../../../packages/code/src/workspace'; import type { WorkspaceMutationQuarantine } from '../../../packages/code/src/worker'; import type { BridgeWorkspaceToolCapabilities } from '../../../packages/code/src/protocol'; -for (const failure of ['execution', 'cleanup', 'post-unlink']) { - const cleanupFailure = failure !== 'execution'; +for (const failure of [ + 'execution', + 'cleanup', + 'post-unlink', + 'hung-cleanup', + 'lost-response', + 'delivery-outage', +]) { + const cleanupFailure = ['cleanup', 'post-unlink', 'hung-cleanup'].includes( + failure, + ); test(`concurrent worker isolates ${failure} failure`, async () => { const redis = new RedisMock() as unknown as Redis; const store = new RedisBridgeStore(redis, 60, 1000, 2); @@ -26,6 +35,10 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { pending.add(root); }, async clear() { + if (failure === 'hung-cleanup' && root === 'a') { + pending.delete(root); + await new Promise(() => {}); + } if (failure === 'post-unlink') pending.delete(root); if (cleanupFailure && root === 'a') throw new Error('injected guard cleanup failure'); @@ -46,6 +59,7 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { }); const started = new Set(); const errors: unknown[] = []; + let quarantineAttempts = 0; let registered!: () => void; const ready = new Promise((resolve) => { registered = resolve; @@ -57,6 +71,7 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { incarnationId, sandboxEndpoint: 'http://sandbox.invalid', leaseWaitMs: 50, + workspaceCleanupTimeoutMs: 20, capabilities: { statefulWorkspace: false, sandboxProfile: 'native-srt', @@ -154,7 +169,15 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { signal, ), }; + } else if (path.endsWith('/workspace-cleanup')) { + await store.confirmWorkspaceCleanup(workerId, id, body, signal); + result = { protocolVersion: 1, accepted: true }; } else { + if (path.endsWith('/quarantine')) { + quarantineAttempts++; + if (failure === 'delivery-outage') + throw new TypeError('injected transport outage'); + } await store.settle( workerId, id, @@ -163,6 +186,12 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { undefined, path.endsWith('/quarantine'), ); + if ( + path.endsWith('/quarantine') && + failure === 'lost-response' && + quarantineAttempts === 1 + ) + throw new TypeError('injected lost response after commit'); result = { protocolVersion: 1, accepted: true }; } } @@ -188,21 +217,34 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { }), ), ); - if (!cleanupFailure) + if (failure === 'delivery-outage') + expect(results[0]).toMatchObject({ + status: 'rejected', + reason: { code: 'ASSIGNMENT_EXPIRED' }, + }); + else if (!cleanupFailure) expect(results[0]).toMatchObject({ status: 'fulfilled', value: { status: 'rejected' }, }); else if (results[0].status === 'fulfilled') - expect(results[0].value).toMatchObject({ status: 'fulfilled', result: { stdout: 'completed' } }); - else expect(results[0].reason).toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + expect(results[0].value).toMatchObject({ + status: 'fulfilled', + result: { stdout: 'completed' }, + }); + else + expect(results[0].reason).toMatchObject({ + code: 'WORKSPACE_QUARANTINED', + }); expect(results[1]).toMatchObject({ status: 'fulfilled', value: { status: 'fulfilled' }, }); - for (let i = 0; i < 100 && errors.length === 0; i++) + for (let i = 0; i < 300 && errors.length === 0; i++) await new Promise((resolve) => setTimeout(resolve, 5)); expect(errors.length).toBe(1); + if (failure === 'lost-response') expect(quarantineAttempts).toBe(2); + if (failure === 'delivery-outage') expect(quarantineAttempts).toBe(3); await expect( store.dispatchWorkspaceTool({ workerId, @@ -215,7 +257,12 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { command: 'must not execute', }, }), - ).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + ).rejects.toMatchObject({ + code: + failure === 'delivery-outage' + ? 'ASSIGNMENT_EXPIRED' + : 'WORKSPACE_QUARANTINED', + }); await expect( store.dispatchWorkspaceTool({ workerId, @@ -229,7 +276,9 @@ for (const failure of ['execution', 'cleanup', 'post-unlink']) { }, }), ).resolves.toMatchObject({ status: 'fulfilled' }); - expect([...pending]).toEqual(failure === 'post-unlink' ? [] : ['a']); + expect([...pending]).toEqual( + ['post-unlink', 'hung-cleanup'].includes(failure) ? [] : ['a'], + ); expect(started.size).toBe(2); } catch (error) { throw new AggregateError( diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 90ded8a8..72e51b4d 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -641,6 +641,7 @@ router.post( [ '/workers/:workerId/assignments/:assignmentId/settle', '/workers/:workerId/assignments/:assignmentId/quarantine', + '/workers/:workerId/assignments/:assignmentId/workspace-cleanup', ], workerAuth, asyncRoute(async (req, res) => { @@ -655,7 +656,11 @@ router.post( req.once('aborted', abortSettlement); res.once('close', abortSettlement); try { - await options.store.settle( + if (req.path.endsWith('/workspace-cleanup')) { + await options.store.confirmWorkspaceCleanup(req.params.workerId, req.params.assignmentId, + settlement, settlementController.signal, + (res.locals.bridgeWorkerAuthorization as {identityId: string} | undefined)?.identityId); + } else await options.store.settle( req.params.workerId, req.params.assignmentId, settlement, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index c15cf425..80d5e9fe 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -1523,7 +1523,7 @@ export class RedisBridgeStore { "if redis.call('GET', KEYS[#KEYS]) ~= ARGV[7] then return -4 end", 'redis.call(\'SET\', KEYS[2], ARGV[1], \"EX\", ARGV[2])', "if redis.call('GET', KEYS[3]) == ARGV[3] then redis.call('DEL', KEYS[3], KEYS[4]) end", - 'if ARGV[6] == "1" and ARGV[4] == "rejected" then', + 'if ARGV[6] == "1" and ARGV[4] == "rejected" and ARGV[8] ~= "1" then', " redis.call('DEL', KEYS[6])", 'end', 'return 1', @@ -1541,6 +1541,7 @@ export class RedisBridgeStore { identityId ?? '', hasWorkspace ? '1' : '0', settlement.incarnationId, + assignment.workspaceLeaseSlot === undefined ? '0' : '1', ), signal, 'Bridge settlement commit', @@ -1582,6 +1583,7 @@ export class RedisBridgeStore { ) { // The dispatcher may already have timed out and finished its cleanup. // Release only this settled reservation, retaining a quarantine marker. + await this.commitPendingWorkspace(assignment, settlement); await this.cleanupWithRetry(workerId, assignmentId, assignment); } } @@ -1653,12 +1655,30 @@ export class RedisBridgeStore { ); } + async confirmWorkspaceCleanup( + workerId: string, + assignmentId: string, + intent: AnyCodeBridgeSettlement, + signal?: AbortSignal, + identityId?: string, + ): Promise { + await this.quarantineSettledWorkspace( + workerId, + assignmentId, + intent, + signal, + identityId, + false, + ); + } + private async quarantineSettledWorkspace( workerId: string, assignmentId: string, settlement: AnyCodeBridgeSettlement, signal?: AbortSignal, identityId?: string, + quarantine = true, ): Promise { // Keep a small, expiring ownership receipt separate from assignment cleanup. // A local guard can fail to clear after the result has already committed. @@ -1699,9 +1719,19 @@ export class RedisBridgeStore { "if redis.call('GET', KEYS[2]) ~= ARGV[2] then return 0 end", "if (redis.call('GET', KEYS[3]) or '') ~= ARGV[3] then return 0 end", "if (redis.call('GET', KEYS[4]) or '0') ~= ARGV[4] then return 0 end", - "redis.call('SET', KEYS[5], 'quarantined:' .. ARGV[5])", + 'if ARGV[8] == "1" then', + // A completed cleanup receipt is terminal: a lost response must + // not let a late quarantine overwrite a newer root owner. + " if redis.call('HGET', KEYS[1], 'localCleanup') == '1' and redis.call('HGET', KEYS[1], 'resultCommitted') == '1' then return 1 end", + " redis.call('SET', KEYS[5], 'quarantined:' .. ARGV[5])", // Never replace a committed result. Before settlement, terminate the waiter. - "redis.call('SET', KEYS[6], ARGV[6], 'EX', ARGV[7], 'NX')", + " redis.call('SET', KEYS[6], ARGV[6], 'EX', ARGV[7], 'NX')", + 'else', + " local fence = redis.call('GET', KEYS[5])", + " if fence and string.sub(fence, 1, 12) == 'quarantined:' then return 0 end", + " redis.call('HSET', KEYS[1], 'localCleanup', '1')", + " if redis.call('HGET', KEYS[1], 'resultCommitted') == '1' and fence == ARGV[5] then redis.call('DEL', KEYS[5]) end", + 'end', "if redis.call('GET', KEYS[7]) == ARGV[5] then redis.call('DEL', KEYS[7]) end", "if redis.call('GET', KEYS[8]) == ARGV[5] then redis.call('DEL', KEYS[8]) end", 'return 1', @@ -1730,6 +1760,7 @@ export class RedisBridgeStore { assignmentId, JSON.stringify(settlement), assignmentTtlSeconds(Date.parse(receipt.expiresAt)), + quarantine ? '1' : '0', ), signal, 'Workspace quarantine fence commit', @@ -2056,6 +2087,35 @@ export class RedisBridgeStore { assignment: StoredAssignment, settlement: AnyCodeBridgeSettlement, ): Promise { + if (assignment.workspaceLeaseSlot !== undefined) { + const committed = Number( + await boundedCommand( + this.redis.eval( + [ + "if redis.call('EXISTS', KEYS[2]) == 0 then return 0 end", + "redis.call('HSET', KEYS[2], 'resultCommitted', '1')", + "if redis.call('HGET', KEYS[2], 'localCleanup') == '1' and redis.call('GET', KEYS[1]) == ARGV[1] then redis.call('DEL', KEYS[1]) end", + 'return 1', + ].join('\n'), + 2, + workspaceQuarantineKey( + assignment.workerId, + assignment.workspaceFence!, + ), + workspaceFenceReceiptKey(assignment.assignmentId), + assignment.assignmentId, + ), + this.redisCommandTimeoutMs, + 'Bridge native workspace result commit', + ), + ); + if (committed !== 1) + throw new BridgeStoreError( + 'WORKSPACE_QUARANTINED', + 'Native workspace cleanup ownership expired', + ); + return; + } if ( assignmentWorkspace(assignment) === undefined || settlement.status !== 'fulfilled' @@ -2144,6 +2204,7 @@ export class RedisBridgeStore { 'if claimed and not acknowledged then', " redis.call('DEL', KEYS[3], KEYS[4])", 'end', + 'if ARGV[3] == "1" and redis.call(\'GET\', KEYS[5]) == ARGV[1] then return -1 end', 'if queued == 0 and acknowledged and ARGV[2] == "1" and redis.call(\'GET\', KEYS[5]) == ARGV[1] then', ' return -1', 'end', @@ -2161,6 +2222,7 @@ export class RedisBridgeStore { ...keys, assignment.assignmentId, assignmentWorkspace(assignment) === undefined ? '0' : '1', + assignment.workspaceLeaseSlot === undefined ? '0' : '1', ), this.redisCommandTimeoutMs, 'Bridge assignment cleanup', From 382b165c41def5d148b9d95e964de1345383d9ab Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 9 Sep 2026 00:06:06 -0400 Subject: [PATCH 6/7] fix: Preserve Slot Liveness and Retry Finalization --- packages/code/README.md | 5 +- packages/code/src/worker.ts | 10 ++- service/src/bridge/concurrent-store.test.ts | 65 ++++++++++++++++++++ service/src/bridge/concurrent-worker.test.ts | 16 +++-- service/src/bridge/store.ts | 7 ++- 5 files changed, 93 insertions(+), 10 deletions(-) diff --git a/packages/code/README.md b/packages/code/README.md index 3573a696..819d9f5c 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -567,8 +567,9 @@ receipts exclude command payloads; explicit reset invalidates old fence requests The server releases a root only after result finalization **and** explicit local cleanup confirmation. Local guard cleanup has a five-second bound; an expired receipt never implies a clean root. Control receipt delivery retries three times. -If delivery remains unavailable, that lease lane stops while healthy lanes keep -running; inspect/reset the affected root and restart the worker to restore capacity. +If delivery remains unavailable, the root remains fenced while every advertised +lane keeps polling. Capacity becomes reusable when its owned reservation is +released or expires; inspect/reset the affected root before using it again. Reset-only registration stays unready and cannot attract new assignments. To recover a quarantined native root: diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index be2badbf..207cbfae 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -931,10 +931,14 @@ export class BridgeWorker { (quarantineError.status === 401 || quarantineError.status === 403 || quarantineError.code === 'WORKER_FENCED') - ) + ) { fail(quarantineError); - else this.options.onError?.(quarantineError); - return; + return; + } + this.options.onError?.(quarantineError); + // Keep every advertised slot polled. The root remains fenced, + // but a committed receipt may already have released this slot. + continue; } } if (controller.signal.aborted) return; diff --git a/service/src/bridge/concurrent-store.test.ts b/service/src/bridge/concurrent-store.test.ts index 43715b37..fe53c44f 100644 --- a/service/src/bridge/concurrent-store.test.ts +++ b/service/src/bridge/concurrent-store.test.ts @@ -167,6 +167,71 @@ test('late quarantine after confirmed cleanup cannot fence a newer assignment', await settle(next); await expect(second).resolves.toMatchObject({ status: 'rejected' }); }); +test('duplicate rejected settlement retries finalization after the dispatcher leaves', async () => { + await register(); + const controller = new AbortController(); + const pending = dispatch('a', controller.signal); + const assignment = (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!; + await store.acknowledgeLease( + workerId, + incarnationId, + assignment.assignmentId, + assignment.generation, + assignment.leaseToken, + ); + controller.abort(); + await pending.catch(() => undefined); + const intent = { + protocolVersion: 1 as const, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'rejected' as const, + error: 'clean rejection', + }; + const originalEval = redis.eval.bind(redis); + let failed = false; + redis.eval = ((script: string, ...args: unknown[]) => { + if (!failed && script.includes("'resultCommitted', '1'")) { + failed = true; + return Promise.reject(new Error('injected finalization outage')); + } + return (originalEval as (...args: unknown[]) => unknown)(script, ...args); + }) as typeof redis.eval; + try { + await expect( + store.settle(workerId, assignment.assignmentId, intent), + ).rejects.toThrow('injected finalization outage'); + } finally { + redis.eval = originalEval; + } + expect(failed).toBe(true); + await store.settle(workerId, assignment.assignmentId, intent); + await store.confirmWorkspaceCleanup( + workerId, + assignment.assignmentId, + intent, + ); + const next = dispatch('a'); + await settle( + (await store.lease( + workerId, + incarnationId, + 1000, + undefined, + undefined, + 0, + ))!, + ); + await expect(next).resolves.toMatchObject({ status: 'rejected' }); +}); test('store routes simultaneous roots through separate acknowledged slots', async () => { await register(); const a = dispatch('a'); diff --git a/service/src/bridge/concurrent-worker.test.ts b/service/src/bridge/concurrent-worker.test.ts index 16fcc7ab..5aec49cf 100644 --- a/service/src/bridge/concurrent-worker.test.ts +++ b/service/src/bridge/concurrent-worker.test.ts @@ -13,6 +13,7 @@ for (const failure of [ 'post-unlink', 'hung-cleanup', 'lost-response', + 'all-responses-lost', 'delivery-outage', ]) { const cleanupFailure = ['cleanup', 'post-unlink', 'hung-cleanup'].includes( @@ -58,6 +59,7 @@ for (const failure of [ startBoth = resolve; }); const started = new Set(); + let failedRootExecutions = 0; const errors: unknown[] = []; let quarantineAttempts = 0; let registered!: () => void; @@ -88,6 +90,7 @@ for (const failure of [ capabilities, mutationFailuresAreAtomic: true, async execute(request) { + if (request.workspaceId === 'a') failedRootExecutions++; started.add(request.workspaceId); if (started.size === 2) startBoth(); await bothStarted; @@ -188,8 +191,8 @@ for (const failure of [ ); if ( path.endsWith('/quarantine') && - failure === 'lost-response' && - quarantineAttempts === 1 + ((failure === 'lost-response' && quarantineAttempts === 1) || + failure === 'all-responses-lost') ) throw new TypeError('injected lost response after commit'); result = { protocolVersion: 1, accepted: true }; @@ -242,9 +245,13 @@ for (const failure of [ }); for (let i = 0; i < 300 && errors.length === 0; i++) await new Promise((resolve) => setTimeout(resolve, 5)); - expect(errors.length).toBe(1); + if (failure === 'delivery-outage') + expect(errors.length).toBeGreaterThanOrEqual(1); + else expect(errors.length).toBe(1); if (failure === 'lost-response') expect(quarantineAttempts).toBe(2); - if (failure === 'delivery-outage') expect(quarantineAttempts).toBe(3); + if (failure === 'delivery-outage') + expect(quarantineAttempts).toBeGreaterThanOrEqual(3); + if (failure === 'all-responses-lost') expect(quarantineAttempts).toBe(3); await expect( store.dispatchWorkspaceTool({ workerId, @@ -280,6 +287,7 @@ for (const failure of [ ['post-unlink', 'hung-cleanup'].includes(failure) ? [] : ['a'], ); expect(started.size).toBe(2); + expect(failedRootExecutions).toBe(1); } catch (error) { throw new AggregateError( [error, ...errors], diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 80d5e9fe..41e11839 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -1426,7 +1426,6 @@ export class RedisBridgeStore { signal, 'Bridge settlement existing read', ); - if (existingSettlement === serializedSettlement) return; if ( existingSettlement != null && existingSettlement !== serializedSettlement @@ -1441,6 +1440,12 @@ export class RedisBridgeStore { signal, 'Bridge settlement assignment read', ); + if ( + existingSettlement === serializedSettlement && + (assignment?.workspaceLeaseSlot === undefined || + settlement.status !== 'rejected') + ) + return; if (assignment == null) { if (existingSettlement === serializedSettlement) return; throw new BridgeStoreError( From 6ebad00deaeba0ec6914d4bb86679929c116ad45 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 9 Sep 2026 00:23:50 -0400 Subject: [PATCH 7/7] fix: Close Startup and Executor Recovery Gaps --- docker-compose.yaml | 2 + packages/code/src/native-pool.test.ts | 42 +++++++++++ packages/code/src/native-pool.ts | 19 +++++ packages/code/src/worker-slots.test.ts | 97 ++++++++++++++++++++++++++ packages/code/src/worker.ts | 33 ++++++++- tests/compose-bridge-config.cjs | 3 + 6 files changed, 195 insertions(+), 1 deletion(-) diff --git a/docker-compose.yaml b/docker-compose.yaml index 3409554c..ab37fd24 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -13,6 +13,7 @@ services: - CODEAPI_BRIDGE_TOKEN=${CODEAPI_BRIDGE_TOKEN:-} - CODEAPI_BRIDGE_AUTH_MODE=${CODEAPI_BRIDGE_AUTH_MODE:-paired} - CODEAPI_BRIDGE_DYNAMIC_WORKERS=${CODEAPI_BRIDGE_DYNAMIC_WORKERS:-true} + - CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS=${CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS:-1} - CODEAPI_BRIDGE_WORKER_ID=${CODEAPI_BRIDGE_WORKER_ID:-} - CODEAPI_AUTH_PROVIDER=${CODEAPI_AUTH_PROVIDER:-} - CODEAPI_ALLOW_AUTH_PROVIDER_NONE=${CODEAPI_ALLOW_AUTH_PROVIDER_NONE:-} @@ -61,6 +62,7 @@ services: - CODEAPI_BRIDGE_TOKEN=${CODEAPI_BRIDGE_TOKEN:-} - CODEAPI_BRIDGE_AUTH_MODE=${CODEAPI_BRIDGE_AUTH_MODE:-paired} - CODEAPI_BRIDGE_DYNAMIC_WORKERS=${CODEAPI_BRIDGE_DYNAMIC_WORKERS:-true} + - CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS=${CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS:-1} - CODEAPI_BRIDGE_WORKER_ID=${CODEAPI_BRIDGE_WORKER_ID:-} - CODEAPI_AUTH_PROVIDER=${CODEAPI_AUTH_PROVIDER:-} - CODEAPI_JWT_SINGLE_TENANT_ID=${CODEAPI_JWT_SINGLE_TENANT_ID:-} diff --git a/packages/code/src/native-pool.test.ts b/packages/code/src/native-pool.test.ts index c37a9b92..3a24db34 100644 --- a/packages/code/src/native-pool.test.ts +++ b/packages/code/src/native-pool.test.ts @@ -1,6 +1,7 @@ import test from 'node:test'; import assert from 'node:assert/strict'; import { NativeWorkspaceCommandPool } from './native-pool.js'; +import { WorkspaceToolError } from './workspace.js'; import type { WorkspaceExecuteCommandRequest } from './protocol.js'; const roots = new Map( @@ -12,6 +13,47 @@ const request = (workspaceId: string): WorkspaceExecuteCommandRequest => ({ workspaceId, command: 'fixture', }); +test('a known-clean executor failure is retired without replaying the command', async () => { + let created = 0; + let executed = 0; + let closed = 0; + const pool = new NativeWorkspaceCommandPool(roots, 2, () => { + const first = ++created === 1; + return { + async prepare() {}, + async close() { + closed++; + }, + async execute(req) { + executed++; + if (first) + throw new WorkspaceToolError( + 'prepare failed', + 'COMMAND_UNAVAILABLE', + false, + ); + return { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: req.workspaceId, + stdout: '', + stderr: '', + exitCode: 0, + truncated: false, + timedOut: false, + }; + }, + }; + }); + await assert.rejects(pool.execute(request('b')), { + code: 'COMMAND_UNAVAILABLE', + }); + assert.equal(executed, 1); + assert.equal(closed, 1); + await pool.execute(request('b')); + assert.equal(created, 2); + await pool.close(); +}); test('native pool reuses roots and evicts only idle processes within its bound', async () => { const created: string[] = []; const closed: string[] = []; diff --git a/packages/code/src/native-pool.ts b/packages/code/src/native-pool.ts index 0cdd1d2e..28f155e2 100644 --- a/packages/code/src/native-pool.ts +++ b/packages/code/src/native-pool.ts @@ -105,13 +105,32 @@ export class NativeWorkspaceCommandPool { signal?: AbortSignal, ): Promise { const entry = await this.allocate(request.workspaceId); + let enteredExecutor = false; try { if (signal?.aborted) throw new WorkspaceToolError( 'Command cancelled before dispatch', 'EXECUTION_ABORTED', ); + enteredExecutor = true; return await entry.sandbox.execute(request, signal); + } catch (error) { + if ( + enteredExecutor && + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted + ) { + // Never retry the command here. Retire a failed executor only after + // close succeeds, allowing a later assignment to create a fresh child. + try { + await entry.sandbox.close(); + if (this.entries.get(request.workspaceId) === entry) + this.entries.delete(request.workspaceId); + } catch { + /* Retain ownership for subsequent cleanup/shutdown. */ + } + } + throw error; } finally { entry.busy = false; } diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts index ba6f37eb..ba65a3cf 100644 --- a/packages/code/src/worker-slots.test.ts +++ b/packages/code/src/worker-slots.test.ts @@ -11,6 +11,103 @@ const capabilities: BridgeWorkspaceToolCapabilities = { operations: ['read_file'], workspaces: [{ id: 'a' }, { id: 'b' }], }; +for (const requestedSlots of [1, 2]) { + test(`mapped quarantine blocks serial readiness with ${requestedSlots} requested slots`, async () => { + const paths: string[] = []; + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + incarnationId: 'incarnation-guard', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + workspaceLeaseSlots: requestedSlots, + requiresReadyConfirmation: true, + workspaceTools: capabilities, + }, + workspaceTools: { + capabilities, + async execute() { + throw new Error('must not execute'); + }, + }, + workspaceQuarantines: new Map([ + [ + 'a', + { + async assertAvailable() { + throw new Error('retained guard'); + }, + async arm() {}, + async clear() {}, + async quarantine() {}, + }, + ], + [ + 'b', + { + async assertAvailable() {}, + async arm() {}, + async clear() {}, + async quarantine() {}, + }, + ], + ]), + fetchImpl: async (url) => { + paths.push(new URL(String(url)).pathname); + return Response.json({ + protocolVersion: 1, + workerId: 'worker', + incarnationId: 'incarnation-guard', + registrationGeneration: 1, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60000, + workspaceLeaseSlots: 1, + }); + }, + }); + await assert.rejects(worker.register(), { code: 'WORKER_QUARANTINED' }); + assert.ok(paths.every((path) => path.endsWith('/register'))); + if (requestedSlots === 1) assert.equal(paths.length, 0); + }); +} +test('clean rejection receipt failure uses lane-local quarantine classification', async () => { + const worker = new BridgeWorker({ + codeApiUrl: 'http://localhost:1', + token: 'fixture', + workerId: 'worker', + sandboxEndpoint: 'http://localhost:2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'fixture', + runtimes: [], + }, + }); + const internals = worker as unknown as { + maintainRegistration: () => Promise; + settleWithRetry: () => Promise; + reportWorkspaceOwnership: () => Promise; + rejectUnexecutedAssignment: ( + assignment: BridgeAssignment, + message: string, + ) => Promise; + }; + internals.maintainRegistration = async () => {}; + internals.settleWithRetry = async () => {}; + internals.reportWorkspaceOwnership = async () => { + throw new TypeError('receipt outage'); + }; + await assert.rejects( + internals.rejectUnexecutedAssignment( + { workspaceLeaseSlot: 0 } as BridgeAssignment, + 'expired', + ), + { name: 'BridgeWorkspaceQuarantinedError' }, + ); +}); test('maintenance registration never advertises readiness or starts leasing', async () => { const paths: string[] = []; const worker = new BridgeWorker({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 207cbfae..545b96ca 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -492,6 +492,13 @@ export class BridgeWorker { if (!allowActiveMutation) { try { await this.options.workspaceMutationQuarantine?.assertAvailable(); + if ( + !this.maintenanceOnly && + (this.options.capabilities.workspaceLeaseSlots ?? 1) === 1 + ) { + for (const guard of this.options.workspaceQuarantines?.values() ?? []) + await guard.assertAvailable(); + } } catch (error) { if ( error instanceof BridgeProtocolError && @@ -595,6 +602,23 @@ export class BridgeWorker { ); } this.negotiatedWorkspaceSlots = slots; + if ( + !this.maintenanceOnly && + !allowActiveMutation && + slots === 1 && + (this.options.capabilities.workspaceLeaseSlots ?? 1) > 1 + ) { + try { + for (const guard of this.options.workspaceQuarantines?.values() ?? []) + await guard.assertAvailable(); + } catch { + throw new BridgeProtocolError( + 'Serial workspace quarantine state could not be verified', + undefined, + 'WORKER_QUARANTINED', + ); + } + } const registeredAtMs = Date.parse(registration.registeredAt); if (Number.isFinite(registeredAtMs)) { this.serverClockOffsetMs = registeredAtMs - registrationStartedAtMs; @@ -1885,7 +1909,14 @@ export class BridgeWorker { ), ); if (assignment.workspaceLeaseSlot !== undefined) { - await this.reportWorkspaceOwnership(assignment, 'workspace-cleanup'); + try { + await this.reportWorkspaceOwnership(assignment, 'workspace-cleanup'); + } catch (error) { + throw new BridgeWorkspaceQuarantinedError( + 'Unexecuted workspace cleanup acknowledgement could not be confirmed', + error, + ); + } } } finally { heartbeatController.abort(); diff --git a/tests/compose-bridge-config.cjs b/tests/compose-bridge-config.cjs index d6c9797a..de14b553 100644 --- a/tests/compose-bridge-config.cjs +++ b/tests/compose-bridge-config.cjs @@ -14,6 +14,7 @@ function render(overrides) { CODEAPI_BRIDGE_DYNAMIC_WORKERS: '', CODEAPI_BRIDGE_WORKER_ID: '', CODEAPI_BRIDGE_TOKEN: '', + CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS: '', ...overrides, }, })); @@ -26,6 +27,7 @@ for (const overrides of [ CODEAPI_BRIDGE_TOKEN: token, CODEAPI_BRIDGE_DYNAMIC_WORKERS: 'false', CODEAPI_BRIDGE_WORKER_ID: 'test-worker', + CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS: '4', }, ]) { const config = render(overrides); @@ -34,6 +36,7 @@ for (const overrides of [ assert.equal(env.CODEAPI_HARDENED_SANDBOX_MODE, 'true'); assert.equal(env.CODEAPI_BRIDGE_AUTH_MODE, 'paired'); assert.equal(env.CODEAPI_BRIDGE_TOKEN, token); + assert.equal(env.CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS, overrides.CODEAPI_BRIDGE_MAX_WORKSPACE_LEASE_SLOTS ?? '1'); assert.equal(env.CODEAPI_BRIDGE_DYNAMIC_WORKERS, overrides.CODEAPI_BRIDGE_DYNAMIC_WORKERS ?? 'true'); assert.equal(env.CODEAPI_BRIDGE_WORKER_ID, overrides.CODEAPI_BRIDGE_WORKER_ID ?? ''); }