diff --git a/docker-compose.yaml b/docker-compose.yaml index 3409554..ab37fd2 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/README.md b/packages/code/README.md index 7e7c0aa..819d9f5 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -520,3 +520,63 @@ 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. +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 +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. 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, 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: + +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 0000000..484f0dc --- /dev/null +++ b/packages/code/src/cli-slots.test.ts @@ -0,0 +1,143 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { spawnSync } from 'node:child_process'; +import { mkdtemp, mkdir, rm, stat } 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/); +}); + +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 67ec396..27a8d2d 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, @@ -467,21 +471,114 @@ 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) => root.root); + 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'); + } + 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, + 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 +754,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 +834,9 @@ async function run( ) .digest('hex'), ...(fileRelayEnabled ? { requiresReadyConfirmation: true } : {}), + ...(workspaceLeaseSlots > 1 + ? { workspaceLeaseSlots, requiresReadyConfirmation: true } + : {}), ...(workspaceTools ? { workspaceTools: workspaceTools.capabilities } : {}), }; if (!isValidBridgeWorkerCapabilities(capabilities)) { @@ -752,44 +863,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 +961,16 @@ async function run( ); return; } + const resetNativeRoot = option(args, '--reset-workspace-quarantine'); + if (resetNativeRoot != null) { + await worker.refreshCredential(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`, + ); + return; + } await worker.run(controller.signal); } finally { try { diff --git a/packages/code/src/native-pool.test.ts b/packages/code/src/native-pool.test.ts new file mode 100644 index 0000000..3a24db3 --- /dev/null +++ b/packages/code/src/native-pool.test.ts @@ -0,0 +1,154 @@ +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( + ['a', 'b', 'c'].map((id) => [id, { workspaceRoot: `/fixture/${id}` }]), +); +const request = (workspaceId: string): WorkspaceExecuteCommandRequest => ({ + protocolVersion: 1, + operation: 'execute_command', + 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[] = []; + 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 new file mode 100644 index 0000000..28f155e --- /dev/null +++ b/packages/code/src/native-pool.ts @@ -0,0 +1,152 @@ +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: Pick< + NativeProcessWorkspaceCommandSandbox, + 'prepare' | 'execute' | 'close' + >; + 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, + private readonly createSandbox: ( + options: NativeProcessSandboxOptions, + ) => Entry['sandbox'] = (options) => + new NativeProcessWorkspaceCommandSandbox(options), + ) { + 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: this.createSandbox(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; + }); + 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 { + 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); + 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; + } + } + + 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 37288d6..44cfdfd 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -398,6 +398,8 @@ const WORKSPACE_SEARCH_MATCH_KEYS = new Set([ ]); export interface BridgeWorkerCapabilities { + /** Opt-in protocol: maximum concurrently leased independent workspace roots. */ + workspaceLeaseSlots?: number; statefulWorkspace: boolean; sandboxProfile: string; runtimes: string[]; @@ -414,6 +416,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 +468,7 @@ export interface BridgeSandboxRequest { } export interface BridgeAssignment { + workspaceLeaseSlot?: number; protocolVersion: BridgeProtocolVersion; assignmentId: string; workerId: string; @@ -1125,6 +1130,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 0000000..ba65a3c --- /dev/null +++ b/packages/code/src/worker-slots.test.ts @@ -0,0 +1,314 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import { BridgeWorker } from './worker.js'; +import type { + BridgeAssignment, + BridgeWorkspaceToolCapabilities, +} from './protocol.js'; + +const capabilities: BridgeWorkspaceToolCapabilities = { + protocolVersion: 1, + 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({ + 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(); + 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/); + }); +} + +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'); + }); +} + +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 1fda70f..545b96c 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -35,10 +35,13 @@ 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; leaseAckTransportTimeoutMs?: number; + workspaceCleanupTimeoutMs?: number; resetTransportTimeoutMs?: number; cancellationPollIntervalMs?: number; cancellationTransportTimeoutMs?: number; @@ -59,9 +62,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 { @@ -343,10 +350,32 @@ 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< + 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 = @@ -410,13 +455,50 @@ 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, + ): 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 { 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 && @@ -457,7 +539,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, ); @@ -503,6 +587,38 @@ 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; + 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; @@ -510,7 +626,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; @@ -571,7 +690,55 @@ 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, + ); + this.quarantinedWorkspaces.delete(workspaceId); + } + + 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 +768,7 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, waitMs, incarnationId: this.incarnationId, + ...(workspaceLeaseSlot === undefined ? {} : { workspaceLeaseSlot }), }, leaseController.signal, ); @@ -610,7 +778,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', @@ -698,11 +867,19 @@ 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 { 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 +911,132 @@ 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.reportWorkspaceOwnership( + assignment, + 'quarantine', + controller.signal, + ); + this.options.onError?.(error); + continue; + } catch (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); + 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; + 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; @@ -833,6 +1129,101 @@ 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. + 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) + 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); + } 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); + 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 +1288,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 +1307,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 +1336,18 @@ export class BridgeWorker { throw new BridgeProtocolError('Invalid workspace tool request'); } 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) { + throw new BridgeWorkspaceQuarantinedError( + 'Workspace is quarantined', + error, + ); + } const advertised = this.activeCapabilities.workspaceTools; if (advertised == null) { throw new BridgeProtocolError( @@ -983,7 +1390,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 +1427,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 +1450,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 +1584,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 +1597,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 +1612,7 @@ export class BridgeWorker { assignment.runtimeSessionId, `Stateful workspace ${assignment.runtimeSessionId} was quarantined after an ambiguous sandbox execution`, ambiguousSandboxError, + assignment, ); } const knownCleanStatefulRejection = @@ -1242,7 +1652,37 @@ export class BridgeWorker { } if (workspaceMutationArmed) { try { - await this.options.workspaceMutationQuarantine!.clear(); + 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) + ) { + this.armedWorkspaces.delete(assignment.request.workspaceId); + } this.mutationGuardArmed = false; } catch (error) { throw new BridgeWorkspaceQuarantinedError( @@ -1251,6 +1691,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 { @@ -1262,6 +1716,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) && @@ -1307,6 +1808,7 @@ export class BridgeWorker { assignment.runtimeSessionId, `Stateful workspace ${assignment.runtimeSessionId} could not release its runtime lease`, error, + assignment, ); } } @@ -1315,13 +1817,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 +1859,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 && @@ -1403,6 +1908,16 @@ export class BridgeWorker { this.options.rejectionAckGraceMs ?? REJECTION_ACK_GRACE_MS, ), ); + if (assignment.workspaceLeaseSlot !== undefined) { + try { + await this.reportWorkspaceOwnership(assignment, 'workspace-cleanup'); + } catch (error) { + throw new BridgeWorkspaceQuarantinedError( + 'Unexecuted workspace cleanup acknowledgement could not be confirmed', + error, + ); + } + } } finally { heartbeatController.abort(); await heartbeat; @@ -1439,6 +1954,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 +1999,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 +2026,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 0000000..98b6bb0 --- /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 1e19286..4002ed6 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 0000000..fe53c44 --- /dev/null +++ b/service/src/bridge/concurrent-store.test.ts @@ -0,0 +1,533 @@ +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(workspaceLeaseSlots = 2) { + const generation = await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + runtimes: [], + sandboxProfile: 'native-srt', + requiresReadyConfirmation: true, + workspaceLeaseSlots, + 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, cleanup = true) { + 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', + }); + 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('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'); + 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('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'); + 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]); +}); + +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', + }); +}); + +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, 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( + 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 new file mode 100644 index 0000000..5aec49c --- /dev/null +++ b/service/src/bridge/concurrent-worker.test.ts @@ -0,0 +1,303 @@ +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'; + +for (const failure of [ + 'execution', + 'cleanup', + 'post-unlink', + 'hung-cleanup', + 'lost-response', + 'all-responses-lost', + '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); + 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 === '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'); + 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(); + let failedRootExecutions = 0; + const errors: unknown[] = []; + let quarantineAttempts = 0; + 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, + workspaceCleanupTimeoutMs: 20, + 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) { + if (request.workspaceId === 'a') failedRootExecutions++; + 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( + 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 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, + body, + signal, + undefined, + path.endsWith('/quarantine'), + ); + if ( + path.endsWith('/quarantine') && + ((failure === 'lost-response' && quarantineAttempts === 1) || + failure === 'all-responses-lost') + ) + throw new TypeError('injected lost response after commit'); + 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.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 (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[1]).toMatchObject({ + status: 'fulfilled', + value: { status: 'fulfilled' }, + }); + for (let i = 0; i < 300 && errors.length === 0; i++) + await new Promise((resolve) => setTimeout(resolve, 5)); + 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).toBeGreaterThanOrEqual(3); + if (failure === 'all-responses-lost') expect(quarantineAttempts).toBe(3); + await expect( + store.dispatchWorkspaceTool({ + workerId, + signal: controller.signal, + deadlineAtMs: Date.now() + 1000, + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: 'a', + command: 'must not execute', + }, + }), + ).rejects.toMatchObject({ + code: + failure === 'delivery-outage' + ? 'ASSIGNMENT_EXPIRED' + : '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( + ['post-unlink', 'hung-cleanup'].includes(failure) ? [] : ['a'], + ); + expect(started.size).toBe(2); + expect(failedRootExecutions).toBe(1); + } 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 08ef1b0..fc409ad 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 f428bdb..72e51b4 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -396,6 +396,9 @@ router.post( registrationGeneration, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, + workspaceLeaseSlots: options.store.workspaceLeaseCapacity( + registration.capabilities.workspaceLeaseSlots, + ), supportedWorkspaceToolOperations: [ 'read_file', 'search_text', @@ -529,7 +532,11 @@ router.post( body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || !validIncarnationId(body.incarnationId) || !Number.isFinite(requestedWait) || - requestedWait < 0 + 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; @@ -557,6 +564,7 @@ router.post( | { identityId: string } | undefined )?.identityId, + body.workspaceLeaseSlot === undefined ? undefined : Number(body.workspaceLeaseSlot), ); if (leaseController.signal.aborted) { if (assignment != null) await options.store.returnLease(assignment); @@ -630,7 +638,11 @@ router.post( ); router.post( - '/workers/:workerId/assignments/:assignmentId/settle', + [ + '/workers/:workerId/assignments/:assignmentId/settle', + '/workers/:workerId/assignments/:assignmentId/quarantine', + '/workers/:workerId/assignments/:assignmentId/workspace-cleanup', + ], workerAuth, asyncRoute(async (req, res) => { const settlement = req.body as unknown; @@ -644,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, @@ -654,6 +670,7 @@ router.post( | { identityId: string } | undefined )?.identityId, + req.path.endsWith('/quarantine'), ); if (!settlementController.signal.aborted) { res.json({ diff --git a/service/src/bridge/slots.test.ts b/service/src/bridge/slots.test.ts new file mode 100644 index 0000000..dfe8949 --- /dev/null +++ b/service/src/bridge/slots.test.ts @@ -0,0 +1,93 @@ +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'); +}); + +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 new file mode 100644 index 0000000..9c9e7d3 --- /dev/null +++ b/service/src/bridge/slots.ts @@ -0,0 +1,139 @@ +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( + [ + 'local latest = 0', + `for slot = 0, ${MAX_WORKSPACE_LEASE_SLOTS - 1} do`, + " 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'), + 3, + ...this.keys(workerId).slice(0, 3), + incarnationId, + assignmentId, + ); + } +} diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index a00e4f2..41e1183 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; @@ -64,6 +65,31 @@ export class BridgeStoreError extends Error { interface StoredAssignment extends CodeBridgeAssignment { leaseTokenHash: string; workerIdentityId?: string; + workspaceFence?: string; +} + +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; } export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { @@ -173,16 +199,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 +321,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, @@ -379,12 +432,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 +471,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 +484,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 +516,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 +532,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 +543,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 +583,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; } @@ -622,12 +708,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( @@ -642,6 +733,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 && @@ -692,44 +792,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 +900,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 +934,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 +1000,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 +1062,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 +1082,32 @@ 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 >= this.maxWorkspaceLeaseSlots || + 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 +1131,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 +1158,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 +1167,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 +1179,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 +1192,7 @@ export class RedisBridgeStore { const { leaseTokenHash: _leaseTokenHash, workerIdentityId: _workerIdentityId, + workspaceFence: _workspaceFence, ...wireAssignment } = assignment; return { @@ -1012,6 +1207,7 @@ export class RedisBridgeStore { workerId, incarnationId, assignmentId, + slot, ); if (signalAborted(signal)) return undefined; throw error; @@ -1061,8 +1257,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 +1278,7 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, identityId?: string, + slot?: number, ): Promise { const result = await this.redis.eval( [ @@ -1099,8 +1296,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 +1308,7 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, assignmentId: string, + slot?: number, ): Promise { await this.redis.eval( [ @@ -1120,8 +1318,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 +1329,7 @@ export class RedisBridgeStore { assignment.workerId, assignment.incarnationId, assignment.assignmentId, + assignment.workspaceLeaseSlot, ); } @@ -1138,6 +1337,7 @@ export class RedisBridgeStore { workerId: string, incarnationId: string, assignmentId: string, + slot?: number, ): Promise { await boundedCommand( this.redis.eval( @@ -1153,9 +1353,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 +1367,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 +1385,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 +1396,7 @@ export class RedisBridgeStore { 1, workspaceQuarantineKey( assignment.workerId, - assignment.runtimeSessionId, + assignmentWorkspace(assignment)!, ), assignment.assignmentId, ); @@ -1207,15 +1408,28 @@ export class RedisBridgeStore { settlement: AnyCodeBridgeSettlement, signal?: AbortSignal, 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) return; - if (existingSettlement != null) { + if ( + existingSettlement != null && + existingSettlement !== serializedSettlement + ) { throw new BridgeStoreError( 'ASSIGNMENT_FENCED', 'Bridge assignment was already settled with a different result', @@ -1226,7 +1440,14 @@ 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( 'ASSIGNMENT_NOT_FOUND', 'Bridge assignment was not found', @@ -1268,38 +1489,48 @@ 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" and ARGV[8] ~= "1" then', + " redis.call('DEL', KEYS[6])", + 'end', 'return 1', ].join('\n'); const accepted = Number( @@ -1315,6 +1546,7 @@ export class RedisBridgeStore { identityId ?? '', hasWorkspace ? '1' : '0', settlement.incarnationId, + assignment.workspaceLeaseSlot === undefined ? '0' : '1', ), signal, 'Bridge settlement commit', @@ -1350,6 +1582,15 @@ 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.commitPendingWorkspace(assignment, settlement); + await this.cleanupWithRetry(workerId, assignmentId, assignment); + } } async cancelled( @@ -1419,6 +1660,126 @@ 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. + 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", + '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')", + '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', + ].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)), + quarantine ? '1' : '0', + ), + 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, @@ -1432,12 +1793,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, @@ -1543,11 +1906,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 +1918,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( @@ -1586,19 +1949,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', ); @@ -1610,33 +1968,63 @@ 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", + '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 = [ 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)!, ), ); } + 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, @@ -1648,6 +2036,7 @@ export class RedisBridgeStore { String(ttlSeconds * 1000), String(Date.parse(assignment.expiresAt)), readyToken ?? '', + JSON.stringify(receipt), ); if (Number(result) === -1) { throw new BridgeStoreError( @@ -1685,7 +2074,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), @@ -1703,16 +2092,45 @@ 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 ( - 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 +2139,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 @@ -1744,7 +2159,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) { @@ -1759,17 +2174,29 @@ export class RedisBridgeStore { throw lastError; } - private async cleanup(assignment: StoredAssignment): Promise { + private async cleanup(assignment: AssignmentOwnership): 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 = [ @@ -1782,6 +2209,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', @@ -1798,7 +2226,8 @@ export class RedisBridgeStore { keys.length, ...keys, assignment.assignmentId, - assignment.runtimeSessionId === undefined ? '0' : '1', + assignmentWorkspace(assignment) === undefined ? '0' : '1', + assignment.workspaceLeaseSlot === undefined ? '0' : '1', ), this.redisCommandTimeoutMs, 'Bridge assignment cleanup', @@ -1806,7 +2235,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', ); diff --git a/service/src/config.ts b/service/src/config.ts index 5384810..94daf5a 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -397,6 +397,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. */ diff --git a/tests/compose-bridge-config.cjs b/tests/compose-bridge-config.cjs index d6c9797..de14b55 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 ?? ''); }