diff --git a/api/README.md b/api/README.md index bbbe36ec..38e4cd98 100644 --- a/api/README.md +++ b/api/README.md @@ -94,6 +94,24 @@ Other package-format-compatible runtimes (Go, Rust, Java, GCC) can be installed Execute code in a sandboxed environment. +When a persisted input file is removed during execution, a complete artifact +scan reports its relative path in `deleted_files`. Callers can use this +explicit list to remove stale file references from their next session request. +The field is omitted when no persisted inputs were removed or when artifact +scanning is incomplete, so truncation or unreadable paths cannot be mistaken +for deletions. + +When supported output files are omitted because the response reaches its file +count limit, nesting or path limits, file-size limit, or a filesystem entry +cannot be read, the response includes `artifact_truncation`. Its `reasons` +object counts detected omissions by cause, `skipped_count` reports the total +detected omissions, and `skipped` contains up to 20 relative paths so callers +can match an expected output. Intentional filters such as unsupported file +extensions, hidden runtime directories, and unchanged session files do not +produce this marker when they can be classified within the bounded scan. A +depth-capped subtree that exceeds the metadata probe budget is reported +conservatively rather than allowing post-execution traversal to run unbounded. + ### `GET /api/v2/runtimes` List available language runtimes. diff --git a/api/src/job.ts b/api/src/job.ts index b7b82148..0939de08 100644 --- a/api/src/job.ts +++ b/api/src/job.ts @@ -36,9 +36,9 @@ import { SANDBOX_DIR_MODE, SANDBOX_FILE_MODE, ValidationError, + checkPathShape, hasRunnableSource, isDirkeep, - isValidPathShape, validateFilePath, isValidFilePath, } from './validation'; @@ -60,6 +60,15 @@ export { const AUTO_LOAD_DIRKEEP_TIMEOUT_MS = 10000; const AUTO_LOAD_DIRKEEP_RETRIES = 2; +const PTC_HISTORY_FILENAME = '_ptc_history.json'; +const TRUNCATION_PROBE_MAX_ENTRIES = 1000; +const TRUNCATION_PROBE_MAX_LEVELS = 10; +const TRUNCATION_PROBE_MAX_HASH_BYTES = 50_000_000; + +interface TruncationProbeState { + remainingEntries: number; + remainingHashBytes: number; +} /** Replaying the same sealed grant cannot repair an authorization denial. */ class InputAuthorizationError extends Error { @@ -646,9 +655,23 @@ interface ExecuteResult { /** Top-level execution session id (one sandbox `/exec` invocation). */ session_id: string; files: FileRef[]; + /** Persisted input paths that no longer exist after this execution. */ + deleted_files?: string[]; artifact_delivery?: ArtifactDeliveryFailure; + artifact_truncation?: ArtifactTruncation; +} + +export type ArtifactTruncationReason = 'max_files' | 'depth' | 'size' | 'path' | 'unreadable'; + +export interface ArtifactTruncation { + code: 'artifact_truncated'; + reasons: Partial>; + skipped: string[]; + skipped_count: number; } +const MAX_REPORTED_TRUNCATED_PATHS = 20; + const jobQueue: Array<() => void> = []; async function acquireJobIdentity(log: Logger): Promise { @@ -693,6 +716,13 @@ export class Job { private pendingSurfaced = new Map(); private sessionFiles: FileRef[] = []; private inheritedRefs: FileRef[] = []; + private presentInputFiles = new Set(); + private deletedFiles: string[] = []; + private artifactTruncation: ArtifactTruncation | undefined; + private truncationProbeState: TruncationProbeState = { + remainingEntries: TRUNCATION_PROBE_MAX_ENTRIES, + remainingHashBytes: TRUNCATION_PROBE_MAX_HASH_BYTES, + }; private inputFileHashes = new Map(); private inputManifest = new Map(); private inputDestinations = new Map(); @@ -1673,6 +1703,10 @@ export class Job { version: this.runtime.version.raw, session_id: this.outputSessionId, files: this.sessionFiles, + ...(this.deletedFiles.length > 0 + ? { deleted_files: this.deletedFiles } + : {}), + ...(this.artifactTruncation ? { artifact_truncation: this.artifactTruncation } : {}), }; } @@ -1680,6 +1714,13 @@ export class Job { this.generatedFiles = []; this.sessionFiles = []; this.inheritedRefs = []; + this.presentInputFiles.clear(); + this.deletedFiles = []; + this.artifactTruncation = undefined; + this.truncationProbeState = { + remainingEntries: TRUNCATION_PROBE_MAX_ENTRIES, + remainingHashBytes: TRUNCATION_PROBE_MAX_HASH_BYTES, + }; const inputByName = new Map(); for (const f of this.files) inputByName.set(f.name, f); @@ -1688,6 +1729,26 @@ export class Job { await this.walkDir(this.submissionDir, 0, inputByName); } catch (error) { this.log.error({ err: error }, 'Error scanning submission directory'); + this.recordArtifactTruncation('unreadable', '.'); + } + + if (this.artifactTruncation == null) { + const returnedNames = new Set([ + ...this.sessionFiles.map(file => file.name), + ...this.inheritedRefs.map(file => file.name), + ]); + for (const file of this.files) { + if ( + file.id != null && + file.storage_session_id != null && + this.inputFileHashes.get(file.name)?.readOnly !== true && + !this.presentInputFiles.has(file.name) && + !returnedNames.has(file.name) + ) { + this.deletedFiles.push(file.name); + this.session?.forgetPrimed(file.name); + } + } } /* Generated files get priority in sessionFiles; fill remaining slots up @@ -1699,6 +1760,23 @@ export class Job { if (remaining > 0 && this.inheritedRefs.length > 0) { this.sessionFiles.push(...this.inheritedRefs.slice(0, remaining)); } + for (const ref of this.inheritedRefs.slice(remaining)) { + this.recordArtifactTruncation('max_files', ref.name); + } + } + + private recordArtifactTruncation(reason: ArtifactTruncationReason, relativePath: string): void { + this.artifactTruncation ??= { + code: 'artifact_truncated', + reasons: {}, + skipped: [], + skipped_count: 0, + }; + this.artifactTruncation.reasons[reason] = (this.artifactTruncation.reasons[reason] ?? 0) + 1; + this.artifactTruncation.skipped_count++; + if (this.artifactTruncation.skipped.length < MAX_REPORTED_TRUNCATED_PATHS) { + this.artifactTruncation.skipped.push(relativePath); + } } /** @@ -1722,6 +1800,7 @@ export class Job { isRegularFile = st.isFile(); } catch (err) { this.log.debug({ path: relativePath, err }, 'walkDir: failed to lstat entry'); + this.recordArtifactTruncation('unreadable', relativePath); return 'skip'; } } @@ -1742,7 +1821,14 @@ export class Job { inputByName: Map, ): Promise<{ collected: boolean; truncated: boolean }> { const keepPath = path.join(relativePath, DIRKEEP); - if (!isValidPathShape(keepPath)) return { collected: false, truncated: false }; + const pathShapeError = checkPathShape(keepPath); + if (pathShapeError) { + this.recordArtifactTruncation( + pathShapeError.includes('nesting depth') ? 'depth' : 'path', + keepPath, + ); + return { collected: false, truncated: true }; + } const keepFullPath = path.join(fullPath, DIRKEEP); const inheritedKeep = inputByName.get(keepPath); @@ -1770,6 +1856,7 @@ export class Job { return this.createDirkeepMarker(keepPath, keepFullPath); } if (this.generatedFiles.length >= config.max_output_files) { + this.recordArtifactTruncation('max_files', keepPath); return { collected: false, truncated: true }; } const id = nanoid(); @@ -1819,6 +1906,7 @@ export class Job { if (!keepModified || keepInfo?.readOnly === true) return this.echoInheritedKeep(keepPath, inheritedKeep); if (this.generatedFiles.length >= config.max_output_files) { + this.recordArtifactTruncation('max_files', keepPath); return { collected: false, truncated: true }; } const refreshedId = nanoid(); @@ -1860,6 +1948,7 @@ export class Job { inheritedKeep: TFile, ): { collected: boolean; truncated: boolean } { if (this.inheritedRefs.length >= config.max_output_files) { + this.recordArtifactTruncation('max_files', keepPath); return { collected: false, truncated: true }; } this.inheritedRefs.push({ @@ -1886,6 +1975,7 @@ export class Job { keepFullPath: string, ): Promise<{ collected: boolean; truncated: boolean }> { if (this.generatedFiles.length >= config.max_output_files) { + this.recordArtifactTruncation('max_files', keepPath); return { collected: false, truncated: true }; } try { @@ -1944,6 +2034,7 @@ export class Job { if (existingFile.id && existingFile.storage_session_id) { if (this.inheritedRefs.length >= config.max_output_files) { + this.recordArtifactTruncation('max_files', relativePath); return { collected: false, truncated: true }; } this.inheritedRefs.push({ @@ -2003,9 +2094,29 @@ export class Job { size = st.size; } catch (err) { this.log.debug({ path: relativePath, err }, 'walkDir: unable to stat file'); + this.recordArtifactTruncation('unreadable', relativePath); return { collected: false, truncated: false, stopLoop: false }; } + + const inputFileInfo = this.inputFileHashes.get(relativePath); + const existingFile = inputByName.get(relativePath); if (size > this.runtime.max_file_size) { + /* Only an inline entrypoint needs hashing to decide whether this is + * intentional request-input suppression. Every other oversized file + * is rejected immediately, preserving the scan's bounded I/O cost. */ + if (!inputFileInfo || existingFile?.id != null || relativePath !== this.entryPointName) { + this.recordArtifactTruncation('size', relativePath); + return { collected: false, truncated: false, stopLoop: false }; + } + try { + const currentHash = await this.computeFileHash(fullPath, true); + if (currentHash === inputFileInfo.hash) { + return { collected: true, truncated: false, stopLoop: false }; + } + } catch (err) { + this.log.debug({ path: relativePath, err }, 'walkDir: failed to hash oversized entrypoint'); + } + this.recordArtifactTruncation('size', relativePath); return { collected: false, truncated: false, stopLoop: false }; } @@ -2015,8 +2126,6 @@ export class Job { * stat-only signature would wrongly suppress. Compute once per session/input * file and reuse for the suppression check, wasModified, and the surfaced * mark; non-session jobs still only hash their inputs. */ - const inputFileInfo = this.inputFileHashes.get(relativePath); - const existingFile = inputByName.get(relativePath); let contentHash: string | undefined; if (inputFileInfo != null || this.session != null) { try { @@ -2058,6 +2167,15 @@ export class Job { if (wasModified) this.log.info({ file: relativePath }, 'Input file was modified'); } + /* The unchanged inline entrypoint is executable request input, not an + * output artifact. Suppress it before applying output-size reporting; + * downloaded inputs still flow through the size limit below, preserving + * the existing response-cap behavior for inherited refs. */ + if (!wasModified && inputFileInfo && existingFile?.id == null + && relativePath === this.entryPointName) { + return { collected: true, truncated: false, stopLoop: false }; + } + const echoed = this.tryEchoUnchangedInput({ wasModified, inputFileInfo, @@ -2067,6 +2185,7 @@ export class Job { if (echoed) return { ...echoed, stopLoop: false }; if (this.generatedFiles.length >= config.max_output_files) { + this.recordArtifactTruncation('max_files', relativePath); return { collected: false, truncated: true, stopLoop: true }; } @@ -2105,10 +2224,170 @@ export class Job { const childStatus = await this.walkDir(fullPath, parentDepth + 1, inputByName); if (childStatus === 'collected') return { collected: true, truncated: false }; if (childStatus === 'skipped') return { collected: false, truncated: true }; - if (this.isOutputCapFull()) return { collected: false, truncated: true }; return this.handleEmptyDirectory(relativePath, fullPath, inputByName); } + /** Finds the first artifact that a scan cap would hide without reading file + * contents. Files below a depth boundary cannot be valid primed inputs, and + * symlinks/unsupported files/hidden runtime directories remain intentional + * exclusions. An empty directory represents a reportable `.dirkeep`. */ + private async findTruncatedArtifact( + dir: string, + inputByName: Map, + state = this.truncationProbeState, + probeDepth = 0, + rootPath = path.relative(this.submissionDir, dir) || '.', + isOutputCapProbe = false, + ): Promise { + /* The state is shared by every probe in this job. Once exhausted, return + * conservatively before opening yet another capped sibling directory. */ + if (state.remainingEntries <= 0) return rootPath; + let directory: fs.Dir; + try { + directory = await fsp.opendir(dir); + } catch (err) { + const relativeDir = path.relative(this.submissionDir, dir) || '.'; + this.log.debug({ dir, err }, 'walkDir: unable to inspect depth-capped directory'); + this.recordArtifactTruncation('unreadable', relativeDir); + return undefined; + } + + let sawVisibleEntry = false; + let sawVisibleNonHiddenEntry = false; + try { + for await (const entry of directory) { + const fullPath = path.join(dir, entry.name); + const relativePath = path.relative(this.submissionDir, fullPath); + const kind = await this.classifyDirent(entry, fullPath, relativePath); + if (kind === 'file' && entry.name === PTC_HISTORY_FILENAME) { + if (inputByName.has(relativePath)) { + this.presentInputFiles.add(relativePath); + } + continue; + } + sawVisibleEntry = true; + state.remainingEntries--; + if (state.remainingEntries < 0) return rootPath; + if (kind === 'skip') { + /* Ordinary walking counts symlinks/special entries as non-empty even + * though it does not surface them, so the probe must not invent a + * parent .dirkeep for that shape. */ + sawVisibleNonHiddenEntry = true; + continue; + } + if (kind === 'file') { + sawVisibleNonHiddenEntry = true; + if (inputByName.has(relativePath)) { + this.presentInputFiles.add(relativePath); + } + if (entry.name !== DIRKEEP && !isSupportedOutputFilename(entry.name)) continue; + const existingFile = inputByName.get(relativePath); + const inputFileInfo = this.inputFileHashes.get(relativePath); + let capProbeStat: fs.Stats | undefined; + if (isOutputCapProbe) { + const pathShapeError = checkPathShape(relativePath); + if (pathShapeError) { + this.recordArtifactTruncation( + pathShapeError.includes('nesting depth') ? 'depth' : 'path', + relativePath, + ); + continue; + } + try { + capProbeStat = await fsp.lstat(fullPath); + if (!capProbeStat.isFile()) continue; + } catch (err) { + this.log.debug({ path: relativePath, err }, 'walkDir: failed during cap-probe stat'); + this.recordArtifactTruncation('unreadable', relativePath); + continue; + } + if (capProbeStat.size > this.runtime.max_file_size) { + /* Match handleRegularFile's one exception: an unchanged inline + * entrypoint is request input rather than an oversized output. */ + if (!inputFileInfo || existingFile?.id != null || relativePath !== this.entryPointName) { + this.recordArtifactTruncation('size', relativePath); + continue; + } + if (capProbeStat.size > state.remainingHashBytes) return rootPath; + state.remainingHashBytes -= capProbeStat.size; + try { + if (await this.computeFileHash(fullPath, true) === inputFileInfo.hash) continue; + } catch (err) { + this.log.debug({ path: relativePath, err }, 'walkDir: failed during oversized entrypoint cap probe'); + } + this.recordArtifactTruncation('size', relativePath); + continue; + } + } + if ( + isOutputCapProbe + && relativePath === this.entryPointName + && existingFile?.id == null + && inputFileInfo + ) { + try { + if (capProbeStat!.size > state.remainingHashBytes) return rootPath; + state.remainingHashBytes -= capProbeStat!.size; + if (await this.computeFileHash(fullPath, true) === inputFileInfo.hash) continue; + } catch (err) { + this.log.debug({ path: relativePath, err }, 'walkDir: failed during entrypoint cap probe'); + this.recordArtifactTruncation('unreadable', relativePath); + continue; + } + } + /* Once generated outputs fill the response cap, a persistent + * workspace may still contain unchanged artifacts from earlier + * turns. Ordinary walking suppresses those via their content hash, + * so the bounded cap probe must do the same or it reports a false + * max_files warning. Current-request inputs remain reportable: they + * would otherwise have been echoed into this response. */ + if (isOutputCapProbe && this.session && !existingFile) { + if (this.session.isPrimedReadOnly(relativePath)) continue; + try { + if (capProbeStat!.size > state.remainingHashBytes) return rootPath; + state.remainingHashBytes -= capProbeStat!.size; + const hash = await this.computeFileHash(fullPath, true); + if (this.session.isSurfaced(relativePath, hash)) continue; + if ( + this.session.isPrimedInput(relativePath) + && this.session.primedHash(relativePath) === hash + ) continue; + } catch (err) { + this.log.debug({ path: relativePath, err }, 'walkDir: failed during cap-probe hashing'); + this.recordArtifactTruncation('unreadable', relativePath); + continue; + } + } + return relativePath; + } + if (isHiddenDirectory(entry.name) && !inputsLiveUnder(inputByName, relativePath)) continue; + sawVisibleNonHiddenEntry = true; + /* The probe exists only to avoid false warnings for small, obviously + * unsupported-only subtrees. Once either budget is exhausted, report + * the capped root conservatively instead of defeating the scan bound. */ + if (probeDepth >= TRUNCATION_PROBE_MAX_LEVELS) return rootPath; + const nested = await this.findTruncatedArtifact( + fullPath, + inputByName, + state, + probeDepth + 1, + rootPath, + isOutputCapProbe, + ); + if (nested) return nested; + } + } catch (err) { + const relativeDir = path.relative(this.submissionDir, dir) || '.'; + this.log.debug({ dir, err }, 'walkDir: failed during bounded directory inspection'); + this.recordArtifactTruncation('unreadable', relativeDir); + return undefined; + } + + return sawVisibleEntry && sawVisibleNonHiddenEntry + ? undefined + : path.join(path.relative(this.submissionDir, dir), DIRKEEP); + } + /** * Recursively scans the submission directory for output files. Returns a * status distinguishing truly empty directories from scans truncated by @@ -2120,14 +2399,30 @@ export class Job { depth: number, inputByName: Map, ): Promise<'collected' | 'empty' | 'skipped'> { - if (depth >= config.max_nesting_depth) return 'skipped'; - if (this.isOutputCapFull()) return 'skipped'; - + const relativeDir = path.relative(this.submissionDir, dir) || '.'; + if (depth >= config.max_nesting_depth) { + const skippedPath = await this.findTruncatedArtifact(dir, inputByName); + if (skippedPath) this.recordArtifactTruncation('depth', skippedPath); + return 'skipped'; + } + if (this.isOutputCapFull()) { + const skippedPath = await this.findTruncatedArtifact( + dir, + inputByName, + this.truncationProbeState, + 0, + relativeDir, + true, + ); + if (skippedPath) this.recordArtifactTruncation('max_files', skippedPath); + return 'skipped'; + } let entries: fs.Dirent[]; try { entries = await fsp.readdir(dir, { withFileTypes: true }); } catch (err) { this.log.debug({ dir, err }, 'walkDir: unable to read directory'); + this.recordArtifactTruncation('unreadable', relativeDir); return 'skipped'; } @@ -2146,7 +2441,6 @@ export class Job { * separate npm packages so we can't import directly; the filename literal * is asserted-equal in `service/scripts/test-ptc-sentinel.ts` to catch * accidental drift in CI. */ - const PTC_HISTORY_FILENAME = '_ptc_history.json'; const isPtcReserved = (name: string): boolean => name === PTC_HISTORY_FILENAME; const nonDirkeepCount = entries.reduce( @@ -2166,16 +2460,21 @@ export class Job { let skippedHiddenDirs = 0; for (const entry of entries) { - if (this.isOutputCapFull()) { truncated = true; break; } - if (isPtcReserved(entry.name)) continue; - const fullPath = path.join(dir, entry.name); const relativePath = path.relative(this.submissionDir, fullPath); - if (!isValidPathShape(relativePath)) continue; - const kind = await this.classifyDirent(entry, fullPath, relativePath); if (kind === 'skip') continue; + if (kind === 'file' && inputByName.has(relativePath)) { + this.presentInputFiles.add(relativePath); + } + + /* A by-reference input may legitimately use the reserved replay-history + * basename on the ordinary execution endpoint. It remains hidden from + * output collection, but must be observed before the runtime fixture is + * skipped so an untouched input is not reported as deleted. */ + if (kind === 'file' && isPtcReserved(entry.name)) continue; + if (kind === 'dir') { /* Skip hidden directories (basename starts with `.`) unless the user * explicitly primed something under them. Matplotlib, pip, and other @@ -2189,9 +2488,42 @@ export class Job { skippedHiddenDirs++; continue; } + const pathShapeError = checkPathShape(relativePath); + if (pathShapeError) { + const skippedPath = await this.findTruncatedArtifact( + fullPath, + inputByName, + this.truncationProbeState, + 0, + relativePath, + ); + if (skippedPath) { + this.recordArtifactTruncation( + pathShapeError.includes('nesting depth') ? 'depth' : 'path', + skippedPath, + ); + truncated = true; + } + continue; + } const res = await this.walkSubdirectory(relativePath, fullPath, depth, inputByName); if (res.collected) hasCollectedChild = true; if (res.truncated) truncated = true; + if (this.isOutputCapFull() && this.artifactTruncation?.reasons.max_files) break; + continue; + } + + /* Check intentional filename filtering before path limits. Unsupported + * files never belong in files[], regardless of how long their path is. */ + if (entry.name !== DIRKEEP && !isSupportedOutputFilename(entry.name)) continue; + + const pathShapeError = checkPathShape(relativePath); + if (pathShapeError) { + this.recordArtifactTruncation( + pathShapeError.includes('nesting depth') ? 'depth' : 'path', + relativePath, + ); + truncated = true; continue; } diff --git a/api/src/session-workspace.test.ts b/api/src/session-workspace.test.ts index 87ca71df..1a5320c6 100644 --- a/api/src/session-workspace.test.ts +++ b/api/src/session-workspace.test.ts @@ -113,6 +113,11 @@ describe('SessionWorkspace state', () => { expect(ws.isPrimedInput('in.csv')).toBe(false); ws.markPrimed('in.csv', 'file_abc'); expect(ws.primedInputId('in.csv')).toBe('file_abc'); + ws.markSurfaced('in.csv', 'old-output'); + ws.forgetPrimed('in.csv'); + expect(ws.primedInputId('in.csv')).toBeUndefined(); + expect(ws.isSurfaced('in.csv', 'old-output')).toBe(false); + ws.markPrimed('in.csv', 'file_abc'); /* read-only primes report as not-primed so the caller re-downloads them * (a reused on-disk copy could have been tampered via the writable dir). */ diff --git a/api/src/session-workspace.ts b/api/src/session-workspace.ts index 3f2105bc..82ed6c89 100644 --- a/api/src/session-workspace.ts +++ b/api/src/session-workspace.ts @@ -202,6 +202,12 @@ export class SessionWorkspace { this.primed.set(relPath, { id: storageFileId, readOnly, hash }); } + /** Clears input lineage after execution proves that the path was deleted. */ + forgetPrimed(relPath: string): void { + this.primed.delete(relPath); + this.forget(relPath); + } + markDirty(reason: string): void { this.dirty = reason; logger.error( diff --git a/api/src/walker.test.ts b/api/src/walker.test.ts index 91382468..44e252a6 100644 --- a/api/src/walker.test.ts +++ b/api/src/walker.test.ts @@ -26,11 +26,29 @@ interface WalkerInternals { generatedFiles: Array<{ id: string; name: string; path: string }>; sessionFiles: Array<{ id: string; name: string; storage_session_id: string; modified_from?: { id: string; storage_session_id: string }; inherited?: true; entity_id?: string }>; inheritedRefs: Array<{ id: string; name: string; storage_session_id: string; inherited?: true; entity_id?: string }>; + presentInputFiles: Set; + deletedFiles: string[]; + artifactTruncation?: { + code: 'artifact_truncated'; + reasons: Partial>; + skipped: string[]; + skipped_count: number; + }; + truncationProbeState: { remainingEntries: number; remainingHashBytes: number }; pendingSurfaced: Map; inputFileHashes: Map; files: TFile[]; reusePrimedInput: (file: TFile) => Promise; writeFile: (file: TFile) => Promise; + computeFileHash: (filePath: string, noFollow?: boolean) => Promise; + findTruncatedArtifact: ( + dir: string, + inputByName: Map, + state?: { remainingEntries: number; remainingHashBytes: number }, + probeDepth?: number, + rootPath?: string, + respectSessionSuppression?: boolean, + ) => Promise; walkDir: (dir: string, depth: number, inputByName: Map) => Promise<'collected' | 'empty' | 'skipped'>; handleSessionFiles: () => Promise; } @@ -743,6 +761,10 @@ describe('walkDir / output caps', () => { await internals.walkDir(tmpDir, 0, new Map()); expect(internals.generatedFiles.length).toBeLessThanOrEqual(cap); + expect(internals.artifactTruncation).toMatchObject({ + code: 'artifact_truncated', + reasons: { max_files: 1 }, + }); }); it('respects max_output_files cap on inherited refs', async () => { @@ -771,6 +793,11 @@ describe('walkDir / output caps', () => { expect(internals.inheritedRefs.length).toBeLessThanOrEqual(cap); expect(internals.generatedFiles).toHaveLength(0); + expect(internals.artifactTruncation).toMatchObject({ + code: 'artifact_truncated', + reasons: { max_files: 5 }, + skipped_count: 5, + }); }); }); @@ -792,6 +819,407 @@ describe('walkDir / depth cap', () => { const deepName = path.relative(tmpDir, path.join(cursor, 'deep.py')); expect(internals.generatedFiles.map(f => f.name)).not.toContain(deepName); + expect(internals.artifactTruncation).toMatchObject({ + code: 'artifact_truncated', + reasons: { depth: 1 }, + skipped_count: 1, + }); + expect(internals.artifactTruncation?.skipped[0]).toBe(deepName); + }); + + it('does not report a depth cap when the skipped subtree has only unsupported files', async () => { + let cursor = tmpDir; + for (let i = 0; i < config.max_nesting_depth; i++) { + cursor = path.join(cursor, `d${i}`); + await fsp.mkdir(cursor); + } + await fsp.writeFile(path.join(cursor, 'cache.bin'), 'ignored'); + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toBeUndefined(); + }); + + it('bounds depth-cap eligibility probes and reports the capped subtree conservatively', async () => { + let cursor = tmpDir; + const totalDepth = config.max_nesting_depth + 12; + for (let i = 0; i < totalDepth; i++) { + cursor = path.join(cursor, `d${i}`); + await fsp.mkdir(cursor); + } + await fsp.writeFile(path.join(cursor, 'cache.bin'), 'ignored'); + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + const cappedRoot = Array.from( + { length: config.max_nesting_depth }, + (_, i) => `d${i}`, + ).join(path.sep); + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { depth: 1 }, + skipped: [cappedRoot], + skipped_count: 1, + }); + }); +}); + +describe('walkDir / artifact truncation details', () => { + it('reports oversized supported outputs while leaving them out of files', async () => { + await fsp.writeFile(path.join(tmpDir, 'large.txt'), 'too large'); + const job = makeJob({ maxFileSize: 3 }); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.generatedFiles).toHaveLength(0); + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { size: 1 }, + skipped: ['large.txt'], + skipped_count: 1, + }); + }); + + it('reports overlong output paths', async () => { + const directory = 'a'.repeat(200); + await fsp.mkdir(path.join(tmpDir, directory)); + const name = path.join(directory, `${'b'.repeat(60)}.txt`); + await fsp.writeFile(path.join(tmpDir, name), 'content'); + const job = makeJob(); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { path: 1 }, + skipped: [name], + skipped_count: 1, + }); + }); + + it('does not report an overlong path for an unsupported output', async () => { + const directory = 'a'.repeat(200); + await fsp.mkdir(path.join(tmpDir, directory)); + await fsp.writeFile(path.join(tmpDir, directory, `${'b'.repeat(60)}.bin`), 'ignored'); + const job = makeJob(); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.generatedFiles).toHaveLength(0); + expect(internals.artifactTruncation).toBeUndefined(); + }); + + it('reports an empty-directory marker whose appended path is too long', async () => { + const directory = 'a'.repeat(config.max_path_length - 6); + await fsp.mkdir(path.join(tmpDir, directory)); + const job = makeJob(); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { path: 1 }, + skipped: [path.join(directory, DIRKEEP)], + skipped_count: 1, + }); + }); + + it('reports an explicit overlong .dirkeep exactly once', async () => { + const directory = 'a'.repeat(config.max_path_length - 6); + await fsp.mkdir(path.join(tmpDir, directory)); + const keepName = path.join(directory, DIRKEEP); + await fsp.writeFile(path.join(tmpDir, keepName), ''); + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { path: 1 }, + skipped: [keepName], + skipped_count: 1, + }); + }); + + it('uses a bounded probe instead of recursively walking an overlong directory', async () => { + const first = 'a'.repeat(200); + const second = 'b'.repeat(60); + const overlongDir = path.join(first, second); + await fsp.mkdir(path.join(tmpDir, overlongDir), { recursive: true }); + for (let i = 0; i < 1001; i++) { + await fsp.writeFile(path.join(tmpDir, overlongDir, `ignored-${i}.bin`), 'ignored'); + } + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { path: 1 }, + skipped: [overlongDir], + skipped_count: 1, + }); + }); + + it('does not report an unchanged oversized inline entrypoint', async () => { + const name = 'main.py'; + const content = 'print(1)'; + const full = path.join(tmpDir, name); + await fsp.writeFile(full, content); + const inline: TFile = { name, content }; + const job = makeJob({ files: [inline], maxFileSize: 3 }); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + internals.entryPointName = name; + internals.inputFileHashes.set(name, { hash: sha256(content), path: full }); + + await internals.walkDir(tmpDir, 0, buildInputByName([inline])); + + expect(internals.generatedFiles).toHaveLength(0); + expect(internals.artifactTruncation).toBeUndefined(); + }); + + it('inspects a capped directory before deciding whether an artifact was omitted', async () => { + const job = makeJob(); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + await fsp.mkdir(path.join(tmpDir, 'ignored')); + await fsp.writeFile(path.join(tmpDir, 'ignored', 'cache.bin'), 'ignored'); + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toBeUndefined(); + }); + + it('reports a capped empty-directory marker', async () => { + const job = makeJob(); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + await fsp.mkdir(path.join(tmpDir, 'empty')); + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { max_files: 1 }, + skipped: [path.join('empty', DIRKEEP)], + skipped_count: 1, + }); + }); + + it('classifies an oversized supported file by size when the output cap is full', async () => { + await fsp.writeFile(path.join(tmpDir, 'oversized.txt'), 'too large'); + const internals = asInternals(makeJob({ maxFileSize: 3 })); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { size: 1 }, + skipped: ['oversized.txt'], + skipped_count: 1, + }); + }); + + it('classifies an overlong supported path by path when the output cap is full', async () => { + const directory = 'a'.repeat(200); + const filename = path.join(directory, `${'b'.repeat(60)}.txt`); + await fsp.mkdir(path.join(tmpDir, directory)); + await fsp.writeFile(path.join(tmpDir, filename), 'output'); + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { path: 1 }, + skipped: [filename], + skipped_count: 1, + }); + }); + + it('does not hash ordinary oversized files in session mode', async () => { + await fsp.writeFile(path.join(tmpDir, 'large.txt'), 'too large'); + const session = new SessionWorkspace({ runtimeSessionId: 'rt_large' }); + const internals = asInternals(makeJob({ maxFileSize: 3, session })); + internals.submissionDir = tmpDir; + let hashCalls = 0; + internals.computeFileHash = async () => { + hashCalls++; + return sha256('too large'); + }; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(hashCalls).toBe(0); + expect(internals.artifactTruncation?.reasons).toEqual({ size: 1 }); + }); + + it('keeps scanning for generated outputs when only inherited refs are capped', async () => { + await fsp.mkdir(path.join(tmpDir, 'a-ignored')); + await fsp.writeFile(path.join(tmpDir, 'a-ignored', 'cache.bin'), 'ignored'); + await fsp.writeFile(path.join(tmpDir, 'z-generated.txt'), 'new'); + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + internals.inheritedRefs = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `inherited-${i}.txt`, + storage_session_id: 'previous', + inherited: true, + })); + internals.artifactTruncation = { + code: 'artifact_truncated', + reasons: { max_files: 1 }, + skipped: ['another-inherited.txt'], + skipped_count: 1, + }; + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.generatedFiles.map(file => file.name)).toContain('z-generated.txt'); + }); + + it('bounds output-cap eligibility probes for wide unsupported-only directories', async () => { + const job = makeJob(); + const internals = asInternals(job); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + for (let i = 0; i < 1001; i++) { + await fsp.writeFile(path.join(tmpDir, `ignored-${i}.bin`), 'ignored'); + } + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { max_files: 1 }, + skipped: ['.'], + skipped_count: 1, + }); + }); + + it('shares the output-cap probe budget across sibling subtrees', async () => { + for (const dirname of ['b-ignored', 'c-ignored']) { + await fsp.mkdir(path.join(tmpDir, dirname)); + for (let i = 0; i < 600; i++) { + await fsp.writeFile(path.join(tmpDir, dirname, `ignored-${i}.bin`), 'ignored'); + } + } + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + + await internals.walkDir(path.join(tmpDir, 'b-ignored'), 1, new Map()); + await internals.walkDir(path.join(tmpDir, 'c-ignored'), 1, new Map()); + + expect(internals.artifactTruncation?.reasons).toEqual({ max_files: 1 }); + expect(internals.artifactTruncation?.skipped).toEqual(['c-ignored']); + }); + + it('does not report surfaced session artifacts during output-cap probing', async () => { + const name = 'old-output.txt'; + const content = 'already returned'; + await fsp.writeFile(path.join(tmpDir, name), content); + const session = new SessionWorkspace({ runtimeSessionId: 'rt_capped' }); + session.markSurfaced(name, sha256(content)); + const internals = asInternals(makeJob({ session })); + internals.submissionDir = tmpDir; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + + await internals.walkDir(tmpDir, 0, new Map()); + + expect(internals.artifactTruncation).toBeUndefined(); + }); + + it('does not reopen capped directories after the shared probe budget is exhausted', async () => { + const internals = asInternals(makeJob()); + internals.submissionDir = tmpDir; + internals.truncationProbeState.remainingEntries = 0; + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + const absentDir = path.join(tmpDir, 'not-opened'); + + await internals.walkDir(absentDir, 1, new Map()); + + expect(internals.artifactTruncation).toEqual({ + code: 'artifact_truncated', + reasons: { max_files: 1 }, + skipped: ['not-opened'], + skipped_count: 1, + }); + }); + + it('does not report an unchanged inline entrypoint during output-cap probing', async () => { + const directory = path.join(tmpDir, 'src'); + const name = path.join('src', 'main.py'); + const content = 'print(1)'; + await fsp.mkdir(directory); + await fsp.writeFile(path.join(tmpDir, name), content); + const inline: TFile = { name, content }; + const internals = asInternals(makeJob({ files: [inline] })); + internals.submissionDir = tmpDir; + internals.entryPointName = name; + internals.inputFileHashes.set(name, { hash: sha256(content), path: path.join(tmpDir, name) }); + internals.generatedFiles = Array.from({ length: config.max_output_files }, (_, i) => ({ + id: `id-${i}`, + name: `file-${i}.txt`, + path: path.join(tmpDir, `file-${i}.txt`), + })); + + await internals.walkDir(directory, 1, buildInputByName([inline])); + + expect(internals.artifactTruncation).toBeUndefined(); }); }); @@ -882,6 +1310,192 @@ describe('handleSessionFiles / priority-fill composition', () => { }); }); +describe('handleSessionFiles / persisted input deletion', () => { + it('reports a persisted input that no longer exists', async () => { + const inherited: TFile = { + id: 'prior-id', + storage_session_id: 'prior-session', + name: 'removed.txt', + }; + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.deletedFiles).toEqual(['removed.txt']); + }); + + it('retains a read-only persisted input when sandbox code removes its local copy', async () => { + const inherited: TFile = { + id: 'skill-id', + storage_session_id: 'skill-session', + name: path.join('skills', 'review', 'SKILL.md'), + }; + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + internals.inputFileHashes.set(inherited.name, { + hash: sha256('trusted-skill'), + path: path.join(tmpDir, inherited.name), + originalId: inherited.id, + originalSessionId: inherited.storage_session_id, + readOnly: true, + }); + + await internals.handleSessionFiles(); + + expect(internals.deletedFiles).toEqual([]); + }); + + it('tracks a persisted input using the reserved PTC history basename', async () => { + const inherited: TFile = { + id: 'history-id', + storage_session_id: 'prior-session', + name: path.join('fixtures', '_ptc_history.json'), + }; + await fsp.mkdir(path.join(tmpDir, 'fixtures')); + await fsp.writeFile(path.join(tmpDir, inherited.name), '{}'); + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.deletedFiles).toEqual([]); + expect(internals.generatedFiles.map(file => file.name)).not.toContain(inherited.name); + }); + + it('traverses a directory that uses the reserved PTC history basename', async () => { + const inherited: TFile = { + id: 'nested-id', + storage_session_id: 'prior-session', + name: path.join('_ptc_history.json', 'data.csv'), + }; + await fsp.mkdir(path.join(tmpDir, '_ptc_history.json')); + await fsp.writeFile(path.join(tmpDir, inherited.name), 'persisted'); + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.deletedFiles).toEqual([]); + expect(internals.presentInputFiles.has(inherited.name)).toBe(true); + }); + + it('tracks a reserved persisted input during a capped subtree probe', async () => { + const inherited: TFile = { + id: 'history-id', + storage_session_id: 'prior-session', + name: path.join('fixtures', '_ptc_history.json'), + }; + const fixtures = path.join(tmpDir, 'fixtures'); + await fsp.mkdir(fixtures); + await fsp.writeFile(path.join(tmpDir, inherited.name), '{}'); + await fsp.writeFile(path.join(fixtures, 'unsupported.bin'), 'binary'); + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + const skipped = await internals.findTruncatedArtifact( + fixtures, + new Map([[inherited.name, inherited]]) + ); + + expect(skipped).toBeUndefined(); + expect(internals.presentInputFiles.has(inherited.name)).toBe(true); + }); + + it('does not report a surviving input that is unsupported as an output artifact', async () => { + const inherited: TFile = { + id: 'prior-id', + storage_session_id: 'prior-session', + name: 'archive.bin', + }; + await fsp.writeFile(path.join(tmpDir, inherited.name), 'binary-placeholder'); + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.generatedFiles).toHaveLength(0); + expect(internals.deletedFiles).toEqual([]); + }); + + it('suppresses deletion reporting when the artifact scan is incomplete', async () => { + const inherited: TFile = { + id: 'prior-id', + storage_session_id: 'prior-session', + name: 'removed.txt', + }; + await fsp.writeFile(path.join(tmpDir, 'too-large.txt'), 'too large'); + const internals = asInternals(makeJob({ files: [inherited], maxFileSize: 3 })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.artifactTruncation?.reasons).toEqual({ size: 1 }); + expect(internals.deletedFiles).toEqual([]); + }); + + it('does not report an inherited marker that is returned for an empty directory', async () => { + const name = path.join('empty', DIRKEEP); + const inherited: TFile = { + id: 'marker-id', + storage_session_id: 'prior-session', + name, + }; + await fsp.mkdir(path.join(tmpDir, 'empty')); + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.deletedFiles).toEqual([]); + expect([ + ...internals.sessionFiles, + ...internals.inheritedRefs, + ].map(file => file.name)).toContain(name); + }); + + it('clears stateful priming lineage when a persisted input is deleted', async () => { + const inherited: TFile = { + id: 'prior-id', + storage_session_id: 'prior-session', + name: 'removed.txt', + }; + const session = new SessionWorkspace({ runtimeSessionId: 'rt_deleted' }); + session.markPrimed(inherited.name, inherited.id!, true, 'old-hash'); + session.markSurfaced(inherited.name, 'old-output-hash'); + const internals = asInternals(makeJob({ files: [inherited], session })); + internals.submissionDir = tmpDir; + + await internals.handleSessionFiles(); + + expect(internals.deletedFiles).toEqual([inherited.name]); + expect(session.isPrimedInput(inherited.name)).toBe(false); + expect(session.isSurfaced(inherited.name, 'old-output-hash')).toBe(false); + }); + + it('tracks surviving persisted inputs during capped subtree probes', async () => { + const inherited: TFile = { + id: 'prior-id', + storage_session_id: 'prior-session', + name: path.join('assets', 'model.bin'), + }; + await fsp.mkdir(path.join(tmpDir, 'assets')); + await fsp.writeFile( + path.join(tmpDir, inherited.name), + 'unsupported-but-persisted', + ); + const internals = asInternals(makeJob({ files: [inherited] })); + internals.submissionDir = tmpDir; + + await internals.findTruncatedArtifact( + tmpDir, + new Map([[inherited.name, inherited]]), + ); + + expect(internals.presentInputFiles.has(inherited.name)).toBe(true); + }); +}); + describe('walkDir / dirent classification', () => { it('ignores symlinks (never classifies them as file or dir)', async () => { await fsp.writeFile(path.join(tmpDir, 'real.py'), 'print(1)'); diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index ab7e2f64..b9de3ac2 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -149,6 +149,17 @@ implementation and an allowlist of workspace IDs, preserves per-workspace operation restrictions, validates bounded results, and treats an unknown command failure as an uncertain mutation. +Selected attached workspaces on macOS, Linux, and WSL2 can also advertise Bash +Programmatic Tool Calling. Native Windows workers do not advertise Bash PTC. +Code API then runs each replay iteration through the same workspace-scoped +native SRT executor. Source code operates in the selected local root, while +replay metadata, injected skills and attachments, and generated artifacts are +staged in an execution-private data directory and removed after settlement. +Only authorized file references and returned artifacts cross the relay; the +repository is never uploaded to Code API. This capability is advertised only +when native SRT commands and a file-relay upstream are both configured, so +older or partially configured workers continue to fail closed. + Native SRT is the MVP and default command backend on a user's chosen laptop or VM. It uses Seatbelt on macOS, bubblewrap/seccomp on Linux, and the SRT restricted-account helper on Windows. It confines writes to the registered @@ -157,9 +168,11 @@ credentials, and denies network egress by default. Startup fails closed when the platform dependencies are unavailable; there is no unsandboxed fallback. Use `LIBRECHAT_CODE_COMMAND_ALLOWED_DOMAINS` for an explicit comma-separated egress allowlist. -Linux hosts must provide Bash at `/bin/bash`, `bubblewrap`, `socat`, and -`ripgrep`; macOS uses system facilities. Windows requires SRT's one-time -restricted-account setup. +Linux hosts must provide `bubblewrap`, `socat`, and `ripgrep`; macOS uses +system facilities. Bash Programmatic Tool Calling additionally requires Bash +5.2 or newer and `jq` on `PATH` on macOS, Linux, and WSL2. The worker resolves +the compatible shell from `PATH` rather than assuming `/bin/bash`. Windows +requires SRT's one-time restricted-account setup. The optional `docker-nsjail` adapter enables a stronger container boundary with `--allow-workspace-commands` (or @@ -243,6 +256,23 @@ execution. the currently registered incarnation. - Request cancellation is polled by the worker and aborts the local sandbox request. +- Replay PTC clients may attach a fresh `X-LibreChat-Code-Request-ID` to each + `/exec/programmatic` request and send that same opaque ID to + `POST /v1/exec/programmatic/cancel`. Code API binds the short-lived request + record to the authenticated principal, durably marks cancellation in Redis, + and publishes it to the worker process holding the BullMQ job. This explicit + path avoids relying on HTTP connection teardown, frees waiting jobs + immediately, and interrupts active remote-bridge assignments without polling + once per active job. + Cancellation and completed-result publication use an atomic Redis decision: + a late cancel returns `already_completed` instead of acknowledging Stop after + completion won. Ambiguous enqueue/cancellation errors retain replay ownership + until a durable fence or the original job deadline. Completed results are + retained temporarily (bounded to 16 MiB) so a lost BullMQ completion reply + does not cause sandbox effects to be repeated. Reconnect reconciliation reads + only small status markers, using one subscriber per process. + Roll out the matching Code API queue-worker processes before enabling this + endpoint on API replicas; pre-cancellation workers do not observe its markers. - A leased assignment remains in a Redis-backed delivery claim until the worker explicitly acknowledges it; reconnecting before acknowledgement redelivers the same fenced assignment instead of losing it after an HTTP disconnect. diff --git a/packages/code/README.md b/packages/code/README.md index d06fbf9a..87d9dded 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -143,14 +143,48 @@ policy. This matches the personal-machine SRT trust model; use the Docker/NsJail backend or a dedicated VM boundary when hard teardown of adversarial process trees is required. -Linux hosts need Bash at `/bin/bash`, `bubblewrap`, `socat`, and `ripgrep`; macOS uses system -facilities. Follow SRT's one-time restricted-account setup when using Windows. +Linux hosts need `bubblewrap`, `socat`, and `ripgrep`; macOS uses system +facilities. Bash Programmatic Tool Calling additionally requires Bash 5.2 or +newer and `jq` on `PATH` on macOS, Linux, and WSL2. The worker resolves that +shell explicitly instead of assuming `/bin/bash`, which remains Bash 3.2 on +many macOS hosts. Follow SRT's one-time restricted-account setup when using Windows. An operator may allow explicit egress destinations with the comma-separated `LIBRECHAT_CODE_COMMAND_ALLOWED_DOMAINS` setting. Treat that as a security policy: an allowed destination can receive workspace data. The normalized allowlist is included in the worker policy digest. Tool approval hooks remain the user-facing allow/deny boundary for each invocation. +When Code API negotiates `bash` programmatic execution for a selected +workspace, the same native SRT executor also supports replay-mode Programmatic +Tool Calling on macOS, Linux, and WSL2 workers. Native Windows does not +advertise this Bash capability. The repository remains the command working directory. Generated +PTC scripts, replay history, skill files, chat attachments, and returned +artifacts use an owner-only per-execution directory under the worker's private +SRT scratch root, exposed to code as `LIBRECHAT_CODE_DATA_DIR`. That directory +is removed after every iteration and is never placed in the repository. + +Replay probes run against a disposable copy-on-write snapshot with network and +socket access denied, including under `trusted-vm`. External effects must not +repeat while discovering pending tools. Use registered tools for network-dependent +replay control flow; the final commit pass runs once under the configured policy. +Each probe's SRT proxy session is revoked before restoring the commit policy; +per-command network overrides alone do not restrict SRT's session-level proxies. +Probe failures do not quarantine the real workspace. Once the commit pass starts, +its fence remains until result restoration succeeds; uncertain finalization +quarantines only that workspace. + +Reference inputs and artifact outputs travel only through the configured +`LIBRECHAT_CODE_FILE_RELAY_UPSTREAM`, using Code API's execution-scoped opaque +egress grant. The worker rejects redirects and bounds each transfer to 10 MiB, +each execution to 100 files and 100 MiB total, and transfer concurrency to four. +Caller inputs are limited to 98 files, reserving two for the script and replay +history. Code API reserves one third of the job budget for all transfer batches +and negotiates each transfer's deadline before signing the request. +Its parent process keeps a 64-entry/32-MiB LRU input cache keyed by a stable, +Code-API-authorized digest; sandboxed commands cannot read that cache. Requests +against one workspace remain serialized, while negotiated lease slots allow +different registered roots to execute concurrently. + The native sandbox preserves standard `HTTP_PROXY`, `HTTPS_PROXY`, `ALL_PROXY`, and `NO_PROXY` names (including lowercase forms), plus Windows process and profile variables on Windows. SRT remains responsible for the final sandbox environment diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index a87f1a5a..289b0dc4 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -408,6 +408,11 @@ async function run( 'LIBRECHAT_CODE_COMMAND_SANDBOX must be native-srt or runtime', ); } + const nativeProgrammaticEnabled = + allowWorkspaceCommands && + commandSandboxMode === 'native-srt' && + process.platform !== 'win32' && + (fileRelayUpstream?.length ?? 0) > 0; const commandPolicy = resolveNativeSrtCommandPolicy( option(args, '--command-policy-preset') ?? process.env.LIBRECHAT_CODE_COMMAND_POLICY_PRESET?.trim().toLowerCase() ?? @@ -780,6 +785,9 @@ async function run( github.privateKeyPath, ].filter((path): path is string => path != null), allowedDomains: commandAllowedDomains, + ...(nativeProgrammaticEnabled + ? { programmaticFileUpstream: fileRelayUpstream } + : {}), ...(github.provider ? { maskedEnvironment: { @@ -819,6 +827,9 @@ async function run( workspaceTools = new SandboxWorkspaceTools({ workspaceTools, commandWorkspaces: roots.map((root) => root.id), + ...(nativeProgrammaticEnabled + ? { programmaticLanguages: ['bash'] } + : {}), commandSandbox: nativeCommandSandbox ?? new RuntimeWorkspaceCommandSandbox({ @@ -878,6 +889,9 @@ async function run( runtimeSupervisor, capabilities, workspaceTools, + ...(nativeProgrammaticEnabled && nativeCommandSandbox + ? { workspaceProgrammatic: nativeCommandSandbox } + : {}), ...(workspaceLeaseSlots > 1 || roots.length > 1 ? { workspaceQuarantines: new Map( diff --git a/packages/code/src/index.ts b/packages/code/src/index.ts index 5363f0c0..6f94b190 100644 --- a/packages/code/src/index.ts +++ b/packages/code/src/index.ts @@ -7,6 +7,7 @@ export * from './workspace.js'; export * from './workspace-runtime.js'; export * from './native-policy.js'; export * from './native-sandbox.js'; +export * from './native-programmatic.js'; export * from './native-process.js'; export * from './github.js'; export * from './worker.js'; diff --git a/packages/code/src/native-pool.test.ts b/packages/code/src/native-pool.test.ts index 3a24db34..cad30a52 100644 --- a/packages/code/src/native-pool.test.ts +++ b/packages/code/src/native-pool.test.ts @@ -13,6 +13,31 @@ const request = (workspaceId: string): WorkspaceExecuteCommandRequest => ({ workspaceId, command: 'fixture', }); + +test('native pool preflights every registered root with bounded concurrency', async () => { + const prepared: string[] = []; + let active = 0; + let peak = 0; + const pool = new NativeWorkspaceCommandPool(roots, 2, (options) => ({ + async prepare() { + active += 1; + peak = Math.max(peak, active); + await new Promise(resolve => setTimeout(resolve, 5)); + prepared.push(options.workspaceRoot); + active -= 1; + }, + async close() {}, + async execute() { + throw new Error('unreachable'); + }, + })); + + await pool.prepare(); + assert.deepEqual(prepared.sort(), ['/fixture/a', '/fixture/b', '/fixture/c']); + assert.equal(peak, 2); + await pool.close(); +}); + test('a known-clean executor failure is retired without replaying the command', async () => { let created = 0; let executed = 0; diff --git a/packages/code/src/native-pool.ts b/packages/code/src/native-pool.ts index 28f155e2..766409b1 100644 --- a/packages/code/src/native-pool.ts +++ b/packages/code/src/native-pool.ts @@ -2,6 +2,7 @@ import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; import { WorkspaceToolError } from './workspace.js'; import type { NativeProcessSandboxOptions } from './native-process.js'; import type { + BridgeWorkspaceProgrammaticRequest, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, } from './protocol.js'; @@ -10,7 +11,10 @@ interface Entry { sandbox: Pick< NativeProcessWorkspaceCommandSandbox, 'prepare' | 'execute' | 'close' - >; + > & + Partial< + Pick + >; busy: boolean; } @@ -92,12 +96,25 @@ export class NativeWorkspaceCommandPool { } async prepare(): Promise { - const entry = await this.allocate(this.roots.keys().next().value!); - try { - await entry.sandbox.prepare(); - } finally { - entry.busy = false; - } + const workspaceIds = [...this.roots.keys()]; + let next = 0; + await Promise.all( + Array.from( + { length: Math.min(this.capacity, workspaceIds.length) }, + async () => { + for (;;) { + const index = next++; + if (index >= workspaceIds.length) return; + const entry = await this.allocate(workspaceIds[index]!); + try { + await entry.sandbox.prepare(); + } finally { + entry.busy = false; + } + } + }, + ), + ); } async execute( @@ -136,6 +153,51 @@ export class NativeWorkspaceCommandPool { } } + async executeProgrammatic( + workspaceId: string, + request: BridgeWorkspaceProgrammaticRequest, + signal?: AbortSignal, + ): Promise { + const entry = await this.allocate(workspaceId); + let enteredExecutor = false; + try { + if (signal?.aborted) + throw new WorkspaceToolError( + 'Programmatic execution cancelled before dispatch', + 'EXECUTION_ABORTED', + ); + enteredExecutor = true; + if (!entry.sandbox.executeProgrammatic) { + throw new WorkspaceToolError( + 'Native programmatic executor is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + return await entry.sandbox.executeProgrammatic( + workspaceId, + request, + signal, + ); + } catch (error) { + if ( + enteredExecutor && + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted + ) { + try { + await entry.sandbox.close(); + if (this.entries.get(workspaceId) === entry) + this.entries.delete(workspaceId); + } catch { + /* Retain ownership for subsequent cleanup/shutdown. */ + } + } + throw error; + } finally { + entry.busy = false; + } + } + async close(): Promise { this.closing = true; await this.allocation; diff --git a/packages/code/src/native-process-child.ts b/packages/code/src/native-process-child.ts index 05ffedfe..2e8beb38 100644 --- a/packages/code/src/native-process-child.ts +++ b/packages/code/src/native-process-child.ts @@ -1,11 +1,16 @@ import { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; +import { NativeWorkspaceProgrammaticExecutor } from './native-programmatic.js'; import { WorkspaceToolError } from './workspace.js'; import type { NativeSrtWorkspaceCommandSandboxOptions } from './native-sandbox.js'; -import type { WorkspaceExecuteCommandRequest } from './protocol.js'; +import type { + BridgeWorkspaceProgrammaticRequest, + WorkspaceExecuteCommandRequest, +} from './protocol.js'; // This entrypoint is private to a forked trusted executor. No HTTP listener, // argv credentials, bridge token, or persisted pairing material is required. let sandbox: NativeSrtWorkspaceCommandSandbox | undefined; +let programmaticExecutor: NativeWorkspaceProgrammaticExecutor | undefined; let active: { id: string; controller: AbortController } | undefined; let busy = false; let credentials: Record = {}; @@ -44,11 +49,14 @@ process.on('message', async (raw: unknown) => { NativeSrtWorkspaceCommandSandboxOptions, 'maskedEnvironment' > & { + programmaticFileUpstream?: string; variables?: NonNullable< NativeSrtWorkspaceCommandSandboxOptions['maskedEnvironment'] >['variables']; }; request: WorkspaceExecuteCommandRequest; + programmaticRequest?: BridgeWorkspaceProgrammaticRequest; + workspaceId?: string; credentials?: Record; wrappedCommand?: string; }; @@ -62,7 +70,8 @@ process.on('message', async (raw: unknown) => { try { let result: unknown; if (message.type === 'prepare' && !sandbox) { - const { variables, ...options } = message.options; + const { variables, programmaticFileUpstream, ...options } = + message.options; sandbox = new NativeSrtWorkspaceCommandSandbox({ ...options, ...(variables @@ -80,11 +89,33 @@ process.on('message', async (raw: unknown) => { : {}), }); await sandbox.prepare(); + programmaticExecutor = programmaticFileUpstream + ? new NativeWorkspaceProgrammaticExecutor({ + sandbox, + upstreamUrl: programmaticFileUpstream, + }) + : undefined; + await programmaticExecutor?.prepare(); } else if (message.type === 'execute' && sandbox) { active = { id: message.id, controller: new AbortController() }; credentials = message.credentials ?? {}; wrappedCommand = message.wrappedCommand; result = await sandbox.execute(message.request, active.controller.signal); + } else if ( + message.type === 'programmatic' && + sandbox && + programmaticExecutor && + message.programmaticRequest && + typeof message.workspaceId === 'string' + ) { + active = { id: message.id, controller: new AbortController() }; + credentials = message.credentials ?? {}; + wrappedCommand = message.wrappedCommand; + result = await programmaticExecutor.execute( + message.programmaticRequest, + message.workspaceId, + active.controller.signal, + ); } else if (message.type === 'close' && sandbox) { await sandbox.close(); } else throw new Error('Invalid executor state'); diff --git a/packages/code/src/native-process.test.ts b/packages/code/src/native-process.test.ts index 800dcf11..2f38d642 100644 --- a/packages/code/src/native-process.test.ts +++ b/packages/code/src/native-process.test.ts @@ -1,13 +1,29 @@ import assert from 'node:assert/strict'; import { EventEmitter } from 'node:events'; import test from 'node:test'; +import { mkdtemp, rm, symlink, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { dirname, join } from 'node:path'; import type { ChildProcess, ForkOptions } from 'node:child_process'; import { NativeProcessWorkspaceCommandSandbox, nativeExecutorEnvironment, + trustedProgrammaticExecutable, } from './native-process.js'; import { WorkspaceToolError } from './workspace.js'; +test('preflight rejects relative and workspace-controlled executables including symlinks', async t => { + const root = await mkdtemp(join(tmpdir(), 'native-ptc-path-')); + const outside = await mkdtemp(join(tmpdir(), 'native-ptc-link-')); + t.after(async () => { await rm(root, { recursive: true, force: true }); await rm(outside, { recursive: true, force: true }); }); + const executable = join(root, 'bash'); + await writeFile(executable, '#!/bin/sh\nexit 0\n', { mode: 0o700 }); + await symlink(executable, join(outside, 'bash')); + await assert.rejects(trustedProgrammaticExecutable('./bash', root), /absolute/); + await assert.rejects(trustedProgrammaticExecutable(executable, root), /outside the workspace/); + await assert.rejects(trustedProgrammaticExecutable(join(outside, 'bash'), root), /outside the workspace/); +}); + const request = { protocolVersion: 1 as const, operation: 'execute_command' as const, @@ -43,13 +59,19 @@ function fixture( queueMicrotask(() => { if (message.type === 'prepare' && prepare) return prepare(child, message); - if (message.type === 'execute' && execute) + if ( + (message.type === 'execute' || message.type === 'programmatic') && + execute + ) return execute(child, message); if (message.type === 'cancel') return; child.emit('message', { id: message.id, ok: true, - ...(message.type === 'execute' ? { result } : {}), + ...(message.type === 'execute' || + message.type === 'programmatic' + ? { result } + : {}), }); }); return true; @@ -171,8 +193,98 @@ test('executor hands credentials over IPC only for the current command', async ( await sandbox.close(); }); +test('programmatic executor resolves and scopes credentials to its command', async () => { + const fake = fixture(); + const sandbox = new NativeProcessWorkspaceCommandSandbox( + { + workspaceRoot: tmpdir(), + programmaticFileUpstream: 'http://127.0.0.1:3190', + environment: { PATH: '/sandbox-only' }, + maskedEnvironment: { + variables: [{ name: 'TOKEN', injectHosts: ['github.com'] }], + async resolve() { + return { TOKEN: 'per-programmatic-secret' }; + }, + wrapCommand(command) { + return `wrapped ${command}`; + }, + }, + }, + fake.fork, + ); + const programmaticRequest = { + headers: {}, + body: { + language: 'bash' as const, + version: '5.2.0', + session_id: 'session', + files: [{ name: 'main.sh', content: 'git status' }], + }, + }; + await sandbox.executeProgrammatic('primary', programmaticRequest); + assert.equal( + JSON.stringify(fake.options).includes('per-programmatic-secret'), + false, + ); + const message = fake.messages.find( + candidate => candidate.type === 'programmatic', + )!; + const prepareMessage = fake.messages.find( + candidate => candidate.type === 'prepare', + )!; + assert.equal(typeof prepareMessage.options.jqPath, 'string'); + assert.equal(prepareMessage.options.jqPath.startsWith('/'), true); + assert.equal( + '/sandbox-only'.split(':').includes(dirname(prepareMessage.options.jqPath)), + false, + ); + assert.deepEqual(message.credentials, { TOKEN: 'per-programmatic-secret' }); + assert.equal( + message.wrappedCommand, + 'wrapped exec "$LIBRECHAT_CODE_BASH_PATH" "$LIBRECHAT_CODE_DATA_DIR/main.sh"', + ); + await sandbox.close(); +}); + +test('programmatic executor preserves a child-reported pre-dispatch failure', async () => { + const fake = fixture((child, message) => + child.emit('message', { + id: message.id, + ok: false, + code: 'COMMAND_UNAVAILABLE', + errorMessage: 'Programmatic input download failed', + mutation: false, + requiresQuarantine: false, + }), + ); + const sandbox = new NativeProcessWorkspaceCommandSandbox( + { + workspaceRoot: tmpdir(), + programmaticFileUpstream: 'http://127.0.0.1:3190', + }, + fake.fork, + ); + + await assert.rejects( + sandbox.executeProgrammatic('primary', { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'session', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + }), + (error: unknown) => + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted && + !error.requiresQuarantine, + ); + await sandbox.close(); +}); + test('executor loss after dispatch is an uncertain mutation and is never replayed', async () => { - const fake = fixture((child) => child.emit('exit', 1)); + const fake = fixture(child => child.emit('exit', 1)); const sandbox = new NativeProcessWorkspaceCommandSandbox( { workspaceRoot: '/workspace' }, fake.fork, @@ -180,16 +292,17 @@ test('executor loss after dispatch is an uncertain mutation and is never replaye await assert.rejects( sandbox.execute(request), (error: unknown) => - error instanceof WorkspaceToolError && error.mutationMayHaveCommitted, + error instanceof WorkspaceToolError && + error.mutationMayHaveCommitted, ); await assert.rejects(sandbox.execute(request), /unavailable/); - assert.equal(fake.messages.filter((m) => m.type === 'execute').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'execute').length, 1); await sandbox.close(); }); test('executor cancellation targets the active request and preserves mutation certainty', async () => { let dispatched!: () => void; - const dispatch = new Promise((resolve) => { + const dispatch = new Promise(resolve => { dispatched = resolve; }); const fake = fixture(() => dispatched()); @@ -202,7 +315,7 @@ test('executor cancellation targets the active request and preserves mutation ce await dispatch; await assert.rejects(sandbox.execute(request), /unavailable/); controller.abort(); - const command = fake.messages.find((m) => m.type === 'execute')!; + const command = fake.messages.find(m => m.type === 'execute')!; assert.deepEqual(fake.messages.at(-1), { type: 'cancel', id: command.id }); fake.child.emit('message', { id: command.id, @@ -224,7 +337,7 @@ test('executor cancellation targets the active request and preserves mutation ce test('executor ignores a cleanup exemption on non-cancellation failures', async () => { let dispatched!: () => void; - const dispatch = new Promise((resolve) => { + const dispatch = new Promise(resolve => { dispatched = resolve; }); const fake = fixture(() => dispatched()); @@ -234,7 +347,7 @@ test('executor ignores a cleanup exemption on non-cancellation failures', async ); const execution = sandbox.execute(request); await dispatch; - const command = fake.messages.find((message) => message.type === 'execute')!; + const command = fake.messages.find(message => message.type === 'execute')!; fake.child.emit('message', { id: command.id, ok: false, @@ -269,7 +382,8 @@ test('executor rejects mismatched results as uncertain and fences subsequent com await assert.rejects( sandbox.execute(request), (error: unknown) => - error instanceof WorkspaceToolError && error.mutationMayHaveCommitted, + error instanceof WorkspaceToolError && + error.mutationMayHaveCommitted, ); await assert.rejects(sandbox.execute(request), /unavailable/); await sandbox.close(); @@ -277,7 +391,7 @@ test('executor rejects mismatched results as uncertain and fences subsequent com test('executor close drains an active command before closing IPC', async () => { let dispatched!: () => void; - const dispatch = new Promise((resolve) => { + const dispatch = new Promise(resolve => { dispatched = resolve; }); const fake = fixture(() => dispatched()); @@ -288,16 +402,16 @@ test('executor close drains an active command before closing IPC', async () => { const execution = sandbox.execute(request); await dispatch; const closing = sandbox.close(); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal( - fake.messages.some((m) => m.type === 'close'), + fake.messages.some(m => m.type === 'close'), false, ); - const command = fake.messages.find((m) => m.type === 'execute')!; + const command = fake.messages.find(m => m.type === 'execute')!; fake.child.emit('message', { id: command.id, ok: true, result }); assert.deepEqual(await execution, result); await closing; - assert.equal(fake.messages.filter((m) => m.type === 'close').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'close').length, 1); await assert.rejects(sandbox.execute(request), /unavailable/); }); @@ -321,7 +435,7 @@ test('executor close resolves when the child exits during the close handshake', }, }); await sandbox.close(); - assert.equal(fake.messages.filter((m) => m.type === 'close').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'close').length, 1); await assert.rejects(sandbox.execute(request), /unavailable/); }); @@ -365,7 +479,7 @@ test('executor close skips the handshake once the child is already lost', async fake.child.emit('exit', 1, null); await sandbox.close(); assert.equal( - fake.messages.some((m) => m.type === 'close'), + fake.messages.some(m => m.type === 'close'), false, ); await assert.rejects(sandbox.execute(request), /unavailable/); @@ -377,17 +491,20 @@ test('executor startup loss is not reported as an applied mutation', async () => { workspaceRoot: '/workspace' }, (path, args, options) => { const child = fake.fork(path, args, options); - queueMicrotask(() => child.emit('error', new Error('startup failed'))); + queueMicrotask(() => + child.emit('error', new Error('startup failed')), + ); return child; }, ); await assert.rejects( sandbox.execute(request), (error: unknown) => - error instanceof WorkspaceToolError && !error.mutationMayHaveCommitted, + error instanceof WorkspaceToolError && + !error.mutationMayHaveCommitted, ); assert.equal( - fake.messages.some((m) => m.type === 'execute'), + fake.messages.some(m => m.type === 'execute'), false, ); await sandbox.close(); @@ -409,7 +526,7 @@ test('executor shutdown receipt fences reuse before the OS exit event', async () ); await assert.rejects(sandbox.execute(request)); await assert.rejects(sandbox.execute(request), /unavailable/); - assert.equal(fake.messages.filter((m) => m.type === 'execute').length, 1); + assert.equal(fake.messages.filter(m => m.type === 'execute').length, 1); await sandbox.close(); }); @@ -431,7 +548,8 @@ test('executor preserves bounded startup diagnostics and conventional host setti ok: false, mutation: false, code: 'COMMAND_UNAVAILABLE', - errorMessage: 'Native sandbox dependencies are unavailable: bubblewrap', + errorMessage: + 'Native sandbox dependencies are unavailable: bubblewrap', }), ); const sandbox = new NativeProcessWorkspaceCommandSandbox( @@ -465,7 +583,10 @@ test('executor matches POSIX names exactly and folds names only on Windows', () https_proxy: 'http://proxy:8080', }); assert.deepEqual( - nativeExecutorEnvironment({ Path: 'C:\\bin', Temp: 'C:\\temp' }, 'win32'), + nativeExecutorEnvironment( + { Path: 'C:\\bin', Temp: 'C:\\temp' }, + 'win32', + ), { Path: 'C:\\bin', Temp: 'C:\\temp' }, ); }); @@ -486,7 +607,8 @@ test('executor classifies every pre-dispatch setup failure as mutation-atomic', return {}; }, wrapCommand(command) { - if (failure === 'wrapper') throw new Error('wrapper failed'); + if (failure === 'wrapper') + throw new Error('wrapper failed'); return command; }, }, @@ -503,10 +625,12 @@ test('executor classifies every pre-dispatch setup failure as mutation-atomic', error instanceof WorkspaceToolError && !error.mutationMayHaveCommitted && error.code === - (failure === 'abort' ? 'EXECUTION_ABORTED' : 'COMMAND_UNAVAILABLE'), + (failure === 'abort' + ? 'EXECUTION_ABORTED' + : 'COMMAND_UNAVAILABLE'), ); assert.equal( - fake.messages.some((m) => m.type === 'execute'), + fake.messages.some(m => m.type === 'execute'), false, ); await sandbox.close(); diff --git a/packages/code/src/native-process.ts b/packages/code/src/native-process.ts index e1a02498..4ed3ffec 100644 --- a/packages/code/src/native-process.ts +++ b/packages/code/src/native-process.ts @@ -1,11 +1,21 @@ -import { fork } from 'node:child_process'; +import { execFile, fork } from 'node:child_process'; import { randomUUID } from 'node:crypto'; +import { constants as fsConstants } from 'node:fs'; +import { access, realpath } from 'node:fs/promises'; +import { isAbsolute, join, relative, sep } from 'node:path'; +import { promisify } from 'node:util'; import { WorkspaceToolError } from './workspace.js'; -import { isWorkspaceToolRequest, isWorkspaceToolResult } from './protocol.js'; +import { NATIVE_PROGRAMMATIC_COMMAND } from './native-programmatic.js'; +import { + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + isWorkspaceToolRequest, + isWorkspaceToolResult, +} from './protocol.js'; import type { ChildProcess, ForkOptions } from 'node:child_process'; import type { NativeSrtWorkspaceCommandSandboxOptions } from './native-sandbox.js'; import type { WorkspaceCommandSandbox } from './workspace.js'; import type { + BridgeWorkspaceProgrammaticRequest, WorkspaceExecuteCommandRequest, WorkspaceExecuteCommandResult, } from './protocol.js'; @@ -13,7 +23,84 @@ import type { export type NativeProcessSandboxOptions = Omit< NativeSrtWorkspaceCommandSandboxOptions, 'manager' | 'spawnCommand' | 'platform' ->; +> & { + /** Hardened Code API egress gateway used for execution-scoped files. */ + programmaticFileUpstream?: string; +}; + +const execFileAsync = promisify(execFile); + +async function systemProgrammaticExecutable( + name: string, + workspaceRoot: string, +): Promise { + // Preflight runs outside SRT. Never execute a workspace-controlled PATH + // entry (including cwd, node_modules/.bin, or a symlink to another root). + for (const directory of ['/opt/homebrew/bin', '/usr/local/bin', '/usr/bin', '/bin', '/home/linuxbrew/.linuxbrew/bin']) { + const candidate = join(directory, name); + try { + const canonical = await trustedProgrammaticExecutable(candidate, workspaceRoot); + if (!['/opt/homebrew/', '/usr/local/', '/usr/bin/', '/bin/', '/home/linuxbrew/.linuxbrew/'].some(root => canonical.startsWith(root))) continue; + return canonical; + } catch { + // Continue through the bounded PATH entries. + } + } +} + +export async function trustedProgrammaticExecutable(candidate: string, workspaceRoot: string): Promise { + if (!isAbsolute(candidate)) throw new Error('Programmatic executable must be absolute'); + const [canonical, root] = await Promise.all([realpath(candidate), realpath(workspaceRoot)]); + const path = relative(root, canonical); + if (path === '' || (!isAbsolute(path) && path !== '..' && !path.startsWith(`..${sep}`))) { + throw new Error('Programmatic executable must be outside the workspace'); + } + await access(canonical, fsConstants.X_OK); + return canonical; +} + +async function resolveProgrammaticShell( + options: NativeProcessSandboxOptions, +): Promise<{ shellPath: string; jqPath: string }> { + const environment = options.environment ?? process.env; + const shellPath = + options.shellPath != null + ? await trustedProgrammaticExecutable(options.shellPath, options.workspaceRoot) + : await systemProgrammaticExecutable('bash', options.workspaceRoot); + const jqPath = await systemProgrammaticExecutable('jq', options.workspaceRoot); + if (!shellPath || !jqPath) { + throw new WorkspaceToolError( + 'Native programmatic execution requires trusted host installations of Bash 5.2 or newer and jq', + 'COMMAND_UNAVAILABLE', + ); + } + try { + const [{ stdout: bashVersion }] = await Promise.all([ + execFileAsync(shellPath, ['--version'], { + env: nativeExecutorEnvironment(environment), + timeout: 5_000, + }), + execFileAsync(jqPath, ['--version'], { + env: nativeExecutorEnvironment(environment), + timeout: 5_000, + }), + ]); + const match = /version\s+(\d+)\.(\d+)/i.exec(bashVersion); + if ( + !match || + Number(match[1]) < 5 || + (Number(match[1]) === 5 && Number(match[2]) < 2) + ) { + throw new Error('unsupported Bash version'); + } + } catch { + throw new WorkspaceToolError( + 'Native programmatic execution requires trusted host installations of Bash 5.2 or newer and jq', + 'COMMAND_UNAVAILABLE', + ); + } + return { shellPath, jqPath }; +} /** Only OS discovery and conventional proxy settings cross into the executor. * In particular, never inherit NODE_OPTIONS, bridge identity, or app secrets. */ @@ -77,20 +164,22 @@ export function nativeExecutorEnvironment( * deadline, as opposed to a failure the executor reported explicitly. */ class NativeExecutorUnavailableError extends WorkspaceToolError { constructor(mutation: boolean) { - super('Native executor is unavailable', 'COMMAND_UNAVAILABLE', mutation); + super( + 'Native executor is unavailable', + 'COMMAND_UNAVAILABLE', + mutation, + ); this.name = 'NativeExecutorUnavailableError'; } } /** One persistent, process-isolated SRT manager per workspace. No automatic * restart/replay: losing IPC after execution starts is an ambiguous mutation. */ -export class NativeProcessWorkspaceCommandSandbox - implements WorkspaceCommandSandbox -{ +export class NativeProcessWorkspaceCommandSandbox implements WorkspaceCommandSandbox { readonly mutationFailuresAreAtomic = true as const; private child?: ChildProcess; private ready?: Promise; - private active?: Promise; + private active?: Promise; private closing?: Promise; private failed = false; private terminationTimer?: ReturnType; @@ -122,12 +211,17 @@ export class NativeProcessWorkspaceCommandSandbox } private async start(): Promise { + const programmaticExecutables = this.options.programmaticFileUpstream + ? await resolveProgrammaticShell(this.options) + : undefined; const child = this.forkExecutor( new URL('./native-process-child.js', import.meta.url), [], { execArgv: [], - env: nativeExecutorEnvironment(this.options.environment ?? process.env), + env: nativeExecutorEnvironment( + this.options.environment ?? process.env, + ), stdio: ['ignore', 'ignore', 'ignore', 'ipc'], serialization: 'json', }, @@ -165,6 +259,8 @@ export class NativeProcessWorkspaceCommandSandbox const processTerminationConfirmed = code === 'EXECUTION_ABORTED' && message.requiresQuarantine === false; + const mutationMayHaveCommitted = + pending.mutation && message.mutation !== false; pending.reject( new WorkspaceToolError( typeof message.errorMessage === 'string' && @@ -172,8 +268,8 @@ export class NativeProcessWorkspaceCommandSandbox ? message.errorMessage : 'Native executor request failed', code, - pending.mutation && message.mutation !== false, - pending.mutation && !processTerminationConfirmed, + mutationMayHaveCommitted, + mutationMayHaveCommitted && !processTerminationConfirmed, ), ); } @@ -192,6 +288,7 @@ export class NativeProcessWorkspaceCommandSandbox allowedDomains, homeDirectory, shellPath, + programmaticFileUpstream, } = this.options; await this.rpc( 'prepare', @@ -202,13 +299,15 @@ export class NativeProcessWorkspaceCommandSandbox protectedPaths, allowedDomains, homeDirectory, - shellPath, + shellPath: programmaticExecutables?.shellPath ?? shellPath, + jqPath: programmaticExecutables?.jqPath, + programmaticFileUpstream, variables: this.options.maskedEnvironment?.variables, }, }, 30_000, false, - ).catch((error) => { + ).catch(error => { this.failed = true; this.terminate(); throw error; @@ -223,7 +322,10 @@ export class NativeProcessWorkspaceCommandSandbox !isWorkspaceToolRequest(request) || request.operation !== 'execute_command' ) { - throw new WorkspaceToolError('Invalid native command', 'INVALID_REQUEST'); + throw new WorkspaceToolError( + 'Invalid native command', + 'INVALID_REQUEST', + ); } if (this.active || this.closing || this.failed) throw this.unavailable(false); @@ -236,12 +338,114 @@ export class NativeProcessWorkspaceCommandSandbox } } + async executeProgrammatic( + workspaceId: string, + request: BridgeWorkspaceProgrammaticRequest, + signal?: AbortSignal, + ): Promise { + if (this.active || this.closing || this.failed) + throw this.unavailable(false); + if (!this.options.programmaticFileUpstream) { + throw new WorkspaceToolError( + 'Native programmatic file transport is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const active = this.executeProgrammaticOnce( + request, + workspaceId, + signal, + ); + this.active = active; + try { + return await active; + } finally { + this.active = undefined; + } + } + + private async executeProgrammaticOnce( + request: BridgeWorkspaceProgrammaticRequest, + workspaceId: string, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted) + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + let credentials: Record | undefined; + let wrappedCommand: string | undefined; + try { + await this.prepare(); + if (signal?.aborted) throw new Error('aborted'); + credentials = await this.options.maskedEnvironment?.resolve(signal); + if (signal?.aborted) throw new Error('aborted'); + wrappedCommand = this.options.maskedEnvironment?.wrapCommand?.( + NATIVE_PROGRAMMATIC_COMMAND, + process.platform, + ); + if (signal?.aborted) throw new Error('aborted'); + } catch (error) { + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + throw error instanceof WorkspaceToolError + ? new WorkspaceToolError(error.message, error.code, false) + : new WorkspaceToolError( + 'Native programmatic executor setup failed before dispatch', + 'COMMAND_UNAVAILABLE', + ); + } + const result = await this.rpc( + 'programmatic', + { + programmaticRequest: request, + workspaceId, + credentials, + wrappedCommand, + }, + (request.body.run_timeout ?? 30_000) * + ((request.body.replay_tool_count ?? 0) > 0 ? 2 : 1) + + (Math.ceil( + request.body.files.filter(file => 'id' in file).length / 4, + ) + + Math.ceil( + (request.body.max_output_files ?? + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES) / 4, + )) * + (request.body.transfer_timeout_ms ?? 30_000) + + 5_000, + true, + signal, + ); + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + true, + ); + } + if (typeof result !== 'object' || result === null) { + this.failed = true; + this.terminate(); + throw this.unavailable(true); + } + return result; + } + private async executeOnce( request: WorkspaceExecuteCommandRequest, signal?: AbortSignal, ): Promise { if (signal?.aborted) - throw new WorkspaceToolError('Command aborted', 'EXECUTION_ABORTED'); + throw new WorkspaceToolError( + 'Command aborted', + 'EXECUTION_ABORTED', + ); let credentials: Record | undefined; let wrappedCommand: string | undefined; try { @@ -258,7 +462,10 @@ export class NativeProcessWorkspaceCommandSandbox // No execute RPC has been sent: setup, token refresh and wrapping cannot // have mutated the workspace. Do not quarantine it for setup failures. if (signal?.aborted) - throw new WorkspaceToolError('Command aborted', 'EXECUTION_ABORTED'); + throw new WorkspaceToolError( + 'Command aborted', + 'EXECUTION_ABORTED', + ); throw error instanceof WorkspaceToolError ? new WorkspaceToolError(error.message, error.code, false) : new WorkspaceToolError( @@ -324,7 +531,7 @@ export class NativeProcessWorkspaceCommandSandbox reject(this.unavailable(mutation)); }; try { - child.send({ type, id, ...payload }, (error) => { + child.send({ type, id, ...payload }, error => { if (error) sendFailed(); }); } catch { @@ -353,9 +560,12 @@ export class NativeProcessWorkspaceCommandSandbox await this.ready?.catch(() => undefined); try { if (this.child?.connected && !this.failed) - await this.rpc('close', {}, 10_000, false).catch((error: unknown) => { - if (!(error instanceof NativeExecutorUnavailableError)) throw error; - }); + await this.rpc('close', {}, 10_000, false).catch( + (error: unknown) => { + if (!(error instanceof NativeExecutorUnavailableError)) + throw error; + }, + ); } finally { this.failed = true; this.terminate(); diff --git a/packages/code/src/native-programmatic-live.test.ts b/packages/code/src/native-programmatic-live.test.ts new file mode 100644 index 00000000..2bc22356 --- /dev/null +++ b/packages/code/src/native-programmatic-live.test.ts @@ -0,0 +1,41 @@ +import assert from 'node:assert/strict'; +import { createServer } from 'node:http'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; +import type { AddressInfo } from 'node:net'; +import { NativeProcessWorkspaceCommandSandbox } from './native-process.js'; +import { resolveNativeSrtCommandPolicy } from './native-policy.js'; + +test('real SRT prevents speculative network effects under trusted-vm', { + skip: process.env.LIBRECHAT_CODE_LIVE_SRT_TESTS !== '1', + timeout: 30_000, +}, async () => { + const root = await mkdtemp(join(tmpdir(), 'native-ptc-effects-')); + let effects = 0; + const server = createServer((_req, res) => { effects += 1; res.end('ok'); }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const port = (server.address() as AddressInfo).port; + const executor = new NativeProcessWorkspaceCommandSandbox({ + workspaceRoot: root, + commandPolicy: resolveNativeSrtCommandPolicy('trusted-vm'), + programmaticFileUpstream: `http://127.0.0.1:${port}`, + }); + try { + await executor.prepare(); + const result = await executor.executeProgrammatic('primary', { headers: {}, body: { + language: 'bash', version: '5.2.0', session_id: 'isolated-canary', replay_tool_count: 1, + run_timeout: 5000, + files: [{ name: 'main.sh', content: `curl --noproxy '*' --connect-timeout 1 --max-time 2 -s -X POST http://127.0.0.1:${port}/effect >/dev/null\nprintf once >> commit.txt\n` }], + } }) as { run: { code: number } }; + assert.equal(result.run.code, 0); + assert.equal(effects, 1, 'probe must not emit a network effect'); + assert.equal(await readFile(join(root, 'commit.txt'), 'utf8'), 'once'); + } finally { + try { await executor.close(); } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(root, { recursive: true, force: true }); + } + } +}); diff --git a/packages/code/src/native-programmatic.test.ts b/packages/code/src/native-programmatic.test.ts new file mode 100644 index 00000000..902c707a --- /dev/null +++ b/packages/code/src/native-programmatic.test.ts @@ -0,0 +1,720 @@ +import assert from 'node:assert/strict'; +import { createHash } from 'node:crypto'; +import { createServer } from 'node:http'; +import { mkdir, mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; + +import { NativeWorkspaceProgrammaticExecutor } from './native-programmatic.js'; +import { WorkspaceToolError } from './workspace.js'; + +import type { AddressInfo } from 'node:net'; +import type { BridgeWorkspaceProgrammaticRequest } from './protocol.js'; + +test('stages skill files privately and returns generated artifacts', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-test-')); + const uploads = new Map(); + let downloadCount = 0; + const server = createServer(async (req, res) => { + assert.equal(req.headers['x-codeapi-egress-grant'], 'grant'); + if (req.method === 'GET') { + downloadCount += 1; + assert.match( + req.url ?? '', + /\/sessions\/input-session\/objects\/skill-file$/, + ); + res.end('skill-value'); + return; + } + assert.equal(req.method, 'PUT'); + assert.equal(req.headers['content-type'], 'text/plain'); + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(Buffer.from(chunk)); + uploads.set( + decodeURIComponent(req.headers['x-original-filename'] as string), + Buffer.concat(chunks), + ); + res.statusCode = 200; + res.end(); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as AddressInfo; + let observedDataDirectory = ''; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(request, dataDirectory) { + observedDataDirectory = dataDirectory; + assert.equal( + await readFile( + join(dataDirectory, 'skills/example/reference.txt'), + 'utf8', + ), + 'skill-value', + ); + await writeFile(join(dataDirectory, 'result.txt'), 'artifact'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: request.workspaceId, + exitCode: 0, + stdout: 'done\n', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + const request: BridgeWorkspaceProgrammaticRequest = { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + execution_id: 'execution-one', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [ + { name: 'main.sh', content: 'printf done' }, + { name: '_ptc_history.json', content: '{}' }, + { + name: 'skills/example/reference.txt', + id: 'skill-file', + storage_session_id: 'input-session', + input_cache_key: createHash('sha256') + .update('stable-authorized-input-identity') + .digest('hex'), + }, + ], + }, + }; + try { + const result = await executor.execute(request, 'primary'); + const replay = await executor.execute(request, 'primary'); + assert.equal(result.run.stdout, 'done\n'); + assert.equal(replay.run.stdout, 'done\n'); + assert.equal(result.session_id, 'output-session'); + assert.equal(downloadCount, 1); + await executor.execute( + { + ...request, + body: { ...request.body, execution_id: 'execution-two' }, + }, + 'primary', + ); + assert.equal(downloadCount, 2); + assert.equal(result.files.length, 1); + assert.equal(result.files[0]?.name, 'result.txt'); + assert.equal(uploads.get('result.txt')?.toString(), 'artifact'); + assert.deepEqual( + await readdir(observedDataDirectory).catch(() => []), + [], + ); + } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('reports persisted inputs deleted by selected-workspace execution', async t => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-delete-test-')); + t.after(() => rm(scratch, { recursive: true, force: true })); + const server = createServer((_req, res) => res.end('persisted input')); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + t.after(() => new Promise(resolve => server.close(() => resolve()))); + const address = server.address() as AddressInfo; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(request, dataDirectory) { + await rm(join(dataDirectory, 'input.txt')); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: request.workspaceId, + exitCode: 0, + stdout: '', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + + const result = await executor.execute({ + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + egress_grant: 'grant', + files: [ + { name: 'main.sh', content: 'rm input.txt' }, + { + name: 'input.txt', + id: 'input-id', + storage_session_id: 'input-session', + }, + ], + }, + }, 'primary'); + + assert.deepEqual(result.files, []); + assert.deepEqual(result.deleted_files, ['input.txt']); +}); + +test('retains read-only persisted inputs removed by selected-workspace execution', async t => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-readonly-delete-test-')); + t.after(() => rm(scratch, { recursive: true, force: true })); + let downloads = 0; + const server = createServer((_req, res) => { + downloads++; + res.setHeader('X-Read-Only', 'true'); + res.end('trusted skill'); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + t.after(() => new Promise(resolve => server.close(() => resolve()))); + const address = server.address() as AddressInfo; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(request, dataDirectory) { + await rm(join(dataDirectory, 'skills', 'review', 'SKILL.md')); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: request.workspaceId, + exitCode: 0, + stdout: '', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + const request = { + headers: {}, + body: { + language: 'bash' as const, + version: '5.2.0', + execution_id: 'readonly-execution', + session_id: 'execution-session', + egress_grant: 'grant', + files: [ + { name: 'main.sh', content: 'rm skills/review/SKILL.md' }, + { + name: 'skills/review/SKILL.md', + id: 'skill-id', + storage_session_id: 'skill-session', + input_cache_key: 'a'.repeat(64), + }, + ], + }, + }; + + const result = await executor.execute(request, 'primary'); + const replay = await executor.execute(request, 'primary'); + + assert.equal(downloads, 1); + assert.equal(result.deleted_files, undefined); + assert.equal(replay.deleted_files, undefined); +}); + +test('reports unsupported and rejected artifacts without invalidating a completed command', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-artifact-test-')); + const uploads = new Map(); + const server = createServer(async (req, res) => { + const name = decodeURIComponent(req.headers['x-original-filename'] as string); + uploads.set(name, req.headers['content-type']); + for await (const _chunk of req) { + // Drain the bounded request body before responding. + } + res.statusCode = name === 'image.png' ? 503 : 200; + res.end(); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as AddressInfo; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(_request, dataDirectory) { + await writeFile(join(dataDirectory, '_ptc_report.csv'), 'a,b\n1,2\n'); + await writeFile(join(dataDirectory, 'image.png'), 'not-a-real-png'); + await writeFile(join(dataDirectory, 'model.bin'), 'unsupported'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: 0, + stdout: 'done\n', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [{ name: 'main.sh', content: 'printf done' }], + }, + }, + 'primary', + ); + assert.equal(result.run.code, 0); + assert.deepEqual(result.files.map(file => file.name), ['_ptc_report.csv']); + assert.deepEqual(result.artifact_delivery, { + code: 'artifact_delivery_failed', + status: 'partial', + attempted: 3, + delivered: 1, + failed: 2, + }); + assert.equal(uploads.get('_ptc_report.csv'), 'text/csv'); + assert.equal(uploads.get('image.png'), 'image/png'); + assert.equal(uploads.has('model.bin'), false); + } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('reports artifact transport failure without quarantining a completed command', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-artifact-transport-test-')); + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + fetchImpl: async () => { + throw new TypeError('transport unavailable'); + }, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(_request, dataDirectory) { + await writeFile(join(dataDirectory, 'result.txt'), 'artifact'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: 0, + stdout: 'done\n', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [{ name: 'main.sh', content: 'printf done' }], + }, + }, + 'primary', + ); + assert.equal(result.run.code, 0); + assert.deepEqual(result.files, []); + assert.deepEqual(result.artifact_delivery, { + code: 'artifact_delivery_failed', + status: 'failed', + attempted: 1, + delivered: 0, + failed: 1, + }); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('preflights copy-on-write isolation and removes its private snapshot', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-preflight-test-')); + let executionDirectory = ''; + let probes = 0; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + executionDirectory = await mkdtemp(join(scratch, 'execution-')); + return executionDirectory; + }, + async createProgrammaticProbeWorkspace(directory) { + probes += 1; + const workspace = join(directory, 'workspace'); + await mkdir(workspace); + return workspace; + }, + async executeProgrammatic() { + throw new Error('unreachable'); + }, + }, + }); + try { + await executor.prepare(); + assert.equal(probes, 1); + assert.deepEqual(await readdir(executionDirectory).catch(() => []), []); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('keeps replay probes read-only and commits the script exactly once', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-probe-test-')); + const phases: boolean[] = []; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async createProgrammaticProbeWorkspace(executionDirectory) { + const workspace = join(executionDirectory, 'workspace'); + await mkdir(workspace); + return workspace; + }, + async executeProgrammatic( + _request, + dataDirectory, + _signal, + options, + ) { + phases.push(options?.probe === true); + if (options?.probe === true) { + assert.match(options.workspaceRoot ?? '', /\/workspace$/); + } + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: options?.probe ? 1 : 0, + stdout: options?.probe ? 'probe\n' : 'commit\n', + stderr: options?.probe ? 'expected probe denial\n' : '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + execution_id: 'probe-then-commit', + replay_tool_count: 1, + session_id: 'execution-session', + files: [ + { name: 'main.sh', content: 'printf done' }, + { name: '_ptc_history.json', content: '{}' }, + ], + }, + }, + 'primary', + ); + assert.deepEqual(phases, [true, false]); + assert.equal(result.run.stdout, 'commit\n'); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('returns pending calls from the private control file even when stdout truncates', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-control-test-')); + let phases = 0; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async createProgrammaticProbeWorkspace(executionDirectory) { + const workspace = join(executionDirectory, 'workspace'); + await mkdir(workspace); + return workspace; + }, + async executeProgrammatic( + _request, + dataDirectory, + _signal, + options, + ) { + assert.equal(options?.probe, true); + phases += 1; + await writeFile( + join(dataDirectory, '_ptc_pending_result.json'), + JSON.stringify({ + pending: [ + { + call_id: 'call_001', + tool_name: 'lookup', + input: {}, + }, + ], + }), + ); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: 'primary', + exitCode: 0, + stdout: 'x'.repeat(256 * 1024), + stderr: '', + truncated: true, + timedOut: false, + }; + }, + }, + }); + try { + const result = await executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + execution_id: 'truncated-control', + replay_tool_count: 1, + session_id: 'execution-session', + files: [ + { name: 'main.sh', content: 'lookup "{}"' }, + { name: '_ptc_history.json', content: '{}' }, + ], + }, + }, + 'primary', + ); + assert.equal(phases, 1); + assert.equal(result.run.stdout, ''); + assert.equal(result.run.stderr, ''); + assert.deepEqual(JSON.parse(result.pending_tool_calls_payload ?? ''), { + pending: [{ call_id: 'call_001', tool_name: 'lookup', input: {} }], + }); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +test('rejects traversal before creating execution state', async () => { + let allocated = false; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { + allocated = true; + return '/unused'; + }, + async executeProgrammatic() { + throw new Error('unreachable'); + }, + }, + }); + await assert.rejects( + executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + files: [{ name: '../main.sh', content: 'echo unsafe' }], + }, + }, + 'primary', + ), + /Invalid selected-workspace programmatic request/, + ); + assert.equal(allocated, false); +}); + +test('rejects artifacts above the negotiated byte ceiling before upload', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-output-limit-')); + let uploads = 0; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + fetchImpl: async () => { + uploads += 1; + return new Response(); + }, + sandbox: { + async createExecutionDirectory() { + return await mkdtemp(join(scratch, 'execution-')); + }, + async executeProgrammatic(request, dataDirectory) { + await writeFile(join(dataDirectory, 'artifact.txt'), 'too large'); + return { + protocolVersion: 1, + operation: 'execute_command' as const, + workspaceId: request.workspaceId, + exitCode: 0, + stdout: '', + stderr: '', + truncated: false, + timedOut: false, + }; + }, + }, + }); + try { + await assert.rejects( + executor.execute( + { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + max_output_file_bytes: 4, + files: [{ name: 'main.sh', content: 'printf done' }], + }, + }, + 'primary', + ), + /exceeds the file limit/, + ); + assert.equal(uploads, 0); + } finally { + await rm(scratch, { recursive: true, force: true }); + } +}); + +for (const failure of ['truncated', 'process-error']) test(`a failed speculative probe does not quarantine the real workspace (${failure})`, async t => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-failed-probe-')); + t.after(() => rm(scratch, { recursive: true, force: true })); + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { return await mkdtemp(join(scratch, 'execution-')); }, + async createProgrammaticProbeWorkspace(directory) { const root = join(directory, 'workspace'); await mkdir(root); return root; }, + async executeProgrammatic(request, _directory, _signal, options) { + assert.equal(options?.probe, true); + if (failure === 'process-error') throw new WorkspaceToolError('probe output exceeded its limit', 'COMMAND_UNAVAILABLE', true, true); + return { protocolVersion: 1, operation: 'execute_command', workspaceId: request.workspaceId, + exitCode: 0, stdout: '', stderr: '', truncated: true, timedOut: false }; + }, + }, + }); + await assert.rejects(executor.execute({ headers: {}, body: { + language: 'bash', version: '5.2.0', session_id: 'session', replay_tool_count: 1, + files: [{ name: 'main.sh', content: 'true' }], + } }, 'primary'), (error: unknown) => { + assert.match(String(error), /probe output exceeded/); + assert.equal((error as { requiresQuarantine: boolean }).requiresQuarantine, false); + assert.equal((error as { mutationMayHaveCommitted: boolean }).mutationMayHaveCommitted, false); + return true; + }); +}); + +test('unchanged inputs do not consume the negotiated output budget', async t => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-unchanged-')); + t.after(() => rm(scratch, { recursive: true, force: true })); + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: 'http://127.0.0.1:1', + sandbox: { + async createExecutionDirectory() { return await mkdtemp(join(scratch, 'execution-')); }, + async executeProgrammatic(request) { + return { protocolVersion: 1, operation: 'execute_command', workspaceId: request.workspaceId, + exitCode: 0, stdout: '', stderr: '', truncated: false, timedOut: false }; + }, + }, + }); + const result = await executor.execute({ headers: {}, body: { + language: 'bash', version: '5.2.0', session_id: 'session', max_output_file_bytes: 1, + files: [{ name: 'main.sh', content: 'true' }, { name: 'input.txt', content: 'unchanged input' }], + } }, 'primary'); + assert.deepEqual(result.files, []); +}); + +test('stops admitting downloads and drains in-flight transfers before cleanup', async () => { + const scratch = await mkdtemp(join(tmpdir(), 'native-ptc-transfer-test-')); + let executionDirectory = ''; + let requestCount = 0; + const server = createServer((req, res) => { + requestCount += 1; + if (requestCount === 1) { + res.statusCode = 503; + res.end(); + return; + } + setTimeout(() => res.end('in-flight'), 25); + }); + await new Promise(resolve => server.listen(0, '127.0.0.1', resolve)); + const address = server.address() as AddressInfo; + const executor = new NativeWorkspaceProgrammaticExecutor({ + upstreamUrl: `http://127.0.0.1:${address.port}`, + sandbox: { + async createExecutionDirectory() { + executionDirectory = await mkdtemp(join(scratch, 'execution-')); + return executionDirectory; + }, + async executeProgrammatic() { + throw new Error('unreachable'); + }, + }, + }); + const request: BridgeWorkspaceProgrammaticRequest = { + headers: {}, + body: { + language: 'bash', + version: '5.2.0', + session_id: 'execution-session', + output_session_id: 'output-session', + egress_grant: 'grant', + files: [ + ...Array.from({ length: 8 }, (_, index) => ({ + name: `inputs/${index}.txt`, + id: `input-${index}`, + storage_session_id: 'input-session', + })), + { name: 'main.sh', content: 'printf done' }, + ], + }, + }; + try { + const startedAt = performance.now(); + await assert.rejects( + executor.execute(request, 'primary'), + /Programmatic input download failed with HTTP 503/, + ); + assert.ok(performance.now() - startedAt >= 20); + assert.ok(requestCount <= 4); + assert.deepEqual(await readdir(executionDirectory).catch(() => []), []); + } finally { + await new Promise(resolve => server.close(() => resolve())); + await rm(scratch, { recursive: true, force: true }); + } +}); diff --git a/packages/code/src/native-programmatic.ts b/packages/code/src/native-programmatic.ts new file mode 100644 index 00000000..bd2700f2 --- /dev/null +++ b/packages/code/src/native-programmatic.ts @@ -0,0 +1,838 @@ +import { createHash, randomBytes } from 'node:crypto'; +import { constants } from 'node:fs'; +import { cp, mkdir, open, readdir, rm, writeFile } from 'node:fs/promises'; +import { dirname, join, relative, sep } from 'node:path'; + +import { + BRIDGE_PROTOCOL_VERSION, + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES, + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES, + bridgeArtifactMediaType, + isBridgeWorkspaceProgrammaticRequest, + isSafePortableRelativePath, + isSupportedBridgeArtifactName, +} from './protocol.js'; +import { validateFileRelayUpstream } from './relay.js'; +import { WorkspaceToolError } from './workspace.js'; + +import type { + BridgeProgrammaticPayloadFile, + BridgeWorkspaceProgrammaticRequest, + WorkspaceExecuteCommandResult, +} from './protocol.js'; +import type { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; + +const EGRESS_GRANT_HEADER = 'X-CodeAPI-Egress-Grant'; +const EXECUTION_MAIN_FILE = 'main.sh'; +const EXECUTION_HISTORY_FILE = '_ptc_history.json'; +const EXECUTION_CONTROL_FILE = '_ptc_pending_result.json'; +export const NATIVE_PROGRAMMATIC_COMMAND = + 'exec "$LIBRECHAT_CODE_BASH_PATH" "$LIBRECHAT_CODE_DATA_DIR/main.sh"'; +const TRANSFER_TIMEOUT_MS = 30_000; +const TRANSFER_CONCURRENCY = 4; +const MAX_WALK_ENTRIES = 2_000; +const INPUT_CACHE_MAX_ENTRIES = 64; +const INPUT_CACHE_MAX_BYTES = 32 * 1024 * 1024; +const CONTROL_PAYLOAD_MAX_BYTES = 512 * 1024; + +type ProgrammaticFileResult = { + id: string; + name: string; + storage_session_id: string; + modified_from?: { id: string; storage_session_id: string }; +}; + +type ProgrammaticResult = { + language: 'bash'; + version: string; + session_id: string; + files: ProgrammaticFileResult[]; + deleted_files?: string[]; + artifact_delivery?: { + code: 'artifact_delivery_failed'; + status: 'partial' | 'failed'; + attempted: number; + delivered: number; + failed: number; + }; + pending_tool_calls_payload?: string; + run: { + stdout: string; + stderr: string; + code: number | null; + signal: string | null; + output: string; + memory: null; + message: string | null; + status: string | null; + cpu_time: null; + wall_time: number; + }; +}; + +type InputBaseline = { + sha256: string; + source?: { id: string; storage_session_id: string }; + readOnly?: boolean; +}; + +type CachedInput = { bytes: Buffer; readOnly: boolean }; + +function sha256(value: Uint8Array): string { + return createHash('sha256').update(value).digest('hex'); +} + +function outputFileId(): string { + return randomBytes(18).toString('base64url').slice(0, 21); +} + +function localPath(root: string, name: string): string { + if (!isSafePortableRelativePath(name)) { + throw new WorkspaceToolError( + 'Invalid programmatic file path', + 'INVALID_PATH', + ); + } + const path = join(root, ...name.split('/')); + const child = relative(root, path); + if (child === '' || child === '..' || child.startsWith(`..${sep}`)) { + throw new WorkspaceToolError( + 'Invalid programmatic file path', + 'INVALID_PATH', + ); + } + return path; +} + +async function readBoundedResponse( + response: Response, + signal: AbortSignal, +): Promise { + const declaredLength = response.headers.get('content-length'); + if ( + declaredLength != null && + (!/^\d+$/.test(declaredLength) || + Number(declaredLength) > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES) + ) { + await response.body?.cancel(); + throw new WorkspaceToolError( + 'Programmatic input exceeds the file limit', + 'READ_LIMIT_EXCEEDED', + ); + } + if (!response.body) return Buffer.alloc(0); + const reader = response.body.getReader(); + const chunks: Buffer[] = []; + let bytes = 0; + try { + for (;;) { + if (signal.aborted) throw signal.reason; + const { done, value } = await reader.read(); + if (done) break; + bytes += value.byteLength; + if (bytes > BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES) { + throw new WorkspaceToolError( + 'Programmatic input exceeds the file limit', + 'READ_LIMIT_EXCEEDED', + ); + } + chunks.push(Buffer.from(value)); + } + } finally { + await reader.cancel().catch(() => undefined); + } + return Buffer.concat(chunks, bytes); +} + +async function mapConcurrent( + values: readonly T[], + concurrency: number, + action: (value: T) => Promise, +): Promise { + const results = new Array(values.length); + let next = 0; + let failed = false; + let failure: unknown; + await Promise.all( + Array.from( + { length: Math.min(concurrency, values.length) }, + async () => { + for (;;) { + if (failed) return; + const index = next++; + if (index >= values.length) return; + try { + results[index] = await action(values[index]!); + } catch (error) { + if (!failed) { + failed = true; + failure = error; + } + return; + } + } + }, + ), + ); + if (failed) throw failure; + return results; +} + +async function listRegularFiles(root: string): Promise { + const files: string[] = []; + const pending = ['']; + let entries = 0; + while (pending.length > 0) { + const directory = pending.pop()!; + for (const entry of await readdir(join(root, directory), { + withFileTypes: true, + })) { + if (++entries > MAX_WALK_ENTRIES) { + throw new WorkspaceToolError( + 'Programmatic output contains too many entries', + 'WRITE_LIMIT_EXCEEDED', + ); + } + const name = directory ? `${directory}/${entry.name}` : entry.name; + if (!isSafePortableRelativePath(name)) continue; + if (entry.isSymbolicLink()) continue; + if (entry.isDirectory()) pending.push(name); + else if (entry.isFile()) files.push(name); + } + } + return files.sort(); +} + +export interface NativeWorkspaceProgrammaticOptions { + sandbox: Pick< + NativeSrtWorkspaceCommandSandbox, + 'createExecutionDirectory' | 'executeProgrammatic' + > & + Partial< + Pick + >; + upstreamUrl: string; + fetchImpl?: typeof fetch; +} + +/** + * Executes one replay-mode Bash PTC iteration in an attached workspace. + * Program code and injected files live in a private SRT scratch directory; + * the selected repository remains the command cwd and is never used as a + * transport cache. + */ +export class NativeWorkspaceProgrammaticExecutor { + private readonly upstream: URL; + private readonly fetchImpl: typeof fetch; + /** Parent-process cache: sandboxed children cannot inspect this memory. */ + private readonly inputCache = new Map< + string, + CachedInput & { lastUsed: number } + >(); + private inputCacheBytes = 0; + + constructor(private readonly options: NativeWorkspaceProgrammaticOptions) { + this.upstream = validateFileRelayUpstream(options.upstreamUrl); + this.fetchImpl = options.fetchImpl ?? fetch; + } + + /** + * Prove copy-on-write isolation before the worker advertises Bash PTC. + * The probe uses the exact registered root and private scratch path that a + * real replay will use, then removes the snapshot before registration. + */ + async prepare(signal?: AbortSignal): Promise { + const createProbeWorkspace = + this.options.sandbox.createProgrammaticProbeWorkspace; + if (createProbeWorkspace == null) { + throw new WorkspaceToolError( + 'Selected-workspace PTC probe isolation is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const executionDirectory = + await this.options.sandbox.createExecutionDirectory(); + try { + await createProbeWorkspace.call( + this.options.sandbox, + executionDirectory, + signal, + ); + } finally { + await rm(executionDirectory, { recursive: true, force: true }); + } + } + + private cacheKey( + executionId: string | undefined, + file: Extract, + ): string | undefined { + return executionId && file.input_cache_key + ? `${executionId}:${file.input_cache_key}` + : undefined; + } + + private cachedInput(key: string): CachedInput | undefined { + const cached = this.inputCache.get(key); + if (!cached) return undefined; + cached.lastUsed = Date.now(); + return { bytes: cached.bytes, readOnly: cached.readOnly }; + } + + private cacheInput(key: string, input: CachedInput): void { + const { bytes } = input; + if (bytes.byteLength > INPUT_CACHE_MAX_BYTES) return; + const existing = this.inputCache.get(key); + if (existing) this.inputCacheBytes -= existing.bytes.byteLength; + while ( + this.inputCache.size >= INPUT_CACHE_MAX_ENTRIES || + this.inputCacheBytes + bytes.byteLength > INPUT_CACHE_MAX_BYTES + ) { + let oldestKey: string | undefined; + let oldestAt = Number.POSITIVE_INFINITY; + for (const [candidate, value] of this.inputCache) { + if (value.lastUsed < oldestAt) { + oldestAt = value.lastUsed; + oldestKey = candidate; + } + } + if (!oldestKey) break; + this.inputCacheBytes -= + this.inputCache.get(oldestKey)!.bytes.byteLength; + this.inputCache.delete(oldestKey); + } + this.inputCache.set(key, { ...input, lastUsed: Date.now() }); + this.inputCacheBytes += bytes.byteLength; + } + + private async downloadInput( + file: Extract, + grant: string, + executionId: string | undefined, + signal?: AbortSignal, + transferTimeoutMs = TRANSFER_TIMEOUT_MS, + ): Promise { + const key = this.cacheKey(executionId, file); + const cached = key ? this.cachedInput(key) : undefined; + if (cached) return cached; + const controller = new AbortController(); + const abort = (): void => controller.abort(signal?.reason); + signal?.addEventListener('abort', abort, { once: true }); + const timer = setTimeout(() => controller.abort(), transferTimeoutMs); + try { + const response = await this.fetchImpl( + new URL( + `sessions/${encodeURIComponent(file.storage_session_id)}/objects/${encodeURIComponent(file.id)}`, + `${this.upstream.toString().replace(/\/+$/, '')}/`, + ), + { + headers: { [EGRESS_GRANT_HEADER]: grant }, + redirect: 'error', + signal: controller.signal, + }, + ); + if (!response.ok) { + await response.body?.cancel(); + throw new WorkspaceToolError( + `Programmatic input download failed with HTTP ${response.status}`, + 'COMMAND_UNAVAILABLE', + ); + } + const bytes = await readBoundedResponse( + response, + controller.signal, + ); + const input = { + bytes, + readOnly: response.headers.get('x-read-only')?.toLowerCase() === 'true', + }; + if (key) this.cacheInput(key, input); + return input; + } finally { + clearTimeout(timer); + signal?.removeEventListener('abort', abort); + } + } + + async execute( + request: BridgeWorkspaceProgrammaticRequest, + workspaceId: string, + signal?: AbortSignal, + ): Promise { + if (!isBridgeWorkspaceProgrammaticRequest(request)) { + throw new WorkspaceToolError( + 'Invalid selected-workspace programmatic request', + 'INVALID_REQUEST', + ); + } + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + const grant = request.body.egress_grant; + const refFiles = request.body.files.filter( + ( + file, + ): file is Extract => + 'id' in file, + ); + if (refFiles.length > 0 && !grant) { + throw new WorkspaceToolError( + 'Programmatic input grant is unavailable', + 'INVALID_REQUEST', + ); + } + const executionDirectory = + await this.options.sandbox.createExecutionDirectory(); + const inputDirectory = join(executionDirectory, 'inputs'); + let dataDirectory = join(executionDirectory, 'final'); + const baselines = new Map(); + let totalInputBytes = 0; + const startedAt = performance.now(); + let commandDispatched = false; + try { + await mkdir(inputDirectory, { mode: 0o700 }); + await mapConcurrent( + request.body.files, + TRANSFER_CONCURRENCY, + async (file): Promise => { + const input = + 'content' in file + ? { bytes: Buffer.from(file.content), readOnly: false } + : await this.downloadInput( + file, + grant!, + request.body.execution_id, + signal, + request.body.transfer_timeout_ms, + ); + const { bytes } = input; + totalInputBytes += bytes.byteLength; + if ( + totalInputBytes > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES + ) { + throw new WorkspaceToolError( + 'Programmatic inputs exceed the total byte limit', + 'READ_LIMIT_EXCEEDED', + ); + } + const path = localPath(inputDirectory, file.name); + await mkdir(dirname(path), { + recursive: true, + mode: 0o700, + }); + await writeFile(path, bytes, { flag: 'wx', mode: 0o600 }); + baselines.set(file.name, { + sha256: sha256(bytes), + ...(input.readOnly ? { readOnly: true } : {}), + ...('id' in file + ? { + source: { + id: file.id, + storage_session_id: + file.storage_session_id, + }, + } + : {}), + }); + }, + ); + + const run = async ( + directory: string, + probe: boolean, + workspaceRoot?: string, + ): Promise => { + await cp(inputDirectory, directory, { + recursive: true, + force: false, + errorOnExist: true, + mode: constants.COPYFILE_FICLONE, + }); + if (!probe) commandDispatched = true; + return await this.options.sandbox.executeProgrammatic( + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'execute_command', + workspaceId, + command: NATIVE_PROGRAMMATIC_COMMAND, + timeoutMs: Math.min( + request.body.run_timeout ?? + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS, + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS, + ), + maxOutputBytes: + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES, + }, + directory, + signal, + { probe, workspaceRoot }, + ); + }; + + const readPending = async ( + directory: string, + ): Promise => { + try { + const path = join(directory, EXECUTION_CONTROL_FILE); + const handle = await open( + path, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + try { + const metadata = await handle.stat(); + if ( + !metadata.isFile() || + metadata.size === 0 || + metadata.size > CONTROL_PAYLOAD_MAX_BYTES + ) { + throw new WorkspaceToolError( + 'Native programmatic control frame is invalid', + 'COMMAND_UNAVAILABLE', + ); + } + return await handle.readFile('utf8'); + } finally { + await handle.close(); + } + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') + return; + throw error; + } + }; + + if ((request.body.replay_tool_count ?? 0) > 0) { + const probeDirectory = join(executionDirectory, 'probe'); + const createProbeWorkspace = + this.options.sandbox.createProgrammaticProbeWorkspace; + if (createProbeWorkspace == null) { + throw new WorkspaceToolError( + 'Selected-workspace PTC probe isolation is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const probeWorkspace = await createProbeWorkspace.call( + this.options.sandbox, + executionDirectory, + signal, + ); + const probeResult = await run( + probeDirectory, + true, + probeWorkspace, + ); + const pending = await readPending(probeDirectory); + if (pending) { + return this.result( + request, + { + ...probeResult, + /** Probe output is speculative and the script will + * run once under its real policy after tool + * resolution. Never duplicate it or expose + * expected read-only policy denials to callers. */ + stdout: '', + stderr: '', + }, + [], + performance.now() - startedAt, + pending, + ); + } + if (probeResult.truncated) { + throw new WorkspaceToolError( + 'Native programmatic probe output exceeded its limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + if ( + probeResult.timedOut || + probeResult.signal + ) { + return this.result( + request, + probeResult, + [], + performance.now() - startedAt, + ); + } + /** A read-only probe commonly exits non-zero after it reaches + * an intentional workspace write denial. With no pending call, + * run the script once under its real policy so ordinary writes + * and their resulting exit status are evaluated exactly once. */ + } + + const commandResult = await run(dataDirectory, false); + if (await readPending(dataDirectory)) { + throw new WorkspaceToolError( + 'Native programmatic commit pass requested an unexpected replay tool', + 'COMMAND_UNAVAILABLE', + true, + true, + ); + } + if (commandResult.truncated) { + throw new WorkspaceToolError( + 'Native programmatic output exceeded its limit', + 'WRITE_LIMIT_EXCEEDED', + true, + true, + ); + } + + const outputSessionId = request.body.output_session_id; + const survivingNames = new Set(await listRegularFiles(dataDirectory)); + const deletedFiles = refFiles + .filter( + file => + baselines.get(file.name)?.readOnly !== true && + !survivingNames.has(file.name), + ) + .map(file => file.name); + const outputNames = [...survivingNames].filter( + name => + name !== EXECUTION_MAIN_FILE && + name !== EXECUTION_HISTORY_FILE && + name !== EXECUTION_CONTROL_FILE && + !name.startsWith('skills/'), + ); + const changed: Array<{ + name: string; + bytes: Buffer; + source?: { id: string; storage_session_id: string }; + }> = []; + let totalOutputBytes = 0; + for (const name of outputNames) { + const path = localPath(dataDirectory, name); + const baseline = baselines.get(name); + const maxOutputFileBytes = Math.min( + request.body.max_output_file_bytes ?? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + ); + let bytes: Buffer; + const handle = await open( + path, + constants.O_RDONLY | constants.O_NOFOLLOW, + ); + try { + const metadata = await handle.stat(); + if (!metadata.isFile()) continue; + // An unchanged input is not an output. It may legitimately + // exceed the negotiated output ceiling, but never the + // protocol's bounded input limit. + if (metadata.size > (baseline ? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES : maxOutputFileBytes)) { + throw new WorkspaceToolError( + 'Programmatic output exceeds the file limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + bytes = await handle.readFile(); + } finally { + await handle.close(); + } + if (baseline?.sha256 === sha256(bytes)) continue; + if (bytes.byteLength > maxOutputFileBytes) { + throw new WorkspaceToolError( + 'Programmatic output exceeds the file limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + totalOutputBytes += bytes.byteLength; + if ( + totalOutputBytes > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES + ) { + throw new WorkspaceToolError( + 'Programmatic outputs exceed the total byte limit', + 'WRITE_LIMIT_EXCEEDED', + ); + } + changed.push({ name, bytes, source: baseline?.source }); + } + const maxOutputFiles = Math.min( + request.body.max_output_files ?? + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + ); + if (changed.length > maxOutputFiles) { + throw new WorkspaceToolError( + 'Programmatic output contains too many files', + 'WRITE_LIMIT_EXCEEDED', + ); + } + const uploadable = changed.filter(({ name }) => + isSupportedBridgeArtifactName(name), + ); + if (uploadable.length > 0 && (!grant || !outputSessionId)) { + throw new WorkspaceToolError( + 'Programmatic output grant is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + const uploadResults = await mapConcurrent( + uploadable, + TRANSFER_CONCURRENCY, + async ({ name, bytes, source }): Promise => { + const id = outputFileId(); + const controller = new AbortController(); + const abort = (): void => controller.abort(signal?.reason); + signal?.addEventListener('abort', abort, { once: true }); + const timer = setTimeout( + () => controller.abort(), + request.body.transfer_timeout_ms ?? TRANSFER_TIMEOUT_MS, + ); + try { + let response: Response; + try { + response = await this.fetchImpl( + new URL( + `sessions/${encodeURIComponent(outputSessionId!)}/objects/${id}`, + `${this.upstream.toString().replace(/\/+$/, '')}/`, + ), + { + method: 'PUT', + headers: { + [EGRESS_GRANT_HEADER]: grant!, + 'Content-Type': bridgeArtifactMediaType(name), + 'Content-Length': String(bytes.byteLength), + 'X-Original-Filename': encodeURIComponent(name), + }, + body: new Uint8Array(bytes), + redirect: 'error', + signal: controller.signal, + }, + ); + } catch (error) { + if (signal?.aborted) throw error; + return undefined; + } + await response.body?.cancel(); + if (!response.ok) { + return undefined; + } + return { + id, + name, + storage_session_id: outputSessionId!, + ...(source ? { modified_from: source } : {}), + }; + } finally { + clearTimeout(timer); + signal?.removeEventListener('abort', abort); + } + }, + ); + const files = uploadResults.filter( + (file): file is ProgrammaticFileResult => file != null, + ); + const artifactDelivery = + files.length < changed.length + ? { + code: 'artifact_delivery_failed' as const, + status: files.length > 0 ? ('partial' as const) : ('failed' as const), + attempted: changed.length, + delivered: files.length, + failed: changed.length - files.length, + } + : undefined; + return this.result( + request, + commandResult, + files, + performance.now() - startedAt, + undefined, + artifactDelivery, + deletedFiles, + ); + } catch (error) { + if (!commandDispatched) { + // The low-level command runner classifies any launched process as a + // possible mutation. A probe can only mutate its disposable snapshot, + // so translate that classification at this ownership boundary. + throw new WorkspaceToolError( + error instanceof Error ? error.message : 'Programmatic preparation failed', + error instanceof WorkspaceToolError ? error.code : 'COMMAND_UNAVAILABLE', + false, + false, + ); + } + if (error instanceof WorkspaceToolError) { + if ( + error.mutationMayHaveCommitted || + error.requiresQuarantine + ) { + throw error; + } + throw new WorkspaceToolError( + error.message, + error.code, + true, + true, + ); + } + throw new WorkspaceToolError( + 'Native programmatic execution failed after dispatch', + 'COMMAND_UNAVAILABLE', + true, + true, + ); + } finally { + try { + await rm(executionDirectory, { recursive: true, force: true }); + } catch { + throw new WorkspaceToolError( + 'Native programmatic execution cleanup failed', + 'COMMAND_UNAVAILABLE', + commandDispatched, + commandDispatched, + ); + } + } + } + + private result( + request: BridgeWorkspaceProgrammaticRequest, + command: WorkspaceExecuteCommandResult, + files: ProgrammaticFileResult[], + elapsedMs: number, + pendingToolCallsPayload?: string, + artifactDelivery?: ProgrammaticResult['artifact_delivery'], + deletedFiles: string[] = [], + ): ProgrammaticResult { + return { + language: 'bash', + version: request.body.version, + // Code API masks the execution session separately from the writable + // output bucket. Sandbox results must identify the output bucket so the + // gateway can restore it to the caller-owned session after upload. + session_id: + request.body.output_session_id ?? request.body.session_id, + files, + ...(deletedFiles.length > 0 ? { deleted_files: deletedFiles } : {}), + ...(artifactDelivery ? { artifact_delivery: artifactDelivery } : {}), + ...(pendingToolCallsPayload + ? { pending_tool_calls_payload: pendingToolCallsPayload } + : {}), + run: { + stdout: command.stdout, + stderr: command.stderr, + code: command.exitCode, + signal: command.signal ?? null, + output: `${command.stdout}${command.stderr}`, + memory: null, + message: command.timedOut ? 'Execution timed out' : null, + status: command.timedOut ? 'timeout' : null, + cpu_time: null, + wall_time: elapsedMs / 1000, + }, + }; + } +} diff --git a/packages/code/src/native-sandbox.test.ts b/packages/code/src/native-sandbox.test.ts index cf790965..3252ee50 100644 --- a/packages/code/src/native-sandbox.test.ts +++ b/packages/code/src/native-sandbox.test.ts @@ -7,6 +7,7 @@ import { mkdtemp, mkdir, open, + readFile, realpath, rename, rm, @@ -25,7 +26,10 @@ import type { } from '@anthropic-ai/sandbox-runtime'; import type { ChildProcessWithoutNullStreams } from 'node:child_process'; -import { NativeSrtWorkspaceCommandSandbox } from './native-sandbox.js'; +import { + CopyOnWriteCloneUnavailableError, + NativeSrtWorkspaceCommandSandbox, +} from './native-sandbox.js'; import { restoreScratchTraversal } from './native-scratch.js'; import { WorkspaceToolError } from './workspace.js'; @@ -54,6 +58,8 @@ function fakeManager( let credentialSeenDuringWrap: string | undefined; let gitLfsRequiredSeenDuringWrap: string | undefined; let scratchSelectorSeenDuringWrap: string | undefined; + let networkSeenDuringWrap: SandboxRuntimeConfig['network'] | undefined; + let customConfigSeenDuringWrap: Partial | undefined; const manager = { isSupportedPlatform: () => true, async checkDependenciesAsync() { @@ -67,19 +73,30 @@ function fakeManager( askCallback = callback; if (options.initializeError) throw options.initializeError; }, - async wrapWithSandboxArgv(command: string) { + updateConfig(value: SandboxRuntimeConfig) { config = value; }, + async wrapWithSandboxArgv( + command: string, + _binShell?: string, + customConfig?: Partial, + ) { await options.beforeWrap?.(); - credentialSeenDuringWrap = process.env.LIBRECHAT_CODE_TEST_CREDENTIAL; + networkSeenDuringWrap = config?.network; + customConfigSeenDuringWrap = customConfig; + credentialSeenDuringWrap = + process.env.LIBRECHAT_CODE_TEST_CREDENTIAL; gitLfsRequiredSeenDuringWrap = process.env.GIT_CONFIG_VALUE_3; scratchSelectorSeenDuringWrap = process.env.CLAUDE_CODE_TMPDIR; const ambientGitEnvironment = Object.fromEntries( Object.entries(process.env).filter( - ([name, value]) => name.startsWith('GIT_CONFIG_') && value != null, + ([name, value]) => + name.startsWith('GIT_CONFIG_') && value != null, ), ); let gitEnvironment = ambientGitEnvironment; if (options.appendGitSafeDirectory) { - const index = Number(ambientGitEnvironment.GIT_CONFIG_COUNT ?? '0'); + const index = Number( + ambientGitEnvironment.GIT_CONFIG_COUNT ?? '0', + ); gitEnvironment = { ...(options.inheritedGitEnvironment ?? {}), GIT_CONFIG_COUNT: String(index + 1), @@ -130,10 +147,118 @@ function fakeManager( get scratchSelectorSeenDuringWrap() { return scratchSelectorSeenDuringWrap; }, + get networkSeenDuringWrap() { return networkSeenDuringWrap; }, + get customConfigSeenDuringWrap() { + return customConfigSeenDuringWrap; + }, }; } -test('exclusive lifecycle rejects a second workspace sharing an SRT manager', async (t) => { +for (const trustedVm of [false, true]) test(`programmatic probe denies real-workspace writes and external effects (trusted=${trustedVm})`, async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + t.after(() => rm(root, { recursive: true, force: true })); + const fake = fakeManager(); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ + workspaceRoot: root, + manager: fake.manager, + allowedDomains: ['api.example.com'], + ...(trustedVm ? { commandPolicy: { version: 1 as const, preset: 'trusted-vm' as const, + network: { outbound: 'unrestricted' as const, allowLocalBinding: true, allowAllUnixSockets: true }, + } } : {}), + }); + const dataDirectory = await sandbox.createExecutionDirectory(); + await sandbox.executeProgrammatic(request, dataDirectory, undefined, { + probe: true, + }); + assert.deepEqual(fake.customConfigSeenDuringWrap?.network, { + allowedDomains: [], + deniedDomains: [], + strictAllowlist: true, + allowUnixSockets: [], + allowAllUnixSockets: false, + allowLocalBinding: false, + }); + assert.deepEqual(fake.networkSeenDuringWrap, fake.customConfigSeenDuringWrap?.network); + assert.equal(fake.config?.network.strictAllowlist, !trustedVm); + assert.equal(fake.reset, true, 'probe proxy session must be revoked before restoring policy'); + assert.deepEqual(fake.customConfigSeenDuringWrap?.filesystem?.allowWrite, [ + await realpath(dataDirectory), + ]); + assert.equal( + fake.scratchSelectorSeenDuringWrap, + await realpath(dataDirectory), + ); + assert.ok( + fake.customConfigSeenDuringWrap?.filesystem?.denyWrite?.includes( + await realpath(root), + ), + ); + await sandbox.close(); +}); + +test('probe network cleanup failure fences executor reuse', async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-probe-cleanup-')); + t.after(() => rm(root, { recursive: true, force: true })); + const fake = fakeManager(); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ workspaceRoot: root, manager: fake.manager }); + const directory = await sandbox.createExecutionDirectory(); + const reset = fake.manager.reset; + fake.manager.reset = async () => { throw new Error('proxy shutdown failed'); }; + await assert.rejects(sandbox.executeProgrammatic(request, directory, undefined, { probe: true }), /probe network cleanup failed/); + await assert.rejects(sandbox.execute(request)); + fake.manager.reset = reset; + await sandbox.close(); +}); + +test('programmatic probes use a copy-on-write workspace without mutating the project', async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'state.txt'), 'original'); + const fake = fakeManager(); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ + workspaceRoot: root, + manager: fake.manager, + }); + t.after(() => sandbox.close()); + const executionDirectory = await sandbox.createExecutionDirectory(); + let snapshot: string; + try { + snapshot = await sandbox.createProgrammaticProbeWorkspace(executionDirectory); + } catch (error) { + if (error instanceof CopyOnWriteCloneUnavailableError) { + t.skip('host filesystem does not support copy-on-write cloning'); + return; + } + throw error; + } + await writeFile(join(snapshot, 'state.txt'), 'probe-only'); + assert.equal(await readFile(join(root, 'state.txt'), 'utf8'), 'original'); + assert.equal(await readFile(join(snapshot, 'state.txt'), 'utf8'), 'probe-only'); +}); + +test('programmatic probes do not hide clone implementation failures as unsupported filesystems', async t => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + t.after(() => rm(root, { recursive: true, force: true })); + const sandbox = new NativeSrtWorkspaceCommandSandbox({ + workspaceRoot: root, + manager: fakeManager().manager, + spawnCommand() { + throw Object.assign(new Error('spawn /bin/cp ENOENT'), { code: 'ENOENT' }); + }, + }); + t.after(() => sandbox.close()); + const executionDirectory = await sandbox.createExecutionDirectory(); + + await assert.rejects( + sandbox.createProgrammaticProbeWorkspace(executionDirectory), + (error: unknown) => + error instanceof WorkspaceToolError && + !(error instanceof CopyOnWriteCloneUnavailableError) && + error.message === 'Copy-on-write workspace clone failed unexpectedly', + ); +}); + +test('exclusive lifecycle rejects a second workspace sharing an SRT manager', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -164,15 +289,15 @@ test('exclusive lifecycle rejects a second workspace sharing an SRT manager', as assert.equal((await second.execute(request)).stdout, 'hello'); }); -test('exclusive lifecycle rejects overlapping commands and waits before resetting', async (t) => { +test('exclusive lifecycle rejects overlapping commands and waits before resetting', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); let entered!: () => void; - const wrapping = new Promise((resolve) => { + const wrapping = new Promise(resolve => { entered = resolve; }); let release!: () => void; - const gate = new Promise((resolve) => { + const gate = new Promise(resolve => { release = resolve; }); const fake = fakeManager({ @@ -191,7 +316,7 @@ test('exclusive lifecycle rejects overlapping commands and waits before resettin await assert.rejects(sandbox.execute(request), /active command/); const closing = sandbox.close(); const secondClose = sandbox.close(); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal(fake.reset, false); await assert.rejects(sandbox.prepare(), /closing/); release(); @@ -200,7 +325,7 @@ test('exclusive lifecycle rejects overlapping commands and waits before resettin assert.equal(fake.reset, true); }); -test('exclusive lifecycle retains ownership after a failed reset until cleanup succeeds', async (t) => { +test('exclusive lifecycle retains ownership after a failed reset until cleanup succeeds', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -226,16 +351,16 @@ test('exclusive lifecycle retains ownership after a failed reset until cleanup s await second.close(); }); -test('exclusive lifecycle waits for initialization before resetting the manager', async (t) => { +test('exclusive lifecycle waits for initialization before resetting the manager', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); let entered!: () => void; - const initializing = new Promise((resolve) => { + const initializing = new Promise(resolve => { entered = resolve; }); let release!: () => void; - const gate = new Promise((resolve) => { + const gate = new Promise(resolve => { release = resolve; }); fake.manager.initialize = async () => { @@ -249,7 +374,7 @@ test('exclusive lifecycle waits for initialization before resetting the manager' const preparing = sandbox.prepare(); await initializing; const closing = sandbox.close(); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal(fake.reset, false); release(); await preparing; @@ -257,7 +382,7 @@ test('exclusive lifecycle waits for initialization before resetting the manager' assert.equal(fake.reset, true); }); -test('initializes SRT with a default-deny network and scrubbed worker credentials', async (t) => { +test('initializes SRT with a default-deny network and scrubbed worker credentials', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const identity = join(tmpdir(), 'librechat-code-identity.json'); t.after(() => rm(root, { recursive: true, force: true })); @@ -301,7 +426,7 @@ test('initializes SRT with a default-deny network and scrubbed worker credential assert.ok(fake.config?.filesystem.denyRead.includes(canonicalHome)); assert.ok(fake.config?.filesystem.denyWrite.includes(canonicalIdentity)); assert.ok( - fake.config?.filesystem.denyWrite.some((path) => + fake.config?.filesystem.denyWrite.some(path => path.endsWith('/tmp/claude'), ), ); @@ -318,7 +443,7 @@ test('initializes SRT with a default-deny network and scrubbed worker credential await assert.rejects(access(scratchDirectory!)); }); -test('trusted-vm permits unmatched egress and local development sockets', async (t) => { +test('trusted-vm permits unmatched egress and local development sockets', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -351,7 +476,7 @@ test('trusted-vm permits unmatched egress and local development sockets', async ]); }); -test('provides an isolated scratch directory to commands and restores the host environment', async (t) => { +test('provides an isolated scratch directory to commands and restores the host environment', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const originalTmpdir = process.env.TMPDIR; @@ -380,7 +505,7 @@ test('provides an isolated scratch directory to commands and restores the host e await assert.rejects(access(result.stdout)); }); -test('removes scratch storage when SRT initialization fails', async (t) => { +test('removes scratch storage when SRT initialization fails', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager({ initializeError: new Error('init failed') }); @@ -396,7 +521,7 @@ test('removes scratch storage when SRT initialization fails', async (t) => { assert.equal(fake.reset, true); }); -test('rejects workspaces nested inside SRT shared scratch storage', async (t) => { +test('rejects workspaces nested inside SRT shared scratch storage', async t => { if (process.platform === 'win32') return; const sharedRoot = '/tmp/claude'; await mkdir(sharedRoot, { recursive: true }); @@ -431,17 +556,17 @@ test('rejects a workspace that contains worker scratch storage', async () => { await sandbox.close(); }); -test('keeps concurrent sandbox scratch directories independent', async (t) => { +test('keeps concurrent sandbox scratch directories independent', async t => { const firstRoot = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const secondRoot = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(firstRoot, { recursive: true, force: true })); t.after(() => rm(secondRoot, { recursive: true, force: true })); let releaseWrap!: () => void; let wrapStarted!: () => void; - const wrapStartedPromise = new Promise((resolve) => { + const wrapStartedPromise = new Promise(resolve => { wrapStarted = resolve; }); - const holdWrap = new Promise((resolve) => { + const holdWrap = new Promise(resolve => { releaseWrap = resolve; }); const firstFake = fakeManager({ @@ -478,7 +603,7 @@ test('keeps concurrent sandbox scratch directories independent', async (t) => { await secondSandbox.close(); }); -test('removes scratch storage after a command revokes traversal permissions', async (t) => { +test('removes scratch storage after a command revokes traversal permissions', async t => { if (process.platform === 'win32') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -498,7 +623,7 @@ test('removes scratch storage after a command revokes traversal permissions', as await assert.rejects(access(result.stdout)); }); -test('scratch traversal never follows a descendant replaced after inspection', async (t) => { +test('scratch traversal never follows a descendant replaced after inspection', async t => { if (process.platform === 'win32') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-scratch-race-')); const outside = await mkdtemp(join(tmpdir(), 'librechat-code-outside-')); @@ -529,7 +654,7 @@ test('scratch traversal never follows a descendant replaced after inspection', a assert.equal((await stat(outsideChild)).mode & 0o777, 0o711); }); -test('scratch traversal removes command-created Darwin ACLs', async (t) => { +test('scratch traversal removes command-created Darwin ACLs', async t => { if (process.platform !== 'darwin') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -548,7 +673,7 @@ test('scratch traversal removes command-created Darwin ACLs', async (t) => { await assert.rejects(access(result.stdout)); }); -test('scratch traversal bounds descriptors and work across a deep tree', async (t) => { +test('scratch traversal bounds descriptors and work across a deep tree', async t => { if (process.platform === 'win32') return; const root = await mkdtemp(join(tmpdir(), 'librechat-code-scratch-depth-')); t.after(() => rm(root, { recursive: true, force: true })); @@ -565,12 +690,17 @@ test('scratch traversal bounds descriptors and work across a deep tree', async ( await restoreScratchTraversal(rootHandle); - assert.equal((await stat(directories[directories.length - 1])).mode & 0o777, 0o700); + assert.equal( + (await stat(directories[directories.length - 1])).mode & 0o777, + 0o700, + ); }); -test('scratch traversal rejects trees beyond its recovery depth limit', async (t) => { +test('scratch traversal rejects trees beyond its recovery depth limit', async t => { if (process.platform === 'win32') return; - const root = await mkdtemp(join(tmpdir(), 'librechat-code-scratch-depth-limit-')); + const root = await mkdtemp( + join(tmpdir(), 'librechat-code-scratch-depth-limit-'), + ); t.after(() => rm(root, { recursive: true, force: true })); let directory = root; for (let depth = 0; depth < 129; depth += 1) { @@ -586,7 +716,7 @@ test('scratch traversal rejects trees beyond its recovery depth limit', async (t ); }); -test('does not replace scratch state while cleanup remains pending', async (t) => { +test('does not replace scratch state while cleanup remains pending', async t => { if (process.platform === 'win32') return; const workspace = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const retained = await mkdtemp(join(tmpdir(), 'librechat-code-retained-')); @@ -642,7 +772,7 @@ const windowsEnvironment = { }; for (const platform of ['darwin', 'linux', 'win32'] as const) { - test(`preserves required ${platform} environment names without allowing credentials`, async (t) => { + test(`preserves required ${platform} environment names without allowing credentials`, async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -658,23 +788,43 @@ for (const platform of ['darwin', 'linux', 'win32'] as const) { LD_PRELOAD: '/host/private.so', }; const sandbox = new NativeSrtWorkspaceCommandSandbox({ - workspaceRoot: root, platform, allowedDomains: ['github.com'], + workspaceRoot: root, + platform, + allowedDomains: ['github.com'], environment: { - ...proxyEnvironment, ...windowsEnvironment, ...credentials, + ...proxyEnvironment, + ...windowsEnvironment, + ...credentials, HtTp_PrOxY: 'http://mixed-case.invalid:8080', - PATH: '/usr/bin', LC_ALL: 'C.UTF-8', + PATH: '/usr/bin', + LC_ALL: 'C.UTF-8', }, manager: fake.manager, }); t.after(() => sandbox.close()); await sandbox.prepare(); - const denied = new Set(fake.config?.credentials?.envVars - ?.filter(({ mode }) => mode === 'deny').map(({ name }) => name)); - for (const name of [...Object.keys(proxyEnvironment), 'PATH', 'LC_ALL']) { - assert.equal(denied.has(name), false, `${name} must remain available`); + const denied = new Set( + fake.config?.credentials?.envVars + ?.filter(({ mode }) => mode === 'deny') + .map(({ name }) => name), + ); + for (const name of [ + ...Object.keys(proxyEnvironment), + 'PATH', + 'LC_ALL', + ]) { + assert.equal( + denied.has(name), + false, + `${name} must remain available`, + ); } for (const name of Object.keys(windowsEnvironment)) { - assert.equal(denied.has(name), platform !== 'win32', `${name} must be platform-specific`); + assert.equal( + denied.has(name), + platform !== 'win32', + `${name} must be platform-specific`, + ); } assert.equal(denied.has('HtTp_PrOxY'), platform !== 'win32'); for (const name of Object.keys(credentials)) { @@ -685,12 +835,14 @@ for (const platform of ['darwin', 'linux', 'win32'] as const) { }); } -test('uses SRT proxy values without restoring inherited proxies or credentials', async (t) => { +test('uses SRT proxy values without restoring inherited proxies or credentials', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const wrappedEnvironment = { - HTTP_PROXY: 'http://localhost:3128', HTTPS_PROXY: 'http://localhost:3128', - ALL_PROXY: 'http://localhost:3128', NO_PROXY: 'localhost', + HTTP_PROXY: 'http://localhost:3128', + HTTPS_PROXY: 'http://localhost:3128', + ALL_PROXY: 'http://localhost:3128', + NO_PROXY: 'localhost', }; const fake = fakeManager({ wrappedEnvironment }); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -700,14 +852,19 @@ test('uses SRT proxy values without restoring inherited proxies or credentials', }); t.after(() => sandbox.close()); const result = await sandbox.execute({ - ...request, maxOutputBytes: 256, - command: 'printf "%s|%s|%s|%s|%s" "$HTTP_PROXY" "$HTTPS_PROXY" "$ALL_PROXY" "$NO_PROXY" "${GITHUB_TOKEN-unset}"', + ...request, + maxOutputBytes: 256, + command: + 'printf "%s|%s|%s|%s|%s" "$HTTP_PROXY" "$HTTPS_PROXY" "$ALL_PROXY" "$NO_PROXY" "${GITHUB_TOKEN-unset}"', }); assert.equal(result.exitCode, 0); - assert.equal(result.stdout, `${Object.values(wrappedEnvironment).join('|')}|unset`); + assert.equal( + result.stdout, + `${Object.values(wrappedEnvironment).join('|')}|unset`, + ); }); -test('masks a host credential for only its injection host and restores the parent environment', async (t) => { +test('masks a host credential for only its injection host and restores the parent environment', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -731,7 +888,8 @@ test('masks a host credential for only its injection host and restores the paren ], async resolve() { return { - LIBRECHAT_CODE_TEST_CREDENTIAL: 'Authorization: Bearer real-secret', + LIBRECHAT_CODE_TEST_CREDENTIAL: + 'Authorization: Bearer real-secret', }; }, }, @@ -759,7 +917,7 @@ test('masks a host credential for only its injection host and restores the paren }); }); -test('serializes credential handoff across concurrent sandbox instances', async (t) => { +test('serializes credential handoff across concurrent sandbox instances', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const original = process.env.LIBRECHAT_CODE_TEST_CREDENTIAL; @@ -770,11 +928,11 @@ test('serializes credential handoff across concurrent sandbox instances', async else process.env.LIBRECHAT_CODE_TEST_CREDENTIAL = original; }); let firstEntered!: () => void; - const firstEnteredPromise = new Promise((resolve) => { + const firstEnteredPromise = new Promise(resolve => { firstEntered = resolve; }); let releaseFirst!: () => void; - const firstGate = new Promise((resolve) => { + const firstGate = new Promise(resolve => { releaseFirst = resolve; }); let secondEntered = false; @@ -817,7 +975,7 @@ test('serializes credential handoff across concurrent sandbox instances', async const secondExecution = sandbox(second.manager, 'second-secret').execute( request, ); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); assert.equal(secondEntered, false); releaseFirst(); await firstExecution; @@ -828,7 +986,7 @@ test('serializes credential handoff across concurrent sandbox instances', async assert.equal(process.env.LIBRECHAT_CODE_TEST_CREDENTIAL, undefined); }); -test('isolates Git from host-level global and system configuration', async (t) => { +test('isolates Git from host-level global and system configuration', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -844,7 +1002,7 @@ test('isolates Git from host-level global and system configuration', async (t) = assert.equal(result.stdout, '/dev/null|1'); }); -test('restores trusted Git LFS filters without reading host Git configuration', async (t) => { +test('restores trusted Git LFS filters without reading host Git configuration', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager({ @@ -884,7 +1042,7 @@ test('restores trusted Git LFS filters without reading host Git configuration', assert.ok(!denied?.includes('GIT_CONFIG_VALUE_0')); }); -test('filters environment names case-insensitively only on Windows', async (t) => { +test('filters environment names case-insensitively only on Windows', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager(); @@ -923,7 +1081,7 @@ test('filters environment names case-insensitively only on Windows', async (t) = assert.ok(!denied?.includes('git_config_count')); }); -test('fails closed when the configured POSIX shell is unavailable', async (t) => { +test('fails closed when the configured POSIX shell is unavailable', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -942,7 +1100,7 @@ test('fails closed when the configured POSIX shell is unavailable', async (t) => ); }); -test('fails closed when SRT dependencies are unavailable', async (t) => { +test('fails closed when SRT dependencies are unavailable', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const fake = fakeManager({ dependencyErrors: ['bubblewrap missing'] }); @@ -960,7 +1118,7 @@ test('fails closed when SRT dependencies are unavailable', async (t) => { ); }); -test('refuses workspace roots that expose worker home or control files', async (t) => { +test('refuses workspace roots that expose worker home or control files', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); const controlDirectory = join(root, '.control'); await mkdir(controlDirectory); @@ -985,7 +1143,7 @@ test('refuses workspace roots that expose worker home or control files', async ( ); }); -test('executes in the canonical workspace and bounds aggregate output', async (t) => { +test('executes in the canonical workspace and bounds aggregate output', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); await mkdir(join(root, 'src')); t.after(() => rm(root, { recursive: true, force: true })); @@ -1015,7 +1173,7 @@ test('executes in the canonical workspace and bounds aggregate output', async (t ); }); -test('rejects an escaping or unavailable command working directory', async (t) => { +test('rejects an escaping or unavailable command working directory', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1026,11 +1184,12 @@ test('rejects an escaping or unavailable command working directory', async (t) = await assert.rejects( sandbox.execute({ ...request, cwd: '..' }), (error: unknown) => - error instanceof WorkspaceToolError && error.code === 'INVALID_REQUEST', + error instanceof WorkspaceToolError && + error.code === 'INVALID_REQUEST', ); }); -test('terminates detached command descendants before returning', async (t) => { +test('terminates detached command descendants before returning', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1043,15 +1202,15 @@ test('terminates detached command descendants before returning', async (t) => { command: '(sleep 0.2; printf late > late.txt) >/dev/null 2>&1 &', }); assert.equal(result.exitCode, 0); - await new Promise((resolve) => setTimeout(resolve, 350)); + await new Promise(resolve => setTimeout(resolve, 350)); await assert.rejects(access(join(root, 'late.txt'))); }); -test('reports cancellation after command start as a potentially committed mutation', async (t) => { +test('reports cancellation after command start as a potentially committed mutation', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); let commandStarted!: () => void; - const commandStartedPromise = new Promise((resolve) => { + const commandStartedPromise = new Promise(resolve => { commandStarted = resolve; }); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1081,7 +1240,7 @@ test('reports cancellation after command start as a potentially committed mutati ); }); -test('closes stdin immediately when the command protocol provides no input', async (t) => { +test('closes stdin immediately when the command protocol provides no input', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const sandbox = new NativeSrtWorkspaceCommandSandbox({ @@ -1098,7 +1257,7 @@ test('closes stdin immediately when the command protocol provides no input', asy assert.equal(result.timedOut, false); }); -test('maps platform-native exit statuses into the bridge protocol range', async (t) => { +test('maps platform-native exit statuses into the bridge protocol range', async t => { const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); t.after(() => rm(root, { recursive: true, force: true })); const spawnCommand = () => { @@ -1123,7 +1282,7 @@ test('maps platform-native exit statuses into the bridge protocol range', async assert.equal(result.exitCode, 1); }); -test('cleans allocated command state exactly once on every execution exit', async (t) => { +test('cleans allocated command state exactly once on every execution exit', async t => { for (const outcome of [ 'abort-before-spawn', 'spawn-throw', @@ -1134,8 +1293,12 @@ test('cleans allocated command state exactly once on every execution exit', asyn 'wrap-throw', ] as const) { for (const cleanupThrows of [false, true]) { - await t.test(`${outcome}, cleanup throws: ${cleanupThrows}`, async (t) => { - const root = await mkdtemp(join(tmpdir(), 'librechat-code-native-')); + await t.test( + `${outcome}, cleanup throws: ${cleanupThrows}`, + async t => { + const root = await mkdtemp( + join(tmpdir(), 'librechat-code-native-'), + ); t.after(() => rm(root, { recursive: true, force: true })); const controller = new AbortController(); let cleanupCalls = 0; @@ -1143,9 +1306,11 @@ test('cleans allocated command state exactly once on every execution exit', asyn let allocated = false; const fake = fakeManager({ async beforeWrap() { - if (outcome === 'wrap-throw') throw new Error('wrap failed'); + if (outcome === 'wrap-throw') + throw new Error('wrap failed'); allocated = true; - if (outcome === 'abort-before-spawn') controller.abort(); + if (outcome === 'abort-before-spawn') + controller.abort(); }, }); fake.manager.cleanupAfterCommand = () => { @@ -1160,13 +1325,17 @@ test('cleans allocated command state exactly once on every execution exit', asyn spawnCommand() { spawnCalls += 1; assert.equal(allocated, true); - if (outcome === 'spawn-throw') throw new Error('spawn failed'); - const child = new EventEmitter() as ChildProcessWithoutNullStreams; + if (outcome === 'spawn-throw') + throw new Error('spawn failed'); + const child = + new EventEmitter() as ChildProcessWithoutNullStreams; let closeQueued = false; const close = () => { if (!closeQueued) { closeQueued = true; - queueMicrotask(() => child.emit('close', null, 'SIGKILL')); + queueMicrotask(() => + child.emit('close', null, 'SIGKILL'), + ); } return true; }; @@ -1180,7 +1349,10 @@ test('cleans allocated command state exactly once on every execution exit', asyn queueMicrotask(() => { assert.equal(cleanupCalls, 0); if (outcome === 'error') { - child.emit('error', new Error('spawn failed')); + child.emit( + 'error', + new Error('spawn failed'), + ); } else if (outcome === 'abort-after-spawn') { controller.abort(); } else if (outcome === 'close') { @@ -1196,28 +1368,43 @@ test('cleans allocated command state exactly once on every execution exit', asyn ); if (outcome === 'close' || outcome === 'timeout') { const result = await execution; - assert.equal(result.exitCode, outcome === 'close' ? 0 : null); + assert.equal( + result.exitCode, + outcome === 'close' ? 0 : null, + ); assert.equal(result.timedOut, outcome === 'timeout'); } else { - await assert.rejects(execution, (error: unknown) => + await assert.rejects( + execution, + (error: unknown) => error instanceof WorkspaceToolError && - error.code === (outcome.startsWith('abort') + error.code === + (outcome.startsWith('abort') ? 'EXECUTION_ABORTED' : 'COMMAND_UNAVAILABLE') && - error.mutationMayHaveCommitted === (outcome === 'abort-after-spawn') && + error.mutationMayHaveCommitted === + (outcome === 'abort-after-spawn') && error.requiresQuarantine === - (outcome === 'abort-after-spawn' && process.platform === 'win32'), + (outcome === 'abort-after-spawn' && + process.platform === 'win32'), ); } assert.equal( spawnCalls, - outcome === 'abort-before-spawn' || outcome === 'wrap-throw' ? 0 : 1, + outcome === 'abort-before-spawn' || + outcome === 'wrap-throw' + ? 0 + : 1, + ); + assert.equal( + cleanupCalls, + outcome === 'wrap-throw' ? 0 : 1, ); - assert.equal(cleanupCalls, outcome === 'wrap-throw' ? 0 : 1); assert.equal(allocated, false); await sandbox.close(); assert.equal(fake.reset, true); - }); + }, + ); } } }); diff --git a/packages/code/src/native-sandbox.ts b/packages/code/src/native-sandbox.ts index 171ea663..91e9832d 100644 --- a/packages/code/src/native-sandbox.ts +++ b/packages/code/src/native-sandbox.ts @@ -11,14 +11,7 @@ import { sep, } from 'node:path'; import { constants as fsConstants } from 'node:fs'; -import { - access, - mkdtemp, - open, - realpath, - rm, - stat, -} from 'node:fs/promises'; +import { access, mkdtemp, open, realpath, rm, stat } from 'node:fs/promises'; import type { FileHandle } from 'node:fs/promises'; import { SandboxManager } from '@anthropic-ai/sandbox-runtime'; @@ -141,8 +134,12 @@ interface NativeSandboxManager { cwd?: string, options?: { commandId?: string; commandText?: string }, ): Promise<{ argv: string[]; env: NodeJS.ProcessEnv }>; - annotateStderrWithSandboxFailures(commandId: string, stderr: string): string; + annotateStderrWithSandboxFailures( + commandId: string, + stderr: string, + ): string; cleanupAfterCommand(): void; + updateConfig?(config: SandboxRuntimeConfig): void; reset(): Promise; } @@ -172,6 +169,8 @@ export interface NativeSrtWorkspaceCommandSandboxOptions { platform?: NodeJS.Platform; /** Trusted shell path used by SRT on POSIX hosts. */ shellPath?: string; + /** Trusted jq path used by generated programmatic scripts. */ + jqPath?: string; /** Host-owned credentials exposed only as SRT sentinels inside the sandbox. */ maskedEnvironment?: { variables: Array<{ @@ -224,13 +223,16 @@ function deniedEnvironmentNames( platform: NodeJS.Platform, ): string[] { return Object.keys(environment) - .filter((name) => { + .filter(name => { const normalized = platform === 'win32' ? name.toUpperCase() : name; return ( normalized.startsWith('LIBRECHAT_CODE_') || (!SAFE_CHILD_ENV_NAMES.has(normalized) && !PROXY_CHILD_ENV_NAMES.has(normalized) && - !(platform === 'win32' && WINDOWS_CHILD_ENV_NAMES.has(normalized)) && + !( + platform === 'win32' && + WINDOWS_CHILD_ENV_NAMES.has(normalized) + ) && !normalized.startsWith('LC_')) ); }) @@ -244,6 +246,36 @@ function normalizedEnvironmentName( return platform === 'win32' ? name.toUpperCase() : name; } +/** Distinguishes an unsupported host filesystem from an implementation fault. */ +export class CopyOnWriteCloneUnavailableError extends WorkspaceToolError { + constructor() { + super( + 'Selected-workspace PTC requires copy-on-write filesystem cloning', + 'COMMAND_UNAVAILABLE', + ); + this.name = 'CopyOnWriteCloneUnavailableError'; + } +} + +function isCopyOnWriteUnsupported( + error: unknown, + platform: NodeJS.Platform, +): boolean { + if (platform === 'win32') return true; + if ( + typeof error === 'object' && + error !== null && + 'code' in error && + (error.code === 'ENOTSUP' || error.code === 'EOPNOTSUPP') + ) { + return true; + } + return ( + error instanceof Error && + error.message.toLowerCase().includes('operation not supported') + ); +} + export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox { readonly mutationFailuresAreAtomic = true as const; private readonly manager: NativeSandboxManager; @@ -252,6 +284,9 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox private readonly platform: NodeJS.Platform; private initialized?: Promise; private canonicalRoot?: string; + private runtimeConfig?: SandboxRuntimeConfig; + private denyReadPaths: string[] = []; + private denyWritePaths: string[] = []; private scratchDirectory?: string; private scratchHandle?: FileHandle; private execution?: Promise; @@ -288,7 +323,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); } managerOwners.set(this.manager, this); - this.initialized = this.initializeOnce().catch(async (error) => { + this.initialized = this.initializeOnce().catch(async error => { await this.manager.reset().catch(() => { this.resetFailed = true; }); @@ -314,7 +349,9 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox 'COMMAND_UNAVAILABLE', ); } - const home = await canonicalPath(this.options.homeDirectory ?? homedir()); + const home = await canonicalPath( + this.options.homeDirectory ?? homedir(), + ); if (isWithin(root, home)) { throw new WorkspaceToolError( 'Native sandbox workspace cannot contain the worker home directory', @@ -324,7 +361,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const protectedPaths = await Promise.all( (this.options.protectedPaths ?? []).map(canonicalPath), ); - if (protectedPaths.some((path) => isWithin(root, path))) { + if (protectedPaths.some(path => isWithin(root, path))) { throw new WorkspaceToolError( 'Native sandbox workspace cannot contain worker control files', 'REGISTRATION_INVALID', @@ -338,13 +375,16 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const inheritedWritablePaths = [ ...sharedScratchPaths, ...(await Promise.all( - [join(home, '.npm', '_logs'), join(home, '.claude', 'debug')].map( - canonicalPath, - ), + [ + join(home, '.npm', '_logs'), + join(home, '.claude', 'debug'), + ].map(canonicalPath), )), ]; - const deniedInheritedWritablePaths = [...new Set(inheritedWritablePaths)]; - if (deniedInheritedWritablePaths.some((path) => isWithin(path, root))) { + const deniedInheritedWritablePaths = [ + ...new Set(inheritedWritablePaths), + ]; + if (deniedInheritedWritablePaths.some(path => isWithin(path, root))) { throw new WorkspaceToolError( 'Native sandbox workspace cannot be inside an inherited writable path', 'REGISTRATION_INVALID', @@ -359,7 +399,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox } if (this.platform !== 'win32') { try { - await access(this.options.shellPath ?? '/bin/bash', fsConstants.X_OK); + await access( + this.options.shellPath ?? '/bin/bash', + fsConstants.X_OK, + ); } catch { throw new WorkspaceToolError( `Native sandbox shell is unavailable: ${this.options.shellPath ?? '/bin/bash'}`, @@ -383,35 +426,40 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); const unrestrictedNetwork = commandPolicy.network.outbound === 'unrestricted'; - const config: SandboxRuntimeConfig = { - network: { + const network: SandboxRuntimeConfig['network'] = { allowedDomains: [...(this.options.allowedDomains ?? [])], deniedDomains: [], strictAllowlist: !unrestrictedNetwork, allowAllUnixSockets: commandPolicy.network.allowAllUnixSockets, allowLocalBinding: commandPolicy.network.allowLocalBinding, ...(this.options.maskedEnvironment ? { tlsTerminate: {} } : {}), - }, + }; + const config: SandboxRuntimeConfig = { + network, filesystem: { denyRead: [ home, - ...sharedScratchPaths.filter((path) => + ...sharedScratchPaths.filter(path => deniedInheritedWritablePaths.includes(path), ), ], allowRead: [ root, - ...(canonicalScratchDirectory ? [canonicalScratchDirectory] : []), + ...(canonicalScratchDirectory + ? [canonicalScratchDirectory] + : []), ], allowWrite: [ root, - ...(canonicalScratchDirectory ? [canonicalScratchDirectory] : []), + ...(canonicalScratchDirectory + ? [canonicalScratchDirectory] + : []), ], denyWrite: [...protectedPaths, ...deniedInheritedWritablePaths], allowGitConfig: false, }, credentials: { - files: protectedPaths.map((path) => ({ + files: protectedPaths.map(path => ({ path, mode: 'deny' as const, })), @@ -424,15 +472,18 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox }, this.platform, ) - .filter((name) => { + .filter(name => { const normalized = normalizedEnvironmentName( name, this.platform, ); return ( - !Object.hasOwn(TRUSTED_GIT_ENVIRONMENT, normalized) && + !Object.hasOwn( + TRUSTED_GIT_ENVIRONMENT, + normalized, + ) && !this.options.maskedEnvironment?.variables.some( - (variable) => + variable => normalizedEnvironmentName( variable.name, this.platform, @@ -440,12 +491,16 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ) ); }) - .map((name) => ({ name, mode: 'deny' as const })), - ...(this.options.maskedEnvironment?.variables.map((variable) => ({ + .map(name => ({ name, mode: 'deny' as const })), + ...(this.options.maskedEnvironment?.variables.map( + variable => ({ ...variable, mode: 'mask' as const, - ...(variable.extract ? { onExtractNoMatch: 'error' as const } : {}), - })) ?? []), + ...(variable.extract + ? { onExtractNoMatch: 'error' as const } + : {}), + }), + ) ?? []), ], }, allowAppleEvents: false, @@ -458,6 +513,17 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox unrestrictedNetwork ? async () => true : undefined, ); this.canonicalRoot = root; + this.runtimeConfig = config; + this.denyReadPaths = [ + home, + ...sharedScratchPaths.filter(path => + deniedInheritedWritablePaths.includes(path), + ), + ]; + this.denyWritePaths = [ + ...protectedPaths, + ...deniedInheritedWritablePaths, + ]; } async execute( @@ -479,9 +545,251 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox } } + /** + * Allocate an owner-only execution directory that is already inside this + * sandbox's allowlist. The caller must remove the returned directory after + * the execution settles. It is intentionally unavailable on native Windows + * until the restricted-account TEMP directory can be opened and verified by + * the trusted parent process. + */ + async createExecutionDirectory(): Promise { + await this.initialize(); + if (!this.scratchDirectory || this.platform === 'win32') { + throw new WorkspaceToolError( + 'Native programmatic execution storage is unavailable', + 'COMMAND_UNAVAILABLE', + ); + } + return await mkdtemp(join(this.scratchDirectory, 'execution-')); + } + + /** + * Clone the current workspace into private scratch for a side-effect- + * equivalent replay probe. Platform clone flags are intentionally strict: + * silently falling back to a byte copy would make every tool-bearing run + * consume time and disk proportional to the repository size. + */ + async createProgrammaticProbeWorkspace( + executionDirectory: string, + signal?: AbortSignal, + ): Promise { + await this.initialize(); + const scratchDirectory = this.scratchDirectory; + const root = this.canonicalRoot; + let parent: string; + try { + parent = await realpath(executionDirectory); + if ( + !scratchDirectory || + !root || + !isWithin(scratchDirectory, parent) || + !(await stat(parent)).isDirectory() + ) { + throw new Error('invalid execution directory'); + } + } catch { + throw new WorkspaceToolError( + 'Programmatic execution directory is unavailable', + 'INVALID_PATH', + ); + } + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + const destination = join(parent, 'workspace'); + try { + if (this.platform === 'win32') { + throw new Error('copy-on-write cloning is unavailable on Windows'); + } + const args = + this.platform === 'darwin' + ? ['-cR', root, destination] + : ['--archive', '--reflink=always', root, destination]; + await new Promise((resolveCopy, rejectCopy) => { + const child = this.spawnCommand('/bin/cp', args, { + env: { + PATH: this.environment.PATH, + LANG: this.environment.LANG, + LC_ALL: this.environment.LC_ALL, + }, + signal, + }); + let stderr = Buffer.alloc(0); + child.stderr.on('data', (chunk: Buffer) => { + if (stderr.byteLength < 4_096) { + stderr = Buffer.concat([stderr, chunk]).subarray(0, 4_096); + } + }); + child.once('error', rejectCopy); + child.once('close', code => { + if (code === 0) resolveCopy(); + else { + rejectCopy( + new Error( + `copy-on-write clone failed (${code ?? 'signal'}): ${boundedUtf8(stderr, 4_096)}`, + ), + ); + } + }); + }); + return await realpath(destination); + } catch (error) { + await rm(destination, { recursive: true, force: true }).catch( + () => undefined, + ); + if (signal?.aborted) { + throw new WorkspaceToolError( + 'Programmatic execution aborted', + 'EXECUTION_ABORTED', + ); + } + if (isCopyOnWriteUnsupported(error, this.platform)) { + throw new CopyOnWriteCloneUnavailableError(); + } + throw new WorkspaceToolError( + 'Copy-on-write workspace clone failed unexpectedly', + 'COMMAND_UNAVAILABLE', + ); + } + } + + /** Run a generated program from a verified private execution directory. */ + async executeProgrammatic( + request: WorkspaceExecuteCommandRequest, + dataDirectory: string, + signal?: AbortSignal, + options?: { probe?: boolean; workspaceRoot?: string }, + ): Promise { + if (this.execution || this.closing) { + throw new WorkspaceToolError( + 'Native sandbox already has an active command or is closing', + 'COMMAND_UNAVAILABLE', + ); + } + await this.initialize(); + const scratchDirectory = this.scratchDirectory; + let canonicalDataDirectory: string; + let canonicalWorkspaceRoot: string | undefined; + try { + canonicalDataDirectory = await realpath(dataDirectory); + canonicalWorkspaceRoot = options?.workspaceRoot + ? await realpath(options.workspaceRoot) + : undefined; + if ( + !scratchDirectory || + !isWithin(scratchDirectory, canonicalDataDirectory) || + !(await stat(canonicalDataDirectory)).isDirectory() || + (canonicalWorkspaceRoot != null && + (!isWithin(scratchDirectory, canonicalWorkspaceRoot) || + !(await stat(canonicalWorkspaceRoot)).isDirectory())) + ) { + throw new Error('invalid execution directory'); + } + } catch { + throw new WorkspaceToolError( + 'Programmatic execution directory is unavailable', + 'INVALID_PATH', + ); + } + const execute = () => this.executeExclusive( + request, + signal, + { + LIBRECHAT_CODE_DATA_DIR: canonicalDataDirectory, + LIBRECHAT_CODE_CONTROL_PATH: join( + canonicalDataDirectory, + '_ptc_pending_result.json', + ), + LIBRECHAT_CODE_BASH_PATH: this.options.shellPath ?? '/bin/bash', + ...(this.options.jqPath + ? { LIBRECHAT_CODE_JQ_PATH: this.options.jqPath } + : {}), + PTC_HISTORY_PATH: join( + canonicalDataDirectory, + '_ptc_history.json', + ), + TMPDIR: canonicalDataDirectory, + }, + options?.probe + ? { + filesystem: { + allowRead: [ + canonicalWorkspaceRoot ?? this.canonicalRoot!, + canonicalDataDirectory, + ], + allowWrite: [ + ...(canonicalWorkspaceRoot != null + ? [canonicalWorkspaceRoot] + : []), + canonicalDataDirectory, + ], + denyRead: this.denyReadPaths, + denyWrite: [ + this.canonicalRoot!, + ...this.denyWritePaths, + ], + }, + network: { + // A probe is speculative, even on a trusted VM. + // Copy-on-write protects files, not remote mutations. + allowedDomains: [], + deniedDomains: [], + strictAllowlist: true, + allowUnixSockets: [], + allowAllUnixSockets: false, + allowLocalBinding: false, + }, + } + : undefined, + canonicalDataDirectory, + canonicalWorkspaceRoot, + ); + const execution = options?.probe ? this.withProbeNetwork(execute) : execute(); + this.execution = execution; + try { + return await execution; + } finally { + this.execution = undefined; + } + } + + private async withProbeNetwork(execute: () => Promise): Promise { + const config = this.runtimeConfig; + if (!config || !this.manager.updateConfig) { + throw new WorkspaceToolError('Native probe network isolation is unavailable', 'COMMAND_UNAVAILABLE'); + } + // SRT's proxies and Unix/local socket rules read session configuration, + // not wrapWithSandboxArgv's per-command override. + this.manager.updateConfig({ ...config, network: { + allowedDomains: [], deniedDomains: [], strictAllowlist: true, + allowUnixSockets: [], allowAllUnixSockets: false, allowLocalBinding: false, + } }); + try { + return await execute(); + } finally { + try { + // Revoke the probe's proxy endpoints and credentials before restoring + // network access. A lingering probe must never inherit the commit's + // permissive proxy session through a live updateConfig. + await this.manager.reset(); + await this.manager.initialize(config, config.network.strictAllowlist ? undefined : async () => true); + } catch { + this.resetFailed = true; + throw new WorkspaceToolError('Native probe network cleanup failed', 'COMMAND_UNAVAILABLE'); + } + } + } + private async executeExclusive( request: WorkspaceExecuteCommandRequest, signal?: AbortSignal, + trustedEnvironment?: NodeJS.ProcessEnv, + customConfig?: Partial, + sandboxScratchDirectory?: string, + workspaceRoot?: string, ): Promise { if ( !isWorkspaceToolRequest(request) || @@ -499,7 +807,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); } await this.initialize(); - const root = this.canonicalRoot!; + const root = workspaceRoot ?? this.canonicalRoot!; let cwd: string; try { cwd = await realpath(resolve(root, request.cwd ?? '.')); @@ -528,7 +836,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox { ...TRUSTED_GIT_ENVIRONMENT, ...(credentialEnvironment ?? {}), - ...this.scratchSelectorEnvironment(), + ...this.scratchSelectorEnvironment(sandboxScratchDirectory), }, () => this.manager.wrapWithSandboxArgv( @@ -536,7 +844,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox this.platform === 'win32' ? undefined : (this.options.shellPath ?? '/bin/bash'), - undefined, + customConfig, signal, cwd, { commandId, commandText: request.command }, @@ -561,7 +869,14 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox 'EXECUTION_ABORTED', ); } - return await this.runWrapped(request, wrapped, cwd, commandId, signal); + return await this.runWrapped( + request, + wrapped, + cwd, + commandId, + signal, + trustedEnvironment, + ); } finally { // A successful wrap owns command state even when no child is spawned. try { @@ -578,7 +893,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ): Promise { const previousMutation = hostEnvironmentMutationQueue; let releaseMutation!: () => void; - hostEnvironmentMutationQueue = new Promise((resolve) => { + hostEnvironmentMutationQueue = new Promise(resolve => { releaseMutation = resolve; }); await previousMutation; @@ -604,31 +919,41 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox cwd: string, commandId: string, signal?: AbortSignal, + trustedEnvironment?: NodeJS.ProcessEnv, ): Promise { const outputLimit = - request.maxOutputBytes ?? BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; + request.maxOutputBytes ?? + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; const timeoutMs = request.timeoutMs ?? BRIDGE_WORKSPACE_COMMAND_DEFAULT_TIMEOUT_MS; return await new Promise( (resolvePromise, reject) => { let child: ChildProcessWithoutNullStreams; try { - child = this.spawnCommand(wrapped.argv[0], wrapped.argv.slice(1), { + child = this.spawnCommand( + wrapped.argv[0], + wrapped.argv.slice(1), + { cwd, env: { ...wrapped.env, ...this.scratchEnvironment(), + ...trustedEnvironment, ...TRUSTED_GIT_CONFIG_ENTRIES, GIT_CONFIG_COUNT: - wrapped.env.GIT_CONFIG_COUNT ?? TRUSTED_GIT_CONFIG_COUNT, + wrapped.env.GIT_CONFIG_COUNT ?? + TRUSTED_GIT_CONFIG_COUNT, GIT_CONFIG_GLOBAL: - this.platform === 'win32' ? 'NUL' : '/dev/null', + this.platform === 'win32' + ? 'NUL' + : '/dev/null', GIT_CONFIG_NOSYSTEM: '1', }, detached: this.platform !== 'win32', shell: false, windowsHide: true, - }); + }, + ); child.stdin.end(); } catch { reject( @@ -654,10 +979,15 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const accepted = chunk.subarray(0, remaining); target.push(accepted); outputBytes += accepted.byteLength; - if (accepted.byteLength !== chunk.byteLength) truncated = true; + if (accepted.byteLength !== chunk.byteLength) + truncated = true; }; - child.stdout.on('data', (chunk: Buffer) => append(stdout, chunk)); - child.stderr.on('data', (chunk: Buffer) => append(stderr, chunk)); + child.stdout.on('data', (chunk: Buffer) => + append(stdout, chunk), + ); + child.stderr.on('data', (chunk: Buffer) => + append(stderr, chunk), + ); const abort = (): void => { if (settled) return; this.killCommandTree(child); @@ -706,7 +1036,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); return; } - const stdoutValue = boundedUtf8(Buffer.concat(stdout), outputLimit); + const stdoutValue = boundedUtf8( + Buffer.concat(stdout), + outputLimit, + ); const stderrBudget = Math.max( 0, outputLimit - Buffer.byteLength(stdoutValue), @@ -714,7 +1047,8 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox const rawStderr = Buffer.concat(stderr).toString('utf8'); let annotatedStderr = rawStderr; try { - annotatedStderr = this.manager.annotateStderrWithSandboxFailures( + annotatedStderr = + this.manager.annotateStderrWithSandboxFailures( commandId, rawStderr, ); @@ -730,12 +1064,15 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox operation: 'execute_command', workspaceId: request.workspaceId, exitCode: - timedOut || childSignal ? null : this.protocolExitCode(code), + timedOut || childSignal + ? null + : this.protocolExitCode(code), ...(childSignal ? { signal: childSignal } : {}), stdout: stdoutValue, stderr: stderrValue, truncated: - truncated || Buffer.byteLength(annotatedStderr) > stderrBudget, + truncated || + Buffer.byteLength(annotatedStderr) > stderrBudget, timedOut, }); }); @@ -775,7 +1112,7 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox ); } const canonicalTemporaryRoot = await canonicalPath(HOST_TEMPORARY_ROOT); - const sharedScratchRoot = sharedScratchPaths.find((path) => + const sharedScratchRoot = sharedScratchPaths.find(path => isWithin(path, canonicalTemporaryRoot), ); const scratchDirectory = await mkdtemp( @@ -798,7 +1135,9 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox true, ); if (((await scratchHandle.stat()).mode & 0o777) !== 0o700) { - throw new Error('Native sandbox scratch directory is not private'); + throw new Error( + 'Native sandbox scratch directory is not private', + ); } this.scratchHandle = scratchHandle; } catch (error) { @@ -829,11 +1168,13 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox : { TMPDIR: scratchDirectory }; } - private scratchSelectorEnvironment(): NodeJS.ProcessEnv { - const scratchDirectory = this.scratchDirectory; + private scratchSelectorEnvironment( + selectedDirectory = this.scratchDirectory, + ): NodeJS.ProcessEnv { + const scratchDirectory = selectedDirectory; if (!scratchDirectory) return {}; return Object.fromEntries( - SRT_SCRATCH_SELECTOR_NAMES.map((name) => [name, scratchDirectory]), + SRT_SCRATCH_SELECTOR_NAMES.map(name => [name, scratchDirectory]), ); } diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts index fd426783..08cba1c7 100644 --- a/packages/code/src/protocol.test.ts +++ b/packages/code/src/protocol.test.ts @@ -1,8 +1,11 @@ import assert from 'node:assert/strict'; import test from 'node:test'; import { + bridgeArtifactMediaType, bridgeWorkerPath, comparePortableRelativePaths, + isBridgeWorkspaceProgrammaticRequest, + isSupportedBridgeArtifactName, isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, isWorkspaceToolRequest, @@ -13,6 +16,23 @@ import type { WorkspacePreviewEditRequest, } from './protocol.js'; +test('accepts gateway directory markers as artifacts', () => { + assert.equal(isSupportedBridgeArtifactName('.dirkeep'), true); + assert.equal(isSupportedBridgeArtifactName('nested/.dirkeep'), true); + assert.equal(isSupportedBridgeArtifactName('nested/.dirkeep.exe'), false); +}); + +test('rejects caller-supplied programmatic control payloads', () => { + for (const name of ['_ptc_pending_result.json', '_PTC_PENDING_RESULT.JSON', 'nested/_ptc_pending_result.json']) { + assert.equal(isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { language: 'bash', version: '5.2.0', session_id: 'session', files: [ + { name: 'main.sh', content: 'true' }, { name, content: '{}' }, + ] }, + }), false); + } +}); + const validSingleEditRequest: WorkspaceEditFileRequest = { protocolVersion: 1, operation: 'edit_file', @@ -592,3 +612,167 @@ test('workspace capabilities allow per-workspace operation restrictions', () => false, ); }); + +test('workspace programmatic capability is closed to Bash command roots', () => { + const workspaceTools = { + protocolVersion: 1, + operations: ['execute_command'], + programmaticLanguages: ['bash'], + workspaces: [{ id: 'project-a' }], + }; + assert.equal( + isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools, + }), + true, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { ...workspaceTools, operations: ['read_file'] }, + }), + false, + ); + assert.equal( + isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { ...workspaceTools, programmaticLanguages: ['python'] }, + }), + false, + ); +}); + +test('workspace programmatic requests accept only stable input cache identities', () => { + const request = { + headers: {}, + body: { + language: 'bash', + version: '5.2', + execution_id: 'execution_1', + replay_tool_count: 2, + max_output_files: 50, + max_output_file_bytes: 10_000_000, + session_id: 'session-1', + files: [ + { name: 'main.sh', content: 'echo ready' }, + { + name: 'skills/example.txt', + id: 'file-1', + storage_session_id: 'storage-1', + input_cache_key: 'a'.repeat(64), + }, + ], + }, + }; + assert.equal(isBridgeWorkspaceProgrammaticRequest(request), true); + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + ...request, + body: { + ...request.body, + files: [request.body.files[0], { ...request.body.files[1], input_cache_key: '../cache' }], + }, + }), + false, + ); + for (const body of [ + { ...request.body, execution_id: '../execution' }, + { ...request.body, replay_tool_count: -1 }, + { ...request.body, replay_tool_count: 257 }, + { ...request.body, max_output_files: -1 }, + { ...request.body, max_output_files: 101 }, + { ...request.body, max_output_file_bytes: 0 }, + { ...request.body, max_output_file_bytes: 10 * 1024 * 1024 + 1 }, + ]) { + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ ...request, body }), + false, + ); + } +}); + +test('workspace programmatic history can use the bounded replay aggregate budget', () => { + const history = 'h'.repeat(10 * 1024 * 1024 + 1); + const body = { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [ + { name: 'main.sh', content: 'echo ready' }, + { name: '_ptc_history.json', content: history }, + ], + }; + assert.equal(isBridgeWorkspaceProgrammaticRequest({ headers: {}, body }), true); + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { + ...body, + files: [ + { name: 'main.sh', content: history }, + { name: '_ptc_history.json', content: '{}' }, + ], + }, + }), + false, + ); +}); + +test('workspace programmatic requests reject non-canonical file paths', () => { + for (const name of ['./main.sh', 'scripts//main.sh', 'scripts/./main.sh', '.']) { + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [ + { name: 'main.sh', content: 'echo ready' }, + { name, content: 'data' }, + ], + }, + }), + false, + name, + ); + } +}); + +test('workspace programmatic requests reject ancestor-descendant input conflicts', () => { + for (const names of [ + ['main.sh', 'main.sh/data.txt'], + ['main.sh', 'assets', 'assets/logo.png'], + ['main.sh', 'deep/path/file.txt', 'deep'], + ]) { + assert.equal( + isBridgeWorkspaceProgrammaticRequest({ + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: names.map(name => ({ name, content: 'data' })), + }, + }), + false, + names.join(', '), + ); + } +}); + +test('bridge artifact policy and media types match the hardened gateway contract', () => { + assert.equal(isSupportedBridgeArtifactName('reports/result.json'), true); + assert.equal(isSupportedBridgeArtifactName('preview.png'), true); + assert.equal(isSupportedBridgeArtifactName('model.bin'), false); + assert.equal(bridgeArtifactMediaType('preview.png'), 'image/png'); + assert.equal(bridgeArtifactMediaType('reports/result.json'), 'application/json'); + assert.equal(bridgeArtifactMediaType('Dockerfile'), 'application/octet-stream'); +}); diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index c1dad949..b92d21ac 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -20,9 +20,144 @@ export const BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS = 5 * 60_000; export const BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES = 256 * 1024; export const BRIDGE_WORKSPACE_COMMAND_MAX_OUTPUT_BYTES = 1024 * 1024; export const BRIDGE_WORKSPACE_COMMAND_SIGNAL_MAX_LENGTH = 32; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES = 100; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES = BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES - 2; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY = 4; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS = 30_000; + +/** Reserve a bounded share for all input/output batches, not per-file grants. */ +export function programmaticTransferReserveMs(jobTimeoutMs: number): number { + return Math.max(1, Math.floor(jobTimeoutMs / 3)); +} +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES = 10 * 1024 * 1024; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_HISTORY_BYTES = 40_000_000; +export const BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES = 100 * 1024 * 1024; /** How long Code API drains a clean rejection after Stop cancels a workspace mutation. */ export const BRIDGE_CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000; +/** + * Artifact names accepted by the hardened egress gateway. Keep this policy in + * the bridge protocol package so a remote worker can reject unsupported output + * locally instead of discovering the mismatch only after mutating a workspace. + */ +const BRIDGE_ARTIFACT_EXTENSIONS = new Set([ + '.c', '.cs', '.cpp', '.go', '.java', '.js', '.kt', '.kts', '.lua', + '.php', '.pl', '.ps1', '.py', '.r', '.rb', '.rs', '.scala', '.sh', + '.sql', '.swift', '.ts', '.jsx', '.tsx', '.groovy', + '.css', '.htm', '.html', '.less', '.sass', '.scss', '.svg', '.svelte', '.vue', + '.adoc', '.asciidoc', '.md', '.rst', '.tex', '.txt', '.wiki', + '.csv', '.json', '.bson', '.json5', '.jsonl', '.parquet', '.tsv', + '.xml', '.yaml', '.yml', + '.ics', '.ical', '.ifb', '.icalendar', + '.conf', '.env', '.gitignore', '.ini', '.properties', '.toml', + '.doc', '.docx', '.pdf', '.ppt', '.pptx', '.xls', '.xlsx', + '.odt', '.ods', '.odp', '.rtf', + '.avif', '.bmp', '.gif', '.ico', '.jpeg', '.jpg', '.png', + '.tif', '.tiff', '.webp', + '.eot', '.ttf', '.woff', '.woff2', + '.7z', '.bz2', '.gz', '.gzip', '.rar', '.tar', '.zip', + '.tf', '.tfvars', '.tfstate', '.hcl', + '.dockerfile', '.Dockerfile', '.dockerignore', + '.helmignore', '.helmfile', '.jenkinsfile', '.vagrantfile', + '.eslintrc', '.prettierrc', '.editorconfig', '.nomad', + '.bat', '.cmd', '.deb', '.log', '.rpm', '.vbs', +]); + +function portableBasename(name: string): string { + return name.slice(name.lastIndexOf('/') + 1); +} + +/** Apply the gateway's extension allowlist without importing service code. */ +export function isSupportedBridgeArtifactName(name: string): boolean { + const basename = portableBasename(name); + if (basename === '.dirkeep') return true; + const dot = basename.lastIndexOf('.'); + const extension = dot > 0 ? basename.slice(dot).toLowerCase() : ''; + const dottedBasename = `.${basename}`; + return ( + (extension !== '' && BRIDGE_ARTIFACT_EXTENSIONS.has(extension)) || + BRIDGE_ARTIFACT_EXTENSIONS.has(basename) || + BRIDGE_ARTIFACT_EXTENSIONS.has(basename.toLowerCase()) || + (extension === '' && + (BRIDGE_ARTIFACT_EXTENSIONS.has(dottedBasename) || + BRIDGE_ARTIFACT_EXTENSIONS.has(dottedBasename.toLowerCase()))) + ); +} + +const BRIDGE_ARTIFACT_MEDIA_TYPES: Readonly> = { + '.avif': 'image/avif', + '.bmp': 'image/bmp', + '.bz2': 'application/x-bzip2', + '.c': 'text/x-c', + '.conf': 'text/plain', + '.cpp': 'text/x-c++src', + '.css': 'text/css', + '.csv': 'text/csv', + '.doc': 'application/msword', + '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', + '.gif': 'image/gif', + '.gz': 'application/gzip', + '.gzip': 'application/gzip', + '.htm': 'text/html', + '.html': 'text/html', + '.ico': 'image/x-icon', + '.ics': 'text/calendar', + '.ifb': 'text/calendar', + '.ical': 'text/calendar', + '.icalendar': 'text/calendar', + '.ini': 'text/plain', + '.java': 'text/x-java-source', + '.jpeg': 'image/jpeg', + '.jpg': 'image/jpeg', + '.js': 'text/javascript', + '.json': 'application/json', + '.json5': 'application/json5', + '.jsonl': 'application/x-ndjson', + '.jsx': 'text/jsx', + '.log': 'text/plain', + '.md': 'text/markdown', + '.odt': 'application/vnd.oasis.opendocument.text', + '.ods': 'application/vnd.oasis.opendocument.spreadsheet', + '.odp': 'application/vnd.oasis.opendocument.presentation', + '.parquet': 'application/vnd.apache.parquet', + '.pdf': 'application/pdf', + '.png': 'image/png', + '.ppt': 'application/vnd.ms-powerpoint', + '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation', + '.py': 'text/x-python', + '.rst': 'text/x-rst', + '.rtf': 'application/rtf', + '.sh': 'application/x-sh', + '.sql': 'application/sql', + '.svg': 'image/svg+xml', + '.tar': 'application/x-tar', + '.tex': 'application/x-tex', + '.tif': 'image/tiff', + '.tiff': 'image/tiff', + '.toml': 'application/toml', + '.ts': 'text/typescript', + '.tsx': 'text/tsx', + '.tsv': 'text/tab-separated-values', + '.txt': 'text/plain', + '.webp': 'image/webp', + '.woff': 'font/woff', + '.woff2': 'font/woff2', + '.xls': 'application/vnd.ms-excel', + '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', + '.xml': 'application/xml', + '.yaml': 'application/yaml', + '.yml': 'application/yaml', + '.zip': 'application/zip', +}; + +/** Infer a safe response media type from an already-validated artifact name. */ +export function bridgeArtifactMediaType(name: string): string { + const basename = portableBasename(name).toLowerCase(); + const dot = basename.lastIndexOf('.'); + const extension = dot > 0 ? basename.slice(dot) : basename; + return BRIDGE_ARTIFACT_MEDIA_TYPES[extension] ?? 'application/octet-stream'; +} + export type BridgeProtocolVersion = typeof BRIDGE_PROTOCOL_VERSION; export type BridgeWorkspaceToolOperation = @@ -38,6 +173,7 @@ export type WorkspaceWriteFileMode = 'replace' | 'create'; export type WorkspaceEditFileMode = 'single' | 'batch'; export type WorkspaceEditFileFeature = 'expected_base_sha256'; export type WorkspaceListFileFeature = 'after_path'; +export type WorkspaceProgrammaticLanguage = 'bash'; export interface BridgeWorkspaceDescriptor { id: string; @@ -58,6 +194,8 @@ export interface BridgeWorkspaceToolCapabilities { editFileFeatures?: WorkspaceEditFileFeature[]; /** Omitted by workers that cannot continue a bounded file listing. */ listFileFeatures?: WorkspaceListFileFeature[]; + /** Languages that can execute PTC replay inside a selected workspace. */ + programmaticLanguages?: WorkspaceProgrammaticLanguage[]; } export interface WorkspaceReadFileRequest { @@ -153,8 +291,7 @@ interface WorkspaceEditFileRequestBase { expectedBaseSha256?: string; } -export interface WorkspaceSingleEditFileRequest - extends WorkspaceEditFileRequestBase { +export interface WorkspaceSingleEditFileRequest extends WorkspaceEditFileRequestBase { /** Legacy single-edit form. */ oldText: string; /** Legacy single-edit form. */ @@ -162,8 +299,7 @@ export interface WorkspaceSingleEditFileRequest edits?: never; } -export interface WorkspaceBatchEditFileRequest - extends WorkspaceEditFileRequestBase { +export interface WorkspaceBatchEditFileRequest extends WorkspaceEditFileRequestBase { /** Ordered exact replacements applied atomically as one file mutation. */ edits: WorkspaceTextEdit[]; oldText?: never; @@ -171,8 +307,7 @@ export interface WorkspaceBatchEditFileRequest } export type WorkspaceEditFileRequest = - | WorkspaceSingleEditFileRequest - | WorkspaceBatchEditFileRequest; + WorkspaceSingleEditFileRequest | WorkspaceBatchEditFileRequest; export interface WorkspaceTextEdit { oldText: string; @@ -195,23 +330,20 @@ interface WorkspacePreviewEditRequestBase { path: string; } -export interface WorkspaceSinglePreviewEditRequest - extends WorkspacePreviewEditRequestBase { +export interface WorkspaceSinglePreviewEditRequest extends WorkspacePreviewEditRequestBase { oldText: string; newText: string; edits?: never; } -export interface WorkspaceBatchPreviewEditRequest - extends WorkspacePreviewEditRequestBase { +export interface WorkspaceBatchPreviewEditRequest extends WorkspacePreviewEditRequestBase { edits: WorkspaceTextEdit[]; oldText?: never; newText?: never; } export type WorkspacePreviewEditRequest = - | WorkspaceSinglePreviewEditRequest - | WorkspaceBatchPreviewEditRequest; + WorkspaceSinglePreviewEditRequest | WorkspaceBatchPreviewEditRequest; export interface WorkspacePreviewEditResult { protocolVersion: BridgeProtocolVersion; @@ -392,12 +524,7 @@ const WORKSPACE_COMMAND_RESULT_KEYS = new Set([ 'truncated', 'timedOut', ]); -const WORKSPACE_SEARCH_MATCH_KEYS = new Set([ - 'path', - 'line', - 'column', - 'text', -]); +const WORKSPACE_SEARCH_MATCH_KEYS = new Set(['path', 'line', 'column', 'text']); export interface BridgeWorkerCapabilities { /** Opt-in protocol: maximum concurrently leased independent workspace roots. */ @@ -437,6 +564,8 @@ export interface BridgeWorkerRegistrationResponse { supportedWorkspaceEditFileFeatures?: WorkspaceEditFileFeature[]; /** Listing features this Code API can safely route to a capability-aware worker. */ supportedWorkspaceListFileFeatures?: WorkspaceListFileFeature[]; + /** PTC languages this Code API can safely route into a selected workspace. */ + supportedWorkspaceProgrammaticLanguages?: WorkspaceProgrammaticLanguage[]; } /** Administrator-visible liveness for a configured worker. Credentials, @@ -469,6 +598,37 @@ export interface BridgeSandboxRequest { headers: Record; } +export type BridgeProgrammaticPayloadFile = + | { name: string; content: string } + | { + name: string; + id: string; + storage_session_id: string; + input_cache_key?: string; + }; + +export interface BridgeWorkspaceProgrammaticBody { + language: 'bash'; + version: string; + /** Stable identity shared by every replay iteration of one execution. */ + execution_id?: string; + /** Declared replay tools; zero allows the worker to skip the probe pass. */ + replay_tool_count?: number; + run_timeout?: number; + transfer_timeout_ms?: number; + /** Manifest-bound upload ceiling negotiated by Code API. */ + max_output_files?: number; + /** Effective per-file upload ceiling negotiated by Code API. */ + max_output_file_bytes?: number; + files: BridgeProgrammaticPayloadFile[]; + session_id: string; + output_session_id?: string; + egress_grant?: string; +} + +export type BridgeWorkspaceProgrammaticRequest = + BridgeSandboxRequest; + export interface BridgeAssignment { workspaceLeaseSlot?: number; protocolVersion: BridgeProtocolVersion; @@ -481,7 +641,9 @@ export interface BridgeAssignment { /** Server-calculated execution budget at lease time; avoids VM clock skew. */ remainingMs?: number; runtimeSessionId?: string; - executionKind?: 'sandbox' | 'workspace_tool'; + executionKind?: 'sandbox' | 'workspace_tool' | 'workspace_programmatic'; + /** Selected workspace for workspace-scoped programmatic execution. */ + workspaceId?: string; request: BridgeSandboxRequest | WorkspaceToolRequest; } @@ -589,6 +751,140 @@ export function isValidBridgeWorkerId(workerId: string): boolean { return BRIDGE_WORKER_ID_PATTERN.test(workerId); } +export function isBridgeWorkspaceProgrammaticRequest( + value: unknown, +): value is BridgeWorkspaceProgrammaticRequest { + if (typeof value !== 'object' || value === null) return false; + const request = value as Record; + if ( + typeof request.headers !== 'object' || + request.headers === null || + !Object.values(request.headers).every( + entry => typeof entry === 'string', + ) || + typeof request.body !== 'object' || + request.body === null + ) { + return false; + } + const body = request.body as Record; + if ( + body.language !== 'bash' || + typeof body.version !== 'string' || + body.version.length === 0 || + body.version.length > BRIDGE_RUNTIME_MAX_LENGTH || + (body.execution_id !== undefined && + (typeof body.execution_id !== 'string' || + !/^[A-Za-z0-9_-]{1,128}$/.test(body.execution_id))) || + (body.replay_tool_count !== undefined && + (!Number.isSafeInteger(body.replay_tool_count) || + Number(body.replay_tool_count) < 0 || + Number(body.replay_tool_count) > 256)) || + (body.max_output_files !== undefined && + (!Number.isSafeInteger(body.max_output_files) || + Number(body.max_output_files) < 0 || + Number(body.max_output_files) > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES)) || + (body.max_output_file_bytes !== undefined && + (!Number.isSafeInteger(body.max_output_file_bytes) || + Number(body.max_output_file_bytes) < 1 || + Number(body.max_output_file_bytes) > + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES)) || + typeof body.session_id !== 'string' || + body.session_id.length === 0 || + body.session_id.length > 32_768 || + /[\0\r\n]/.test(body.session_id) || + (body.output_session_id !== undefined && + (typeof body.output_session_id !== 'string' || + body.output_session_id.length === 0 || + body.output_session_id.length > 32_768 || + /[\0\r\n]/.test(body.output_session_id))) || + (body.egress_grant !== undefined && + (typeof body.egress_grant !== 'string' || + body.egress_grant.length === 0 || + body.egress_grant.length > 256 * 1024)) || + (body.transfer_timeout_ms !== undefined && + (!Number.isSafeInteger(body.transfer_timeout_ms) || + Number(body.transfer_timeout_ms) < 1 || + Number(body.transfer_timeout_ms) > BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS)) || + (body.run_timeout !== undefined && + (!Number.isSafeInteger(body.run_timeout) || + Number(body.run_timeout) < 1 || + Number(body.run_timeout) > + BRIDGE_WORKSPACE_COMMAND_MAX_TIMEOUT_MS)) || + !Array.isArray(body.files) || + body.files.length < 1 || + body.files.length > BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES + ) { + return false; + } + let inlineBytes = 0; + const names = new Set(); + for (const rawFile of body.files) { + if (typeof rawFile !== 'object' || rawFile === null) return false; + const file = rawFile as Record; + if ( + !isSafePortableRelativePath(file.name) || + file.name === '.' || + portableBasename(file.name).toLowerCase() === '_ptc_pending_result.json' || + normalizePortableRelativePath(file.name) !== file.name || + names.has(file.name) + ) { + return false; + } + names.add(file.name); + if (typeof file.content === 'string') { + inlineBytes += Buffer.byteLength(file.content); + if ( + Object.keys(file).some( + key => key !== 'name' && key !== 'content', + ) || + Buffer.byteLength(file.content) > + (file.name === '_ptc_history.json' + ? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_HISTORY_BYTES + : BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES) + ) { + return false; + } + continue; + } + if ( + typeof file.id !== 'string' || + file.id.length === 0 || + file.id.length > 32_768 || + /[\0\r\n]/.test(file.id) || + typeof file.storage_session_id !== 'string' || + file.storage_session_id.length === 0 || + file.storage_session_id.length > 32_768 || + /[\0\r\n]/.test(file.storage_session_id) || + Object.keys(file).some( + key => + key !== 'name' && + key !== 'id' && + key !== 'storage_session_id' && + key !== 'input_cache_key', + ) || + (file.input_cache_key !== undefined && + (typeof file.input_cache_key !== 'string' || + !/^[a-f0-9]{64}$/.test(file.input_cache_key))) + ) { + return false; + } + } + for (const name of names) { + const segments = name.split('/'); + let ancestor = ''; + for (let index = 0; index < segments.length - 1; index += 1) { + ancestor = ancestor ? `${ancestor}/${segments[index]}` : segments[index]!; + if (names.has(ancestor)) return false; + } + } + return ( + names.has('main.sh') && + inlineBytes <= BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_TOTAL_BYTES + ); +} + export function isSafePortableRelativePath(value: unknown): value is string { if ( typeof value !== 'string' || @@ -602,20 +898,23 @@ export function isSafePortableRelativePath(value: unknown): value is string { ) { return false; } - return value.split('/').every((segment) => segment !== '..'); + return value.split('/').every(segment => segment !== '..'); } function normalizePortableRelativePath(value: string): string { return ( value .split('/') - .filter((segment) => segment.length > 0 && segment !== '.') + .filter(segment => segment.length > 0 && segment !== '.') .join('/') || '.' ); } /** Compare path segments in ripgrep's sorted, depth-first traversal order. */ -export function comparePortableRelativePaths(left: string, right: string): number { +export function comparePortableRelativePaths( + left: string, + right: string, +): number { const encoder = new TextEncoder(); const leftSegments = left.split('/'); const rightSegments = right.split('/'); @@ -645,9 +944,14 @@ function isWithinRequestedPath(candidate: string, requested?: string): boolean { ); } -function isValidWorkspaceEditRequest(request: Record): boolean { +function isValidWorkspaceEditRequest( + request: Record, +): boolean { const hasBatch = request.edits !== undefined; - if (hasBatch && (request.oldText !== undefined || request.newText !== undefined)) { + if ( + hasBatch && + (request.oldText !== undefined || request.newText !== undefined) + ) { return false; } const edits = hasBatch @@ -665,7 +969,10 @@ function isValidWorkspaceEditRequest(request: Record): boolean if ( typeof edit !== 'object' || edit === null || - !hasOnlyKeys(edit as Record, WORKSPACE_TEXT_EDIT_KEYS) + !hasOnlyKeys( + edit as Record, + WORKSPACE_TEXT_EDIT_KEYS, + ) ) { return false; } @@ -673,9 +980,11 @@ function isValidWorkspaceEditRequest(request: Record): boolean if ( typeof candidate.oldText !== 'string' || candidate.oldText.length === 0 || - Buffer.from(candidate.oldText).toString('utf8') !== candidate.oldText || + Buffer.from(candidate.oldText).toString('utf8') !== + candidate.oldText || typeof candidate.newText !== 'string' || - Buffer.from(candidate.newText).toString('utf8') !== candidate.newText + Buffer.from(candidate.newText).toString('utf8') !== + candidate.newText ) { return false; } @@ -698,7 +1007,7 @@ function hasOnlyKeys( value: Record, allowed: ReadonlySet, ): boolean { - return Object.keys(value).every((key) => allowed.has(key)); + return Object.keys(value).every(key => allowed.has(key)); } export function isWorkspaceToolRequest( @@ -723,7 +1032,8 @@ export function isWorkspaceToolRequest( (request.maxLines === undefined || (Number.isSafeInteger(request.maxLines) && Number(request.maxLines) >= 1 && - Number(request.maxLines) <= BRIDGE_WORKSPACE_READ_MAX_LINES)) + Number(request.maxLines) <= + BRIDGE_WORKSPACE_READ_MAX_LINES)) ); } if (request.operation === 'search_text') { @@ -743,7 +1053,8 @@ export function isWorkspaceToolRequest( (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && Number(request.maxResults) >= 1 && - Number(request.maxResults) <= BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS)) + Number(request.maxResults) <= + BRIDGE_WORKSPACE_SEARCH_MAX_RESULTS)) ); } if (request.operation === 'list_files') { @@ -753,12 +1064,14 @@ export function isWorkspaceToolRequest( isSafePortableRelativePath(request.path)) && (request.afterPath === undefined || (isSafePortableRelativePath(request.afterPath) && - normalizePortableRelativePath(request.afterPath) === request.afterPath && + normalizePortableRelativePath(request.afterPath) === + request.afterPath && isWithinRequestedPath(request.afterPath, request.path))) && (request.maxResults === undefined || (Number.isSafeInteger(request.maxResults) && Number(request.maxResults) >= 1 && - Number(request.maxResults) <= BRIDGE_WORKSPACE_LIST_MAX_RESULTS)) + Number(request.maxResults) <= + BRIDGE_WORKSPACE_LIST_MAX_RESULTS)) ); } if (request.operation === 'write_file') { @@ -799,7 +1112,8 @@ export function isWorkspaceToolRequest( !request.command.includes('\0') && new TextEncoder().encode(request.command).byteLength <= BRIDGE_WORKSPACE_COMMAND_MAX_BYTES && - (request.cwd === undefined || isSafePortableRelativePath(request.cwd)) && + (request.cwd === undefined || + isSafePortableRelativePath(request.cwd)) && (request.timeoutMs === undefined || (Number.isSafeInteger(request.timeoutMs) && Number(request.timeoutMs) >= 1 && @@ -838,13 +1152,19 @@ export function isWorkspaceToolResult( if (request.operation === 'read_file') { const startLine = request.startLine ?? 1; const maxLines = request.maxLines ?? 200; - const content = typeof result.content === 'string' ? result.content : null; + const content = + typeof result.content === 'string' ? result.content : null; const reportedLineCount = - Number.isSafeInteger(result.endLine) && Number(result.endLine) >= startLine - 1 + Number.isSafeInteger(result.endLine) && + Number(result.endLine) >= startLine - 1 ? Number(result.endLine) - startLine + 1 : -1; const actualLineCount = - content === null ? -1 : content.length === 0 ? reportedLineCount : content.split('\n').length; + content === null + ? -1 + : content.length === 0 + ? reportedLineCount + : content.split('\n').length; return ( hasOnlyKeys(result, WORKSPACE_READ_RESULT_KEYS) && result.path === request.path && @@ -899,7 +1219,10 @@ export function isWorkspaceToolResult( (enforcesPaginationContract && (normalizedPath !== path || (previousPath !== undefined && - comparePortableRelativePaths(normalizedPath, previousPath) <= 0))) + comparePortableRelativePaths( + normalizedPath, + previousPath, + ) <= 0))) ) { return false; } @@ -909,7 +1232,8 @@ export function isWorkspaceToolResult( if (!enforcesPaginationContract) { return result.nextAfterPath === undefined; } - if (result.truncated !== true) return result.nextAfterPath === undefined; + if (result.truncated !== true) + return result.nextAfterPath === undefined; return ( result.paths.length > 0 && result.nextAfterPath === result.paths[result.paths.length - 1] @@ -942,7 +1266,8 @@ export function isWorkspaceToolResult( if (request.operation === 'preview_edit') { const replacements = request.edits?.length ?? 1; - const content = typeof result.content === 'string' ? result.content : null; + const content = + typeof result.content === 'string' ? result.content : null; return ( hasOnlyKeys(result, WORKSPACE_PREVIEW_EDIT_RESULT_KEYS) && result.path === request.path && @@ -964,7 +1289,8 @@ export function isWorkspaceToolResult( const stdout = typeof result.stdout === 'string' ? result.stdout : null; const stderr = typeof result.stderr === 'string' ? result.stderr : null; const outputLimit = - request.maxOutputBytes ?? BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; + request.maxOutputBytes ?? + BRIDGE_WORKSPACE_COMMAND_DEFAULT_OUTPUT_BYTES; return ( hasOnlyKeys(result, WORKSPACE_COMMAND_RESULT_KEYS) && stdout !== null && @@ -980,7 +1306,8 @@ export function isWorkspaceToolResult( Number(result.exitCode) <= 255)) && (result.signal === undefined || (typeof result.signal === 'string' && - result.signal.length <= BRIDGE_WORKSPACE_COMMAND_SIGNAL_MAX_LENGTH && + result.signal.length <= + BRIDGE_WORKSPACE_COMMAND_SIGNAL_MAX_LENGTH && /^SIG[A-Z0-9]+$/.test(result.signal))) && typeof result.truncated === 'boolean' && typeof result.timedOut === 'boolean' && @@ -995,7 +1322,7 @@ export function isWorkspaceToolResult( return ( hasOnlyKeys(result, WORKSPACE_SEARCH_RESULT_KEYS) && result.matches.length <= maxResults && - result.matches.every((match) => { + result.matches.every(match => { if (typeof match !== 'object' || match === null) return false; const candidate = match as Record; return ( @@ -1007,7 +1334,8 @@ export function isWorkspaceToolResult( Number.isSafeInteger(candidate.column) && Number(candidate.column) >= 1 && typeof candidate.text === 'string' && - candidate.text.length <= BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH && + candidate.text.length <= + BRIDGE_WORKSPACE_SEARCH_TEXT_MAX_LENGTH && candidate.text.includes(request.query) ); }) @@ -1025,7 +1353,7 @@ export function isValidBridgeWorkspaceToolCapabilities( capabilities.operations.length < 1 || capabilities.operations.length > 7 || !capabilities.operations.every( - (operation) => + operation => operation === 'read_file' || operation === 'search_text' || operation === 'list_files' || @@ -1034,7 +1362,8 @@ export function isValidBridgeWorkspaceToolCapabilities( operation === 'edit_file' || operation === 'execute_command', ) || - new Set(capabilities.operations).size !== capabilities.operations.length || + new Set(capabilities.operations).size !== + capabilities.operations.length || !Array.isArray(capabilities.workspaces) || capabilities.workspaces.length < 1 || capabilities.workspaces.length > BRIDGE_WORKSPACE_MAX_COUNT @@ -1049,7 +1378,7 @@ export function isValidBridgeWorkspaceToolCapabilities( capabilities.writeFileModes.length > 2 || !capabilities.operations.includes('write_file') || !capabilities.writeFileModes.every( - (mode) => mode === 'replace' || mode === 'create', + mode => mode === 'replace' || mode === 'create', ) || new Set(capabilities.writeFileModes).size !== capabilities.writeFileModes.length) @@ -1065,7 +1394,7 @@ export function isValidBridgeWorkspaceToolCapabilities( (!capabilities.operations.includes('edit_file') && !capabilities.operations.includes('preview_edit')) || !capabilities.editFileModes.every( - (mode) => mode === 'single' || mode === 'batch', + mode => mode === 'single' || mode === 'batch', ) || new Set(capabilities.editFileModes).size !== capabilities.editFileModes.length) @@ -1093,13 +1422,23 @@ export function isValidBridgeWorkspaceToolCapabilities( return false; } + if ( + capabilities.programmaticLanguages !== undefined && + (!Array.isArray(capabilities.programmaticLanguages) || + capabilities.programmaticLanguages.length !== 1 || + !capabilities.operations.includes('execute_command') || + capabilities.programmaticLanguages[0] !== 'bash') + ) { + return false; + } + const workspaceIds = new Set(); - return capabilities.workspaces.every((workspace) => { + return capabilities.workspaces.every(workspace => { if (typeof workspace !== 'object' || workspace === null) return false; const descriptor = workspace as Record; if ( Object.keys(descriptor).some( - (key) => key !== 'id' && key !== 'name' && key !== 'operations', + key => key !== 'id' && key !== 'name' && key !== 'operations', ) || typeof descriptor.id !== 'string' || !isValidBridgeWorkerId(descriptor.id) || @@ -1107,17 +1446,21 @@ export function isValidBridgeWorkspaceToolCapabilities( (descriptor.name !== undefined && (typeof descriptor.name !== 'string' || descriptor.name.trim().length === 0 || - descriptor.name.length > BRIDGE_WORKSPACE_NAME_MAX_LENGTH)) || + descriptor.name.length > + BRIDGE_WORKSPACE_NAME_MAX_LENGTH)) || (descriptor.operations !== undefined && (!Array.isArray(descriptor.operations) || descriptor.operations.length < 1 || descriptor.operations.length > (capabilities.operations as unknown[]).length || descriptor.operations.some( - (operation) => - !(capabilities.operations as unknown[]).includes(operation), + operation => + !(capabilities.operations as unknown[]).includes( + operation, + ), ) || - new Set(descriptor.operations).size !== descriptor.operations.length)) + new Set(descriptor.operations).size !== + descriptor.operations.length)) ) { return false; } @@ -1139,11 +1482,12 @@ export function isValidBridgeWorkerCapabilities( typeof capabilities.statefulWorkspace === 'boolean' && typeof capabilities.sandboxProfile === 'string' && capabilities.sandboxProfile.trim().length > 0 && - capabilities.sandboxProfile.length <= BRIDGE_SANDBOX_PROFILE_MAX_LENGTH && + capabilities.sandboxProfile.length <= + BRIDGE_SANDBOX_PROFILE_MAX_LENGTH && Array.isArray(capabilities.runtimes) && capabilities.runtimes.length <= BRIDGE_RUNTIME_MAX_COUNT && capabilities.runtimes.every( - (runtime) => + runtime => typeof runtime === 'string' && runtime.length > 0 && runtime.length <= BRIDGE_RUNTIME_MAX_LENGTH, diff --git a/packages/code/src/worker-slots.test.ts b/packages/code/src/worker-slots.test.ts index ba65a3cf..6e8cc8b1 100644 --- a/packages/code/src/worker-slots.test.ts +++ b/packages/code/src/worker-slots.test.ts @@ -312,3 +312,51 @@ test('a local cleanup handoff preserves the new assignment owner and remaining b assert.equal(executed, true); assert.equal(internals.activeWorkspaceAssignments.size, 0); }); + +test('programmatic work on an independent workspace bypasses another root cleanup', 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; + }; + internals.activeWorkspaceAssignments.set('a', { + id: 'previous', + done: new Promise(() => {}), + }); + let executed = false; + internals.executeOwned = async () => { + executed = true; + assert.equal(internals.activeWorkspaceAssignments.get('b')?.id, 'next'); + }; + await worker.executeAndSettle({ + assignmentId: 'next', + executionKind: 'workspace_programmatic', + workspaceId: 'b', + remainingMs: 1_000, + request: { + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'session', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + }, + } as BridgeAssignment); + assert.equal(executed, true); + assert.equal(internals.activeWorkspaceAssignments.has('b'), false); + assert.equal(internals.activeWorkspaceAssignments.get('a')?.id, 'previous'); +}); diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 68d95b28..e26f212a 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -252,7 +252,10 @@ test('worker asks its supervisor to quarantine an ambiguous stateful runtime', a const quarantined: Array<{ sessionId: string; reason: string }> = []; const supervisor: RuntimeSupervisor = { async acquire() { - return { endpoint: 'http://127.0.0.1:3000/runtime', sessionId: 'rt-user-1' }; + return { + endpoint: 'http://127.0.0.1:3000/runtime', + sessionId: 'rt-user-1', + }; }, async reset() {}, async quarantine(sessionId, reason) { @@ -388,7 +391,10 @@ test('worker continues after an assignment-scoped settlement conflict', async () registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (init?.signal?.aborted === true) { @@ -397,15 +403,25 @@ test('worker continues after an assignment-scoped settlement conflict', async () if (url.endsWith('/lease')) { leases += 1; return new Response( - JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0, assignment }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 0, + assignment, + }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/ack')) { leaseAcknowledged = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -511,7 +527,10 @@ test('worker refreshes its registration during a long assignment', async () => { registeredAt: new Date().toISOString(), leaseTtlMs: 100, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -576,7 +595,10 @@ test('worker schedules registration freshness from request start', async () => { registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -589,7 +611,10 @@ test('worker schedules registration freshness from request start', async () => { if (url.endsWith('/cancelled')) { return new Response( JSON.stringify({ protocolVersion: 1, cancelled: false }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } return new Response( @@ -658,10 +683,13 @@ test('worker continues cancellation polling after a stalled response', async () }); } settlementAttempted = true; - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', @@ -694,6 +722,175 @@ test('worker continues cancellation polling after a stalled response', async () assert.equal(settlementAttempted, true); }); +test('worker stops an outstanding cancellation request before settling completed work', async () => { + let startCancellation!: () => void; + const cancellationStarted = new Promise((resolve) => { + startCancellation = resolve; + }); + let cancellationAborted = false; + let settlementAttempted = false; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + await cancellationStarted; + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/cancellation')) { + startCancellation(); + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => { + cancellationAborted = true; + reject(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + }); + } + settlementAttempted = true; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 1, + cancellationTransportTimeoutMs: 10_000, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'complete-while-cancellation-polling', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(cancellationAborted, true); + assert.equal(settlementAttempted, true); +}); + +test('worker aborts a retryable cancellation error body before settling completed work', async () => { + let cancellationBodyStarted!: () => void; + const cancellationStarted = new Promise((resolve) => { + cancellationBodyStarted = resolve; + }); + let cancellationBodyAborted = false; + let settlementAttempted = false; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + await cancellationStarted; + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/cancellation')) { + return new Response( + new ReadableStream({ + start(controller) { + cancellationBodyStarted(); + init?.signal?.addEventListener( + 'abort', + () => { + cancellationBodyAborted = true; + controller.error(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + }, + }), + { status: 500 }, + ); + } + settlementAttempted = true; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 1, + cancellationTransportTimeoutMs: 10_000, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'complete-during-retryable-cancellation-response', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(cancellationBodyAborted, true); + assert.equal(settlementAttempted, true); +}); + +test('worker stops its cancellation delay before settling immediately completed work', async () => { + let settlementAttempted = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 10_000, + fetchImpl: async (input) => { + const url = String(input); + if (url.endsWith('/execute')) { + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/cancellation')) { + throw new Error('cancellation transport should not start'); + } + settlementAttempted = true; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'complete-before-cancellation-polling', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(settlementAttempted, true); +}); + test('worker routes a hintless assignment to an ephemeral template session', async () => { let executeUrl = ''; let runtimeSessionHeader = ''; @@ -852,10 +1049,13 @@ test('worker preserves status for a non-JSON settlement rejection', async () => }, fetchImpl: async (input) => { if (String(input).endsWith('/execute')) { - return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ session_id: 'run-1', files: [] }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); } settlementAttempts += 1; return new Response('assignment fenced', { @@ -1032,7 +1232,10 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/execute')) { @@ -1046,18 +1249,20 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', controller.abort(); throw new TypeError('connection reset'); } - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1097,8 +1302,7 @@ test('worker preserves a definite rejection when its heartbeat fails', async () token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1118,7 +1322,10 @@ test('worker preserves a definite rejection when its heartbeat fails', async () registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/execute')) { @@ -1139,7 +1346,10 @@ test('worker preserves a definite rejection when its heartbeat fails', async () } return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1170,8 +1380,7 @@ test('worker quarantines a stateful workspace after a sandbox 5xx response', asy token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1187,7 +1396,10 @@ test('worker quarantines a stateful workspace after a sandbox 5xx response', asy settlementAttempted = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1217,8 +1429,7 @@ test('worker treats a non-JSON sandbox 4xx as a definite rejection', async () => token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1235,7 +1446,10 @@ test('worker treats a non-JSON sandbox 4xx as a definite rejection', async () => JSON.parse(String(init?.body) || '{}').status === 'rejected'; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1268,18 +1482,20 @@ test('worker quarantines a stateful workspace after the sandbox request aborts', }); } settlementAttempted = true; - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1331,13 +1547,23 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/lease')) { return new Response( - JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0, assignment }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 0, + assignment, + }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -1351,18 +1577,20 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async ); }); } - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1510,7 +1738,10 @@ test('worker subtracts lease response transit from the server budget', async () if (String(input).endsWith('/ack')) { return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } now += 50; @@ -1527,10 +1758,16 @@ test('worker subtracts lease response transit from the server budget', async () leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(0).toISOString(), remainingMs: 1_000, - request: { body: { language: 'bash' }, headers: {} }, + request: { + body: { language: 'bash' }, + headers: {}, + }, }, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1572,21 +1809,28 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/ack')) { now += 10; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/settle')) { settlementAttempts += 1; - abandonedSettlement = JSON.parse( - String(init?.body), - ) as Record; + abandonedSettlement = JSON.parse(String(init?.body)) as Record< + string, + unknown + >; if (settlementAttempts === 1) { return new Response(JSON.stringify({ error: 'unavailable' }), { status: 503, @@ -1595,7 +1839,10 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( } return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } return new Response( @@ -1611,15 +1858,24 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(0).toISOString(), remainingMs: 10, - request: { body: { language: 'bash' }, headers: {} }, + request: { + body: { language: 'bash' }, + headers: {}, + }, }, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); - await assert.rejects(worker.lease(), /expired during lease acknowledgement/); + await assert.rejects( + worker.lease(), + /expired during lease acknowledgement/, + ); assert.equal(abandonedSettlement?.status, 'rejected'); assert.ok(registrations > 0); assert.equal(settlementAttempts, 2); @@ -1653,7 +1909,10 @@ test('worker rejects an assignment after ambiguous acknowledgement delivery', as JSON.parse(String(init?.body) || '{}').status === 'rejected'; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/workers/register')) { @@ -1665,7 +1924,10 @@ test('worker rejects an assignment after ambiguous acknowledgement delivery', as registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } return new Response( @@ -1685,7 +1947,10 @@ test('worker rejects an assignment after ambiguous acknowledgement delivery', as request: { body: { language: 'bash' }, headers: {} }, }, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1702,8 +1967,7 @@ test('worker clamps rejected settlement errors to the protocol limit', async () token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1716,11 +1980,16 @@ test('worker clamps rejected settlement errors to the protocol limit', async () headers: { 'Content-Type': 'application/json' }, }); } - const settlement = JSON.parse(String(init?.body)) as { error: string }; + const settlement = JSON.parse(String(init?.body)) as { + error: string; + }; rejection = settlement.error; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1747,8 +2016,7 @@ test('worker quarantines an explicitly dirty stateful sandbox response', async ( token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1761,13 +2029,19 @@ test('worker quarantines an explicitly dirty stateful sandbox response', async ( error: 'session_workspace_dirty', message: 'restore required', }), - { status: 409, headers: { 'Content-Type': 'application/json' } }, + { + status: 409, + headers: { 'Content-Type': 'application/json' }, + }, ); } settlementAttempted = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1859,15 +2133,21 @@ test('worker uses the server-relative lease budget despite VM clock skew', async }, fetchImpl: async (input) => { if (String(input).endsWith('/execute')) { - return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ session_id: 'run-1', files: [] }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); } settlementAttempted = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1886,7 +2166,6 @@ test('worker uses the server-relative lease budget despite VM clock skew', async assert.equal(settlementAttempted, true); }); - test('worker continues after an expired assignment settlement conflict', async () => { const controller = new AbortController(); let registrations = 0; @@ -1921,18 +2200,22 @@ test('worker continues after an expired assignment settlement conflict', async ( leases += 1; return Response.json({ protocolVersion: 1, - assignment: leases === 1 - ? { - protocolVersion: 1, - assignmentId: 'assignment-expired', - workerId: 'vm-1', - incarnationId, - generation: 1, - leaseToken: 'lease-token-that-is-long-enough-for-testing', - expiresAt: new Date(Date.now() + 10_000).toISOString(), - request: { body: { language: 'bash' }, headers: {} }, - } - : undefined, + assignment: + leases === 1 + ? { + protocolVersion: 1, + assignmentId: 'assignment-expired', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + request: { + body: { language: 'bash' }, + headers: {}, + }, + } + : undefined, }); } if (url.endsWith('/execute')) { @@ -1940,7 +2223,10 @@ test('worker continues after an expired assignment settlement conflict', async ( } if (url.endsWith('/settle')) { return Response.json( - { error: 'Bridge assignment has expired', code: 'ASSIGNMENT_EXPIRED' }, + { + error: 'Bridge assignment has expired', + code: 'ASSIGNMENT_EXPIRED', + }, { status: 409 }, ); } @@ -2172,7 +2458,10 @@ test('paired worker rotates credentials throughout a long assignment', async () } if (url.endsWith('/execute')) { await new Promise((resolve) => setTimeout(resolve, 55)); - return Response.json({ session_id: 'run-long-rotation', files: [] }); + return Response.json({ + session_id: 'run-long-rotation', + files: [], + }); } return Response.json({ protocolVersion: 1, accepted: true }); }; @@ -2276,6 +2565,117 @@ test('paired worker cancels a stalled credential refresh after execution', async assert.equal(refreshAborted, true); }); +test('one concurrent caller cannot abort a credential refresh another caller still needs', async () => { + const key = createBridgeIdentity(); + const first = new AbortController(); + const second = new AbortController(); + let releaseRefresh!: () => void; + let refreshStarted!: () => void; + const started = new Promise((resolve) => { + refreshStarted = resolve; + }); + const released = new Promise((resolve) => { + releaseRefresh = resolve; + }); + let transportAborted = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-shared-refresh', + expiresAt: new Date(Date.now() + 5).toISOString(), + }, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + assert.match(String(input), /credentials\/refresh$/); + refreshStarted(); + init?.signal?.addEventListener('abort', () => { + transportAborted = true; + }); + await released; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-shared-refresh-value', + expiresAt: new Date(Date.now() + 120_000).toISOString(), + }); + }, + }); + + const firstRefresh = worker.refreshCredential(first.signal); + await started; + const secondRefresh = worker.refreshCredential(second.signal); + first.abort(); + + await assert.rejects(firstRefresh, { name: 'AbortError' }); + assert.equal(transportAborted, false); + releaseRefresh(); + await secondRefresh; + assert.equal(transportAborted, false); +}); + +test('a new caller starts a fresh credential refresh after the last waiter aborts', async () => { + const key = createBridgeIdentity(); + const first = new AbortController(); + let refreshCount = 0; + let firstRefreshStarted!: () => void; + const started = new Promise((resolve) => { + firstRefreshStarted = resolve; + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-replacement-refresh', + expiresAt: new Date(Date.now() + 5).toISOString(), + }, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (_input, init) => { + refreshCount += 1; + if (refreshCount === 1) { + firstRefreshStarted(); + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-replacement-refresh', + expiresAt: new Date(Date.now() + 120_000).toISOString(), + }); + }, + }); + + const abandoned = worker.refreshCredential(first.signal, Date.now() + 1_000); + await started; + first.abort(); + await assert.rejects(abandoned, { name: 'AbortError' }); + await worker.refreshCredential(undefined, Date.now() + 1_000); + + assert.equal(refreshCount, 2); +}); + test('paired worker refreshes conservatively before server clock calibration', async () => { const key = createBridgeIdentity(); let refreshCount = 0; @@ -2347,8 +2747,7 @@ test('paired worker charges initial credential refresh against the assignment de sandboxStarted = true; } if (url.endsWith('/settle')) { - rejected = - JSON.parse(String(init?.body)).status === 'rejected'; + rejected = JSON.parse(String(init?.body)).status === 'rejected'; } return Response.json({ protocolVersion: 1, @@ -2389,8 +2788,7 @@ test('paired worker rechecks the deadline after request serialization', async () codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', incarnationId, - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', identity: { privateKey: key.privateKey, credential: 'credential-valid-during-serialization', @@ -2408,8 +2806,7 @@ test('paired worker rechecks the deadline after request serialization', async () sandboxStarted = true; } if (url.endsWith('/settle')) { - rejected = - JSON.parse(String(init?.body)).status === 'rejected'; + rejected = JSON.parse(String(init?.body)).status === 'rejected'; } return Response.json({ protocolVersion: 1, @@ -2471,8 +2868,7 @@ test('paired worker keeps endpoint validation failures known-clean', async () => const url = String(input); if (url.endsWith('/execute')) sandboxStarted = true; if (url.endsWith('/settle')) { - rejected = - JSON.parse(String(init?.body)).status === 'rejected'; + rejected = JSON.parse(String(init?.body)).status === 'rejected'; } return Response.json({ protocolVersion: 1, @@ -2758,11 +3154,13 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn } if (url.endsWith('/execute')) { await refreshStartedPromise; - return Response.json({ session_id: 'run-rotation-race', files: [] }); + return Response.json({ + session_id: 'run-rotation-race', + files: [], + }); } - settleAuthorization = ( - init?.headers as Record - ).Authorization; + settleAuthorization = (init?.headers as Record) + .Authorization; return Response.json({ protocolVersion: 1, accepted: true }); }; const worker = new BridgeWorker({ @@ -2799,8 +3197,41 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn ); }); +test('settlement does not drain another lane credential renewal', async () => { + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', token: 'fixture', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', runtimes: ['bash'] }, + fetchImpl: async (input) => String(input).endsWith('/execute') + ? Response.json({ session_id: 'independent-lane', files: [] }) + : Response.json({ protocolVersion: 1, accepted: true }), + }); + // A different lane owns this pending renewal. The settling lane has no + // maintenance waiter and must not consume its own lease on that promise. + Object.assign(worker, { credentialInFlight: { + promise: new Promise(() => {}), controller: new AbortController(), waiters: 1, + }, refreshCredential: async () => {} }); + const startedAt = Date.now(); + await worker.executeAndSettle({ + protocolVersion: 1, assignmentId: 'independent-lane', workerId: 'vm-1', + incarnationId, generation: 5, leaseToken: 'independent-lane-lease-token', + expiresAt: new Date(Date.now() + 600_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + assert.ok(Date.now() - startedAt < 500, 'unrelated renewal must not add a one-second drain'); +}); + test('reconnect delay uses bounded exponential jitter', () => { - assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 0), 500); - assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 1), 1_000); - assert.equal(reconnectDelayMs(10, 1_000, 30_000, () => 1), 30_000); + assert.equal( + reconnectDelayMs(0, 1_000, 30_000, () => 0), + 500, + ); + assert.equal( + reconnectDelayMs(0, 1_000, 30_000, () => 1), + 1_000, + ); + assert.equal( + reconnectDelayMs(10, 1_000, 30_000, () => 1), + 30_000, + ); }); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 19f219da..ffe1b350 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -5,6 +5,7 @@ import { BRIDGE_PROTOCOL_VERSION, BridgeProtocolError, bridgeWorkerPath, + isBridgeWorkspaceProgrammaticRequest, isWorkspaceToolResult, } from './protocol.js'; import { EndpointRuntimeSupervisor } from './runtime.js'; @@ -21,6 +22,7 @@ import type { BridgeWorkerCredentialResponse, BridgeWorkerRegistrationResponse, BridgeWorkspaceToolOperation, + BridgeWorkspaceProgrammaticRequest, } from './protocol.js'; import type { RuntimeLease, RuntimeSupervisor } from './runtime.js'; import type { WorkspaceToolExecutor } from './workspace.js'; @@ -35,6 +37,18 @@ export interface BridgeWorkerOptions { runtimeSupervisor?: RuntimeSupervisor; capabilities: BridgeWorkerCapabilities; workspaceTools?: WorkspaceToolExecutor; + workspaceProgrammatic?: { + /** + * True when a WorkspaceToolError without mutation uncertainty proves the + * selected workspace was not changed. + */ + mutationFailuresAreAtomic?: true; + executeProgrammatic( + workspaceId: string, + request: BridgeWorkspaceProgrammaticRequest, + signal?: AbortSignal, + ): Promise; + }; workspaceMutationQuarantine?: WorkspaceMutationQuarantine; /** Required per-root durable guards when opting into concurrent workspace leases. */ workspaceQuarantines?: ReadonlyMap; @@ -90,6 +104,7 @@ const DEFAULT_REGISTRATION_TRANSPORT_TIMEOUT_MS = 10_000; const DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS = 10_000; const DEFAULT_CANCELLATION_POLL_INTERVAL_MS = 500; const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; +const CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS = 1_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; const REGISTRATION_RETRY_DELAY_MS = 100; const CREDENTIAL_REFRESH_RETRY_DELAY_MS = 100; @@ -149,21 +164,29 @@ function workspaceCapabilitiesMatch( advertised.writeFileModes?.length === executor.writeFileModes?.length && (advertised.writeFileModes?.every( (mode, index) => mode === executor.writeFileModes?.[index], - ) ?? executor.writeFileModes == null) && + ) ?? + executor.writeFileModes == null) && advertised.editFileModes?.length === executor.editFileModes?.length && (advertised.editFileModes?.every( (mode, index) => mode === executor.editFileModes?.[index], - ) ?? executor.editFileModes == null) && - advertised.editFileFeatures?.length === - executor.editFileFeatures?.length && + ) ?? + executor.editFileModes == null) && + advertised.editFileFeatures?.length === executor.editFileFeatures?.length && (advertised.editFileFeatures?.every( (feature, index) => feature === executor.editFileFeatures?.[index], - ) ?? executor.editFileFeatures == null) && - advertised.listFileFeatures?.length === - executor.listFileFeatures?.length && + ) ?? + executor.editFileFeatures == null) && + advertised.listFileFeatures?.length === executor.listFileFeatures?.length && (advertised.listFileFeatures?.every( (feature, index) => feature === executor.listFileFeatures?.[index], - ) ?? executor.listFileFeatures == null) && + ) ?? + executor.listFileFeatures == null) && + advertised.programmaticLanguages?.length === + executor.programmaticLanguages?.length && + (advertised.programmaticLanguages?.every( + (language, index) => language === executor.programmaticLanguages?.[index], + ) ?? + executor.programmaticLanguages == null) && advertised.workspaces.length === executor.workspaces.length && advertised.workspaces.every( (workspace, index) => @@ -175,7 +198,8 @@ function workspaceCapabilitiesMatch( (operation, operationIndex) => operation === executor.workspaces[index]?.operations?.[operationIndex], - ) ?? executor.workspaces[index]?.operations == null), + ) ?? + executor.workspaces[index]?.operations == null), ) ); } @@ -187,8 +211,7 @@ function registrationCompatibleCapabilities( if ( workspaceTools == null || (workspaceTools.operations.every( - (operation) => - operation === 'read_file' || operation === 'search_text', + (operation) => operation === 'read_file' || operation === 'search_text', ) && workspaceTools.workspaces.every( (workspace) => workspace.operations == null, @@ -197,8 +220,7 @@ function registrationCompatibleCapabilities( return capabilities; } const operations = workspaceTools.operations.filter( - (operation) => - operation === 'read_file' || operation === 'search_text', + (operation) => operation === 'read_file' || operation === 'search_text', ); if (operations.length === 0) { const { workspaceTools: _workspaceTools, ...compatible } = capabilities; @@ -207,7 +229,9 @@ function registrationCompatibleCapabilities( const workspaces = workspaceTools.workspaces.flatMap((workspace) => { if ( workspace.operations != null && - !operations.every((operation) => workspace.operations?.includes(operation)) + !operations.every((operation) => + workspace.operations?.includes(operation), + ) ) { return []; } @@ -223,6 +247,7 @@ function registrationCompatibleCapabilities( editFileModes: _editFileModes, editFileFeatures: _editFileFeatures, listFileFeatures: _listFileFeatures, + programmaticLanguages: _programmaticLanguages, ...compatibleWorkspaceTools } = workspaceTools; return { @@ -302,11 +327,16 @@ function supportedWorkspaceCapabilities( const listFileFeatures = desired.listFileFeatures?.filter((feature) => registration.supportedWorkspaceListFileFeatures?.includes(feature), ); + const programmaticLanguages = desired.programmaticLanguages?.filter( + (language) => + registration.supportedWorkspaceProgrammaticLanguages?.includes(language), + ); const { writeFileModes: _writeFileModes, editFileModes: _editFileModes, editFileFeatures: _editFileFeatures, listFileFeatures: _listFileFeatures, + programmaticLanguages: _programmaticLanguages, ...compatibleDesired } = desired; return { @@ -327,6 +357,10 @@ function supportedWorkspaceCapabilities( ...(operations.includes('list_files') && listFileFeatures?.length ? { listFileFeatures } : {}), + ...(operations.includes('execute_command') && + programmaticLanguages?.length + ? { programmaticLanguages } + : {}), }, }; } @@ -362,7 +396,11 @@ export class BridgeWorker { private negotiatedWorkspaceSlots = 1; private concurrentRunning = false; private registrationInFlight?: Promise; - private credentialInFlight?: Promise; + private credentialInFlight?: { + promise: Promise; + controller: AbortController; + waiters: number; + }; private serverClockOffsetMs = MAX_PROOF_CLOCK_SKEW_MS; constructor(private readonly options: BridgeWorkerOptions) { @@ -406,6 +444,17 @@ export class BridgeWorker { 'Workspace tool capabilities require a matching executor', ); } + if ( + (options.workspaceProgrammatic != null) !== + (options.capabilities.workspaceTools?.programmaticLanguages?.includes( + 'bash', + ) === + true) + ) { + throw new BridgeProtocolError( + 'Workspace programmatic capability requires a matching executor', + ); + } if ( options.capabilities.workspaceTools?.operations.some( (operation) => @@ -519,7 +568,9 @@ export class BridgeWorker { if (signal?.aborted) { abortRegistration(); } else { - signal?.addEventListener('abort', abortRegistration, { once: true }); + signal?.addEventListener('abort', abortRegistration, { + once: true, + }); } const timeoutMs = Math.min( Math.max(1, this.registrationTtlMs - 1), @@ -541,7 +592,10 @@ export class BridgeWorker { workerId: this.options.workerId, incarnationId: this.incarnationId, capabilities: this.maintenanceOnly - ? { ...capabilities, requiresReadyConfirmation: true } + ? { + ...capabilities, + requiresReadyConfirmation: true, + } : capabilities, }, registrationController.signal, @@ -675,7 +729,9 @@ export class BridgeWorker { } await this.runtimeSupervisor.reset(runtimeSessionId, signal); await this.timedRequest( - `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + `${this.codeApiUrl}${bridgeWorkerPath( + this.options.workerId, + )}/workspaces/reset`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, incarnationId: this.incarnationId, @@ -711,7 +767,9 @@ export class BridgeWorker { // machine-local guard before the remote fence can be removed. await guard.assertAvailable(); await this.timedRequest( - `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + `${this.codeApiUrl}${bridgeWorkerPath( + this.options.workerId, + )}/workspaces/reset`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, incarnationId: this.incarnationId, @@ -1017,20 +1075,58 @@ export class BridgeWorker { transportTimeoutMs = Number.POSITIVE_INFINITY, ): Promise { while (this.credentialInFlight) { - await this.credentialInFlight; + await this.waitForCredentialRefresh(this.credentialInFlight, signal); // A longer-lived caller may still need another refresh after this one. } + const controller = new AbortController(); const pending = this.refreshCredentialOwned( - signal, + controller.signal, validThroughMs, transportTimeoutMs, ); - this.credentialInFlight = pending; + const entry = { promise: pending, controller, waiters: 0 }; + this.credentialInFlight = entry; + void pending.then( + () => { + if (this.credentialInFlight === entry) + this.credentialInFlight = undefined; + }, + () => { + if (this.credentialInFlight === entry) + this.credentialInFlight = undefined; + }, + ); + await this.waitForCredentialRefresh(entry, signal); + } + + private async waitForCredentialRefresh( + entry: NonNullable, + signal?: AbortSignal, + ): Promise { + entry.waiters += 1; + let removeAbortListener = (): void => {}; + const aborted = new Promise((_, reject) => { + if (signal == null) return; + const abort = (): void => + reject( + signal.reason instanceof Error + ? signal.reason + : new DOMException('aborted', 'AbortError'), + ); + removeAbortListener = (): void => + signal.removeEventListener('abort', abort); + signal.addEventListener('abort', abort, { once: true }); + if (signal.aborted) abort(); + }); try { - await pending; + await Promise.race([entry.promise, aborted]); } finally { - if (this.credentialInFlight === pending) + removeAbortListener(); + entry.waiters -= 1; + if (entry.waiters === 0 && this.credentialInFlight === entry) { this.credentialInFlight = undefined; + entry.controller.abort(); + } } } @@ -1084,7 +1180,7 @@ export class BridgeWorker { assignment: BridgeAssignment, stopSignal: AbortSignal, serverClockOffsetMs: number, - requestSignal?: AbortSignal, + maintenance: { refresh?: Promise }, ): Promise { const identity = this.options.identity; if (identity == null) return; @@ -1102,18 +1198,18 @@ export class BridgeWorker { await abortableDelay(waitMs, stopSignal); if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return; try { - await this.refreshCredential( - requestSignal, + maintenance.refresh = this.refreshCredential( + stopSignal, Date.now() + serverClockOffsetMs + refreshWindowMs, ); + await maintenance.refresh; } catch (error) { if (stopSignal.aborted) return; const terminal = error instanceof BridgeProtocolError && (error.status === 401 || error.status === 403); const credentialRemainingMs = - Date.parse(identity.expiresAt) - - (Date.now() + serverClockOffsetMs); + Date.parse(identity.expiresAt) - (Date.now() + serverClockOffsetMs); if (terminal || credentialRemainingMs <= 0) throw error; await abortableDelay( Math.min( @@ -1122,6 +1218,8 @@ export class BridgeWorker { ), stopSignal, ); + } finally { + maintenance.refresh = undefined; } } } @@ -1130,11 +1228,7 @@ export class BridgeWorker { assignment: BridgeAssignment, signal?: AbortSignal, ): Promise { - const root = - assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) - ? assignment.request.workspaceId - : undefined; + const root = this.assignmentWorkspaceId(assignment); const waitingAt = Date.now(); while (root != null && this.activeWorkspaceAssignments.has(root)) { const active = this.activeWorkspaceAssignments.get(root)!; @@ -1212,14 +1306,31 @@ export class BridgeWorker { private workspaceGuard( assignment: BridgeAssignment, ): WorkspaceMutationQuarantine | undefined { - return assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) - ? (this.options.workspaceQuarantines?.get( - assignment.request.workspaceId, - ) ?? this.options.workspaceMutationQuarantine) + const workspaceId = this.assignmentWorkspaceId(assignment); + return workspaceId != null + ? (this.options.workspaceQuarantines?.get(workspaceId) ?? + this.options.workspaceMutationQuarantine) : this.options.workspaceMutationQuarantine; } + private assignmentWorkspaceId( + assignment: BridgeAssignment, + ): string | undefined { + if ( + assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) + ) { + return assignment.request.workspaceId; + } + if ( + assignment.executionKind === 'workspace_programmatic' && + typeof assignment.workspaceId === 'string' + ) { + return assignment.workspaceId; + } + return undefined; + } + private async executeOwned( assignment: BridgeAssignment, signal?: AbortSignal, @@ -1305,6 +1416,7 @@ export class BridgeWorker { ); let credentialMaintenanceError: unknown; let credentialMaintenance: Promise | undefined; + const ownCredentialMaintenance: { refresh?: Promise } = {}; let settlement: BridgeSettlement; let ambiguousSandboxError: unknown; let ambiguousWorkspaceMutationError: unknown; @@ -1321,7 +1433,7 @@ export class BridgeWorker { assignment, credentialController.signal, serverClockOffsetMs, - signal, + ownCredentialMaintenance, ).catch((error) => { credentialMaintenanceError = error; executionController.abort(); @@ -1472,6 +1584,76 @@ export class BridgeWorker { 'Bridge assignment expired during workspace execution', ); } + } else if (assignment.executionKind === 'workspace_programmatic') { + const workspaceId = assignment.workspaceId; + if ( + workspaceId == null || + this.options.workspaceProgrammatic == null || + !isBridgeWorkspaceProgrammaticRequest(assignment.request) + ) { + throw new BridgeProtocolError( + 'Worker does not provide valid selected-workspace programmatic execution', + ); + } + try { + if (this.quarantinedWorkspaces.has(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; + const workspace = advertised?.workspaces.find( + (candidate) => candidate.id === workspaceId, + ); + if ( + workspace == null || + !advertised?.operations.includes('execute_command') || + (workspace.operations != null && + !workspace.operations.includes('execute_command')) || + !advertised.programmaticLanguages?.includes('bash') + ) { + throw new BridgeProtocolError( + 'Selected-workspace programmatic execution is not advertised', + ); + } + this.mutationGuardArmed = true; + try { + this.armedWorkspaces.add(workspaceId); + await guard!.arm( + 'Workspace programmatic execution is pending settlement', + assignment.assignmentId, + ); + workspaceMutationArmed = true; + } catch (error) { + this.mutationGuardArmed = false; + throw new BridgeWorkspaceQuarantinedError( + 'Workspace mutation quarantine could not be armed before execution', + error, + ); + } + payload = await this.options.workspaceProgrammatic.executeProgrammatic( + workspaceId, + assignment.request, + executionController.signal, + ); + workspaceMutationApplied = true; + if (executionController.signal.aborted) { + throw ( + executionController.signal.reason ?? + new DOMException('aborted', 'AbortError') + ); + } + if (Date.now() >= localDeadlineAtMs) { + throw new BridgeProtocolError( + 'Bridge assignment expired during programmatic execution', + ); + } } else { runtimeLease = await this.runtimeSupervisor.acquire( assignment, @@ -1557,14 +1739,22 @@ export class BridgeWorker { ) { workspaceMutationGuardError = error; } + const knownAtomicWorkspaceToolFailure = + assignment.executionKind === 'workspace_tool' && + error instanceof WorkspaceToolError && + this.options.workspaceTools?.mutationFailuresAreAtomic === true && + !error.requiresQuarantine; + const knownAtomicProgrammaticFailure = + assignment.executionKind === 'workspace_programmatic' && + error instanceof WorkspaceToolError && + this.options.workspaceProgrammatic?.mutationFailuresAreAtomic === + true && + !error.requiresQuarantine; if ( workspaceMutationApplied || (workspaceMutationArmed && - !( - error instanceof WorkspaceToolError && - this.options.workspaceTools?.mutationFailuresAreAtomic === true && - !error.requiresQuarantine - )) + !knownAtomicWorkspaceToolFailure && + !knownAtomicProgrammaticFailure) ) { ambiguousWorkspaceMutationError = error; } @@ -1581,7 +1771,8 @@ export class BridgeWorker { leaseToken: assignment.leaseToken, incarnationId: this.incarnationId, status: 'rejected', - ...(assignment.executionKind === 'workspace_tool' && + ...((assignment.executionKind === 'workspace_tool' || + assignment.executionKind === 'workspace_programmatic') && error instanceof WorkspaceToolError ? { errorCode: error.code } : {}), @@ -1595,12 +1786,30 @@ export class BridgeWorker { clearTimeout(deadlineTimer); cancellationController.abort(); await cancellationWatcher; + // Only drain renewal joined by this assignment, never an unrelated lane's + // refresh. Leave settlement time inside the original assignment budget. + const credentialInFlight = ownCredentialMaintenance.refresh; + if (credentialInFlight != null && !credentialController.signal.aborted) { + let drainTimer: ReturnType | undefined; + await Promise.race([ + credentialInFlight.catch(() => undefined), + new Promise((resolve) => { + drainTimer = setTimeout( + resolve, + Math.min(CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS, + Math.max(0, Date.parse(assignment.expiresAt) - serverClockOffsetMs - Date.now() - 5_000)), + ); + }), + ]); + if (drainTimer != null) clearTimeout(drainTimer); + } credentialController.abort(); await credentialMaintenance; try { if (workspaceMutationGuardError != null) throw workspaceMutationGuardError; if (ambiguousWorkspaceMutationError != null) { + this.options.onError?.(ambiguousWorkspaceMutationError); throw await this.quarantineWorkspace( undefined, 'Worker stopped after a workspace mutation completed without a fulfilled settlement', @@ -1693,12 +1902,8 @@ export class BridgeWorker { clearTimeout(timer); } } - if ( - assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) - ) { - this.armedWorkspaces.delete(assignment.request.workspaceId); - } + const workspaceId = this.assignmentWorkspaceId(assignment); + if (workspaceId != null) this.armedWorkspaces.delete(workspaceId); this.mutationGuardArmed = false; } catch (error) { throw new BridgeWorkspaceQuarantinedError( @@ -1799,7 +2004,9 @@ export class BridgeWorker { return await lease.execute({ body, headers, signal }); } if (lease.endpoint == null) { - throw new BridgeProtocolError('Runtime lease does not provide an execution transport'); + throw new BridgeProtocolError( + 'Runtime lease does not provide an execution transport', + ); } const endpoint = lease.endpoint.replace(/\/+$/, ''); const response = await this.fetchImpl(`${endpoint}/execute`, { @@ -1957,11 +2164,12 @@ export class BridgeWorker { const fulfilledWorkspaceMutation = workspaceMutationApplied && settlement.status === 'fulfilled' && - assignment.executionKind === 'workspace_tool' && - isWorkspaceToolRequest(assignment.request) && - (assignment.request.operation === 'write_file' || - assignment.request.operation === 'edit_file' || - assignment.request.operation === 'execute_command'); + (assignment.executionKind === 'workspace_programmatic' || + (assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) && + (assignment.request.operation === 'write_file' || + assignment.request.operation === 'edit_file' || + assignment.request.operation === 'execute_command'))); if (signal?.aborted === true) { if (assignment.runtimeSessionId != null || fulfilledWorkspaceMutation) { throw await this.quarantineWorkspace( @@ -2055,17 +2263,23 @@ export class BridgeWorker { signal: AbortSignal, ): Promise { while (!signal.aborted && !executionController.signal.aborted) { - await this.delay( - Math.max( - 1, - this.options.cancellationPollIntervalMs ?? - DEFAULT_CANCELLATION_POLL_INTERVAL_MS, - ), - signal, - ); + try { + await this.delay( + Math.max( + 1, + this.options.cancellationPollIntervalMs ?? + DEFAULT_CANCELLATION_POLL_INTERVAL_MS, + ), + signal, + ); + } catch (error) { + if (signal.aborted || executionController.signal.aborted) return; + throw error; + } if (signal.aborted || executionController.signal.aborted) return; const pollController = new AbortController(); const abortPoll = (): void => pollController.abort(); + signal.addEventListener('abort', abortPoll, { once: true }); executionController.signal.addEventListener('abort', abortPoll, { once: true, }); @@ -2085,6 +2299,15 @@ export class BridgeWorker { incarnationId: this.incarnationId, }, pollController.signal, + (response) => { + // Once response headers arrive, drain the bounded body before a + // successful execution can settle. Otherwise a cancellation=true + // response racing command completion can be discarded. The + // transport timer and execution signal still cap the drain. + if (response.ok || response.status === 404) { + signal.removeEventListener('abort', abortPoll); + } + }, ); if (response.cancelled) { executionController.abort(); @@ -2098,6 +2321,7 @@ export class BridgeWorker { if (signal.aborted) return; } finally { clearTimeout(timeout); + signal.removeEventListener('abort', abortPoll); executionController.signal.removeEventListener('abort', abortPoll); } } @@ -2107,6 +2331,7 @@ export class BridgeWorker { url: string, body: object, signal?: AbortSignal, + onResponseHeaders?: (response: Response) => void, ): Promise { const requestBody = JSON.stringify(body); const response = await this.fetchImpl(url, { @@ -2118,6 +2343,7 @@ export class BridgeWorker { body: requestBody, signal, }); + onResponseHeaders?.(response); let payload: unknown; try { payload = await response.json(); diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index 8f9fcd45..f6205cd4 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -950,6 +950,332 @@ test('worker executes a workspace tool assignment locally without acquiring a sa }); }); +test('worker executes programmatic Bash in the selected workspace and preserves its fence', async () => { + const programmaticRequests: object[] = []; + const quarantineEvents: string[] = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + async executeProgrammatic(workspaceId, request) { + programmaticRequests.push({ workspaceId, request }); + return { + session_id: 'session-1', + language: 'bash', + version: '5.2', + files: [], + run: { stdout: 'ready\n', stderr: '', code: 0, signal: null }, + }; + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine( + (reason) => quarantineEvents.push(`quarantine:${reason}`), + (reason) => quarantineEvents.push(`arm:${reason}`), + () => quarantineEvents.push('clear'), + ), + ], + ]), + fetchImpl: async () => Response.json({ protocolVersion: 1, accepted: true }), + }); + const request = { + body: { + language: 'bash' as const, + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + headers: {}, + }; + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-1', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_programmatic', + workspaceId: 'primary', + request, + }); + + assert.deepEqual(programmaticRequests, [{ workspaceId: 'primary', request }]); + assert.deepEqual(quarantineEvents, [ + 'arm:Workspace programmatic execution is pending settlement', + 'clear', + ]); +}); + +test('worker keeps a selected workspace usable after an atomic programmatic setup failure', async () => { + const lifecycle: string[] = []; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + mutationFailuresAreAtomic: true, + async executeProgrammatic() { + throw new WorkspaceToolError( + 'Programmatic input download failed', + 'COMMAND_UNAVAILABLE', + ); + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + ], + ]), + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-setup-failure', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_programmatic', + workspaceId: 'primary', + request: { + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + headers: {}, + }, + }); + + assert.deepEqual(lifecycle, ['arm', 'clear']); + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.errorCode, 'COMMAND_UNAVAILABLE'); +}); + +test('worker keeps a selected workspace usable after confirmed programmatic cancellation cleanup', async () => { + const lifecycle: string[] = []; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + mutationFailuresAreAtomic: true, + async executeProgrammatic() { + throw new WorkspaceToolError( + 'Workspace command execution aborted', + 'EXECUTION_ABORTED', + true, + false, + ); + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + ], + ]), + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-cancelled-cleanly', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_programmatic', + workspaceId: 'primary', + request: { + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'sleep 30' }], + }, + headers: {}, + }, + }); + + assert.deepEqual(lifecycle, ['arm', 'clear']); + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.errorCode, 'EXECUTION_ABORTED'); +}); + +test('worker reports the underlying cause before quarantining an uncertain programmatic mutation', async () => { + const rootCause = new WorkspaceToolError( + 'Programmatic output upload failed', + 'COMMAND_UNAVAILABLE', + true, + true, + ); + let reported: unknown; + let quarantined = false; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + mutationFailuresAreAtomic: true, + async executeProgrammatic() { + throw rootCause; + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine(() => { + quarantined = true; + }), + ], + ]), + onError(error) { + reported = error; + }, + fetchImpl: async () => { + throw new Error('settlement must not run'); + }, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-uncertain-failure', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_programmatic', + workspaceId: 'primary', + request: { + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + headers: {}, + }, + }), + BridgeWorkspaceQuarantinedError, + ); + + assert.equal(reported, rootCause); + assert.equal(quarantined, true); +}); + test('worker stops after Code API rejects a fulfilled workspace mutation', async () => { let quarantinedReason: string | undefined; let armed = 0; diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index 5bfc07fe..dfda49fc 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -112,6 +112,8 @@ export interface SandboxWorkspaceToolsOptions { commandSandbox: WorkspaceCommandSandbox; /** Workspace IDs whose sandbox is configured and may run commands. */ commandWorkspaces: string[]; + /** Optional execution-scoped languages supplied by the same command sandbox. */ + programmaticLanguages?: BridgeWorkspaceToolCapabilities['programmaticLanguages']; } const MAX_SEARCH_CANDIDATE_BYTES = 1024 * 1024; @@ -1657,6 +1659,9 @@ export class SandboxWorkspaceTools implements WorkspaceToolExecutor { ...(base.listFileFeatures != null ? { listFileFeatures: base.listFileFeatures } : {}), + ...(options.programmaticLanguages?.length + ? { programmaticLanguages: [...options.programmaticLanguages] } + : {}), workspaces: base.workspaces.map((workspace) => ({ ...workspace, operations: [ diff --git a/service/src/bridge/concurrent-worker.test.ts b/service/src/bridge/concurrent-worker.test.ts index 5aec49cf..ea5ac601 100644 --- a/service/src/bridge/concurrent-worker.test.ts +++ b/service/src/bridge/concurrent-worker.test.ts @@ -243,11 +243,25 @@ for (const failure of [ status: 'fulfilled', value: { status: 'fulfilled' }, }); - for (let i = 0; i < 300 && errors.length === 0; i++) + const diagnosticCount = cleanupFailure ? 1 : 2; + const quarantineAttemptCount = + failure === 'lost-response' + ? 2 + : failure === 'all-responses-lost' || failure === 'delivery-outage' + ? 3 + : 1; + for ( + let i = 0; + i < 300 && + (errors.length < diagnosticCount || + (!cleanupFailure && quarantineAttempts < quarantineAttemptCount)); + i++ + ) { await new Promise((resolve) => setTimeout(resolve, 5)); + } if (failure === 'delivery-outage') - expect(errors.length).toBeGreaterThanOrEqual(1); - else expect(errors.length).toBe(1); + expect(errors.length).toBeGreaterThanOrEqual(2); + else expect(errors.length).toBe(diagnosticCount); if (failure === 'lost-response') expect(quarantineAttempts).toBe(2); if (failure === 'delivery-outage') expect(quarantineAttempts).toBeGreaterThanOrEqual(3); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 25b0bdaf..369b306c 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -414,6 +414,7 @@ router.post( supportedWorkspaceEditFileModes: ['single', 'batch'], supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], supportedWorkspaceListFileFeatures: ['after_path'], + supportedWorkspaceProgrammaticLanguages: ['bash'], }); } catch (error) { if (error instanceof BridgeStoreError) { diff --git a/service/src/bridge/selection.ts b/service/src/bridge/selection.ts index 0959279f..f3926e8a 100644 --- a/service/src/bridge/selection.ts +++ b/service/src/bridge/selection.ts @@ -1,4 +1,5 @@ export const CODEAPI_BRIDGE_WORKER_HEADER = 'X-LibreChat-Code-Worker-ID'; +export const CODEAPI_BRIDGE_WORKSPACE_HEADER = 'X-LibreChat-Code-Workspace-ID'; export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/; export class BridgeWorkerSelectionError extends Error { diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index b09d3792..d6b91469 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -151,6 +151,26 @@ function supportsWorkspaceTool( return true; } +function supportsWorkspaceProgrammatic( + registration: RegisteredBridgeWorker, + workspaceId: string, + language: string, +): boolean { + const capabilities = registration.capabilities.workspaceTools; + const workspace = capabilities?.workspaces.find( + (candidate) => candidate.id === workspaceId, + ); + return ( + workspace != null && + capabilities?.operations.includes('execute_command') === true && + (workspace.operations == null || + workspace.operations.includes('execute_command')) && + capabilities.programmaticLanguages?.includes( + language as 'bash', + ) === true + ); +} + function workerKey(workerId: string): string { return `${PREFIX}:worker:${encodeURIComponent(workerId)}`; } @@ -700,6 +720,7 @@ export class RedisBridgeStore { body: t.PayloadBody; headers: Record; workspaceRequest?: WorkspaceToolRequest; + workspaceId?: string; runtimeSessionId?: string; deadlineAtMs: number; executionTimeoutMs?: number; @@ -709,6 +730,12 @@ export class RedisBridgeStore { registration: RegisteredBridgeWorker, ) => Promise; }): Promise { + if (args.workspaceRequest != null && args.workspaceId != null) { + throw new BridgeStoreError( + 'ASSIGNMENT_INVALID', + 'A bridge assignment cannot be both a workspace tool and programmatic execution', + ); + } if ( args.executionTimeoutMs !== undefined && (args.workspaceRequest == null || @@ -773,6 +800,19 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} does not advertise the requested workspace tool`, ); } + if ( + args.workspaceId != null && + !supportsWorkspaceProgrammatic( + registration, + args.workspaceId, + args.body.language, + ) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} does not advertise programmatic execution for the selected workspace`, + ); + } if ( args.runtimeSessionId !== undefined && (await this.dispatchCommand( @@ -799,14 +839,16 @@ export class RedisBridgeStore { const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; let workspaceLeaseSlot: number | undefined; + const selectedWorkspaceId = + args.workspaceRequest?.workspaceId ?? args.workspaceId; const workspaceSlots = - args.workspaceRequest != null && + selectedWorkspaceId != null && (registration.capabilities.workspaceLeaseSlots ?? 1) > 1 ? new BridgeWorkspaceSlots(this.redis) : undefined; let resultCommitted = false; const admission = - args.workspaceRequest == null + selectedWorkspaceId == null ? undefined : new BridgeAdmissionQueue(this.redis); try { @@ -820,7 +862,7 @@ export class RedisBridgeStore { args.deadlineAtMs, workspaceSlots == null ? undefined - : args.workspaceRequest?.workspaceId, + : selectedWorkspaceId, ), args, 'Bridge admission enqueue', @@ -855,7 +897,7 @@ export class RedisBridgeStore { workerId: args.workerId, incarnationId: lockIncarnationId, assignmentId, - workspaceId: args.workspaceRequest!.workspaceId, + workspaceId: selectedWorkspaceId!, capacity: registration.capabilities.workspaceLeaseSlots!, expiresAtMs: Date.now() + ttlSeconds * 1000, }), @@ -910,7 +952,14 @@ export class RedisBridgeStore { ); } if ( - !supportsWorkspaceTool(current.registration, args.workspaceRequest!) + (args.workspaceRequest != null && + !supportsWorkspaceTool(current.registration, args.workspaceRequest)) || + (args.workspaceId != null && + !supportsWorkspaceProgrammatic( + current.registration, + args.workspaceId, + args.body.language, + )) ) { throw new BridgeStoreError( 'WORKER_MISMATCH', @@ -935,11 +984,14 @@ export class RedisBridgeStore { generation, leaseToken, leaseTokenHash: tokenHash(leaseToken), + ...(selectedWorkspaceId == null ? {} : { + workspaceFence: `native-workspace:${selectedWorkspaceId}`, + }), ...(workspaceLeaseSlot === undefined ? {} : { workspaceLeaseSlot, - workspaceFence: `native-workspace:${args.workspaceRequest!.workspaceId}`, + workspaceFence: `native-workspace:${selectedWorkspaceId!}`, }), ...(registration.identityId != null ? { workerIdentityId: registration.identityId } @@ -951,6 +1003,15 @@ export class RedisBridgeStore { executionKind: 'workspace_tool' as const, request: args.workspaceRequest, } + : args.workspaceId != null + ? { + executionKind: 'workspace_programmatic' as const, + workspaceId: args.workspaceId, + request: { + body: args.body, + headers: args.headers, + }, + } : { request: { body: args.body, @@ -1011,6 +1072,19 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} no longer advertises the requested workspace tool`, ); } + if ( + args.workspaceId != null && + !supportsWorkspaceProgrammatic( + replacement.registration, + args.workspaceId, + args.body.language, + ) + ) { + throw new BridgeStoreError( + 'WORKER_MISMATCH', + `Bridge worker ${args.workerId} no longer advertises programmatic execution for the selected workspace`, + ); + } registration = replacement.registration; readyToken = replacement.readyToken; } @@ -1034,11 +1108,24 @@ export class RedisBridgeStore { resultCommitted = true; return result; } catch (error) { - if (args.runtimeSessionId !== undefined) { + if (assignment.workspaceFence != null) { + // Native roots retain their own fence through result restoration. + // Do not quarantine unrelated roots or invalidate the worker lease. + await boundedCommand(this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", + "redis.call('SET', KEYS[1], 'quarantined:' .. ARGV[1])", + 'return 1', + ].join('\n'), + 1, + workspaceQuarantineKey(args.workerId, assignment.workspaceFence), + assignment.assignmentId, + ), this.redisCommandTimeoutMs, 'Bridge native workspace finalization quarantine'); + } else if (assignmentWorkspace(assignment) !== undefined) { await this.quarantine( args.workerId, assignment.incarnationId, - args.runtimeSessionId, + assignmentWorkspace(assignment)!, ); } throw error; @@ -1904,10 +1991,11 @@ export class RedisBridgeStore { : undefined; const cancelledMutation = signal.aborted && - workspaceRequest != null && - (workspaceRequest.operation === 'write_file' || - workspaceRequest.operation === 'edit_file' || - workspaceRequest.operation === 'execute_command'); + (assignment.executionKind === 'workspace_programmatic' || + (workspaceRequest != null && + (workspaceRequest.operation === 'write_file' || + workspaceRequest.operation === 'edit_file' || + workspaceRequest.operation === 'execute_command'))); if (cancelledMutation) { try { // Keep the acknowledged assignment available long enough for the diff --git a/service/src/bridge/workspace-store.test.ts b/service/src/bridge/workspace-store.test.ts index 5b0186d0..f248aa40 100644 --- a/service/src/bridge/workspace-store.test.ts +++ b/service/src/bridge/workspace-store.test.ts @@ -77,6 +77,123 @@ test('dispatches a workspace tool only to a worker advertising its workspace and }); }); +for (const finalizationFails of [false, true]) test(`single-slot programmatic finalization retains the workspace fence (failure=${finalizationFails})`, async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['execute_command'], + programmaticLanguages: ['bash'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const body = { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }; + const completion = store.dispatch({ + workerId: 'workspace-worker', + body, + headers: {}, + workspaceId: 'primary', + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + finalize: async settlement => { + if (finalizationFails) throw new Error('artifact restoration failed'); + return settlement; + }, + }); + + const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + expect(assignment).toMatchObject({ + executionKind: 'workspace_programmatic', + workspaceId: 'primary', + request: { body }, + }); + await store.settle('workspace-worker', assignment?.assignmentId ?? '', { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'fulfilled', + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2', + files: [], + run: { + stdout: 'ready\n', + stderr: '', + code: 0, + signal: null, + output: 'ready\n', + memory: null, + message: null, + status: null, + cpu_time: null, + wall_time: 0.01, + }, + }, + }); + + if (finalizationFails) { + await expect(completion).rejects.toThrow('artifact restoration failed'); + await expect(store.dispatch({ workerId: 'workspace-worker', body, headers: {}, + workspaceId: 'primary', deadlineAtMs: Date.now() + 1000, + signal: new AbortController().signal, + })).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED' }); + return; + } + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + result: { session_id: 'session-1' }, + }); +}); + +test('rejects programmatic execution without the workspace capability', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['execute_command'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatch({ + workerId: 'workspace-worker', + body: { + language: 'bash', + version: '5.2', + session_id: 'session-2', + files: [{ name: 'main.sh', content: 'echo denied' }], + }, + headers: {}, + workspaceId: 'primary', + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); +}); + test('drains an acknowledged workspace mutation cancellation before releasing it', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/config.spec.ts b/service/src/config.spec.ts index 44b64b8b..69b416d8 100644 --- a/service/src/config.spec.ts +++ b/service/src/config.spec.ts @@ -99,6 +99,11 @@ describe('egress grant TTL configuration', () => { }); describe('job deadline accounting', () => { + it('never extends the producer deadline when worker configuration differs', () => { + expect(jobDeadlineAtMs(1_000, 300_000, 50_000, 91_000)).toBe(91_000); + expect(jobDeadlineAtMs(1_000, 30_000, 50_000, 91_000)).toBe(31_000); + expect(jobDeadlineAtMs(1_000, 300_000, 50_000, Number.NaN)).toBe(0); + }); it('counts time spent waiting in BullMQ against JOB_TIMEOUT', () => { expect(jobDeadlineAtMs(1_000, 300_000, 50_000)).toBe(301_000); }); diff --git a/service/src/config.ts b/service/src/config.ts index d025831e..55570dfe 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -95,10 +95,17 @@ export function jobDeadlineAtMs( enqueuedAtMs: number | undefined, timeoutMs: number, nowMs: number = Date.now(), + producerDeadlineAtMs?: number, ): number { - return Number.isFinite(enqueuedAtMs) && (enqueuedAtMs as number) > 0 + const localDeadline = Number.isFinite(enqueuedAtMs) && (enqueuedAtMs as number) > 0 ? (enqueuedAtMs as number) + timeoutMs : nowMs + timeoutMs; + if (producerDeadlineAtMs === undefined) return localDeadline; + // A worker with a larger JOB_TIMEOUT must not outlive the admission fence + // retained by its API producer. Malformed explicit deadlines fail closed. + return Number.isFinite(producerDeadlineAtMs) + ? Math.min(localDeadline, producerDeadlineAtMs) + : 0; } /** The worker stops user work at JOB_TIMEOUT, then may still need to terminate diff --git a/service/src/egress-gateway.ts b/service/src/egress-gateway.ts index f86ce656..873db173 100644 --- a/service/src/egress-gateway.ts +++ b/service/src/egress-gateway.ts @@ -45,6 +45,7 @@ import { parseBoundedContentLength } from './http-limits'; import { validateEgressGatewayHardenedConfig } from './secure-startup'; import { isOpaqueObjectContentDisposition } from './file-metadata'; import { mapObjectDetails } from './file-object-resolver'; +import { isSupportedBridgeArtifactName } from '../../packages/code/src/protocol'; export const app: Express = express(); app.disable('x-powered-by'); @@ -52,29 +53,6 @@ validateEgressGatewayHardenedConfig(); app.use(traceHttpRequest('codeapi.egress_gateway.request')); app.use(httpMetricsMiddleware); -const SUPPORTED_OUTPUT_EXTENSIONS = new Set([ - '.c', '.cs', '.cpp', '.go', '.java', '.js', '.kt', '.kts', '.lua', - '.php', '.pl', '.ps1', '.py', '.r', '.rb', '.rs', '.scala', '.sh', - '.sql', '.swift', '.ts', '.jsx', '.tsx', '.groovy', - '.css', '.htm', '.html', '.less', '.sass', '.scss', '.svg', '.svelte', '.vue', - '.adoc', '.asciidoc', '.md', '.rst', '.tex', '.txt', '.wiki', - '.csv', '.json', '.bson', '.json5', '.jsonl', '.parquet', '.tsv', - '.xml', '.yaml', '.yml', - '.ics', '.ical', '.ifb', '.icalendar', - '.conf', '.env', '.gitignore', '.ini', '.properties', '.toml', - '.doc', '.docx', '.pdf', '.ppt', '.pptx', '.xls', '.xlsx', - '.odt', '.ods', '.odp', '.rtf', - '.avif', '.bmp', '.gif', '.ico', '.jpeg', '.jpg', '.png', - '.tif', '.tiff', '.webp', - '.eot', '.ttf', '.woff', '.woff2', - '.7z', '.bz2', '.gz', '.gzip', '.rar', '.tar', '.zip', - '.tf', '.tfvars', '.tfstate', '.hcl', - '.dockerfile', '.Dockerfile', '.dockerignore', - '.helmignore', '.helmfile', '.jenkinsfile', '.vagrantfile', - '.eslintrc', '.prettierrc', '.editorconfig', '.nomad', - '.bat', '.cmd', '.deb', '.log', '.rpm', '.vbs', -]); - type EgressAuditFields = { execHash?: string; requestExecHash?: string; @@ -248,18 +226,7 @@ function assertOutputFilenameAllowed(name: string): void { throw new EgressGrantError('malformed', 'Output filename must be canonical'); } if (!isDirkeepName(name)) { - const basename = path.posix.basename(name); - const ext = path.posix.extname(basename).toLowerCase(); - const dottedBasename = `.${basename}`; - const allowed = - (ext !== '' && SUPPORTED_OUTPUT_EXTENSIONS.has(ext)) || - SUPPORTED_OUTPUT_EXTENSIONS.has(basename) || - SUPPORTED_OUTPUT_EXTENSIONS.has(basename.toLowerCase()) || - (ext === '' && ( - SUPPORTED_OUTPUT_EXTENSIONS.has(dottedBasename) || - SUPPORTED_OUTPUT_EXTENSIONS.has(dottedBasename.toLowerCase()) - )); - if (!allowed) { + if (!isSupportedBridgeArtifactName(name)) { throw new EgressGrantError('scope_mismatch', 'Output filename extension is not supported'); } } diff --git a/service/src/egress-grant.test.ts b/service/src/egress-grant.test.ts index 32ef2988..cd26fcd7 100644 --- a/service/src/egress-grant.test.ts +++ b/service/src/egress-grant.test.ts @@ -5,6 +5,7 @@ import { env } from './config'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, + normalizeSelectedWorkspaceProgrammaticTimeoutMs, prepareSandboxJobSecurity, refreshEgressGrantClaims, timeoutMsToGrantSeconds, @@ -478,6 +479,13 @@ describe('egress encrypted grants and handles', () => { expect(() => normalizeProgrammaticTimeoutMs(0, 300000)).toThrow('timeout must be a positive number'); }); + test('budgets both selected-workspace replay passes inside the worker deadline', () => { + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(undefined, 300_000)).toBe(82_500); + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(120_000, 300_000)).toBe(82_500); + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(300_000, 300_000)).toBe(82_500); + expect(normalizeSelectedWorkspaceProgrammaticTimeoutMs(10_000, 20_000)).toBe(2_167); + }); + test('normalizes the gateway callback URL for sandbox-originated PTC', () => { expect(normalizeEgressGatewayUrl(' http://egress-gateway:3190/// ')).toBe('http://egress-gateway:3190'); expect(() => normalizeEgressGatewayUrl(' ')).toThrow('EGRESS_GATEWAY_URL is required'); diff --git a/service/src/execution-log.test.ts b/service/src/execution-log.test.ts index f04bad90..0f538125 100644 --- a/service/src/execution-log.test.ts +++ b/service/src/execution-log.test.ts @@ -28,6 +28,12 @@ describe('execution log summaries', () => { failed: 1, detail: 'private storage failure', }, + artifact_truncation: { + code: 'artifact_truncated', + reasons: { max_files: 2 }, + skipped: ['secret-one.txt', 'secret-two.txt'], + skipped_count: 2, + }, run: { code: 0, stdout: 'top secret stdout', @@ -42,6 +48,7 @@ describe('execution log summaries', () => { expect(JSON.stringify(summary)).not.toContain('sensitive stderr'); expect(JSON.stringify(summary)).not.toContain('combined output'); expect(JSON.stringify(summary)).not.toContain('private storage failure'); + expect(JSON.stringify(summary)).not.toContain('secret-one.txt'); expect(summary).toMatchObject({ session_id: 'sess_123', files: { count: 2, inheritedCount: 1, modifiedCount: 1 }, @@ -52,6 +59,12 @@ describe('execution log summaries', () => { delivered: 2, failed: 1, }, + artifact_truncation: { + code: 'artifact_truncated', + reasons: { max_files: 2 }, + skipped_count: 2, + reported_paths: 2, + }, run: { stdout: { length: 17, present: true }, stderr: { length: 16, present: true }, diff --git a/service/src/execution-log.ts b/service/src/execution-log.ts index 48a93cce..143e5541 100644 --- a/service/src/execution-log.ts +++ b/service/src/execution-log.ts @@ -19,6 +19,7 @@ type SandboxResponseLike = { version?: unknown; files?: unknown; artifact_delivery?: unknown; + artifact_truncation?: unknown; run?: RunLike; }; @@ -40,6 +41,22 @@ function summarizeArtifactDelivery(value: unknown): Record | un }; } +function summarizeArtifactTruncation(value: unknown): Record | undefined { + if (value == null || typeof value !== 'object' || Array.isArray(value)) return undefined; + const truncation = value as { + code?: unknown; + reasons?: unknown; + skipped?: unknown; + skipped_count?: unknown; + }; + return { + code: truncation.code, + reasons: truncation.reasons, + skipped_count: truncation.skipped_count, + reported_paths: Array.isArray(truncation.skipped) ? truncation.skipped.length : undefined, + }; +} + export function summarizeText(value: unknown): { length: number; present: boolean } { if (typeof value !== 'string') { return { length: 0, present: false }; @@ -87,6 +104,7 @@ export function summarizeSandboxResponse(data: SandboxResponseLike): Record>; +beforeEach(async () => { + redis = await startTestRedis(); +}); +afterEach(async () => { + await redis.closeTestServer(); +}); +const target = { queueName: 'other', jobId: 'commit-race' }; + +for (const outcome of ['commit', 'stop', 'duplicate']) + test(`native mutation handoff commits or quarantines before root release (${outcome})`, async () => { + const store = new RedisBridgeStore(redis); + const workerId = 'handoff-worker'; + const incarnationId = 'incarnation-handoff-01'; + await store.register({ + protocolVersion: 1, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: 1, + operations: ['execute_command'], + programmaticLanguages: ['bash'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const controller = new AbortController(); + const dispatchArgs = { + workerId, + workspaceId: 'primary', + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'handoff-session', + files: [{ name: 'main.sh', content: 'echo mutation' }], + }, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }; + const completion = store.dispatch({ + ...dispatchArgs, + finalize: async settlement => { + if (outcome === 'stop') await requestJobCancellation(redis, target, 60); + if (outcome === 'duplicate') + await commitJobResult( + redis, + target, + { stdout: 'first mutation' }, + 60, + ); + if ( + (await commitJobResult( + redis, + target, + { stdout: 'mutation settled' }, + 60, + )) !== 'committed' + ) + throw new Error('handoff did not win'); + // This represents Stop during post-handoff egress cleanup. It must no + // longer turn the applied mutation into an acknowledged cancellation. + expect(await requestJobCancellation(redis, target, 60)).toBe(false); + return settlement; + }, + }); + void completion.catch(() => undefined); + const assignment = await store.lease(workerId, incarnationId, 1_000); + if (assignment == null) throw new Error('Missing assignment'); + await store.settle(workerId, assignment.assignmentId, { + protocolVersion: 1, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'fulfilled', + result: { + session_id: 'handoff-session', + language: 'bash', + version: '5.2', + files: [], + }, + }); + if (outcome !== 'commit') { + await expect(completion).rejects.toThrow('handoff did not win'); + await expect(store.dispatch(dispatchArgs)).rejects.toMatchObject({ + code: 'WORKSPACE_QUARANTINED', + }); + } else { + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + }); + expect(await readCommittedJobResult(redis, target)).toEqual({ + result: { stdout: 'mutation settled' }, + }); + } + }); + +test('concurrent stalled-job redelivery claims at most one sandbox execution', async () => { + let executions = 0; + const attempt = async () => { + const claim = await claimJobExecution(redis, target, 60); + if (claim.status === 'claimed') executions += 1; + return claim; + }; + const results = await Promise.allSettled([attempt(), attempt()]); + expect(executions).toBe(1); + expect(results.filter(result => result.status === 'fulfilled')).toHaveLength( + 1, + ); + expect(results.filter(result => result.status === 'rejected')).toHaveLength( + 1, + ); + await expect(attempt()).rejects.toThrow('already claimed'); + expect(executions).toBe(1); + await commitJobResult(redis, target, { stdout: 'first result' }, 60); + expect(await attempt()).toEqual({ + status: 'completed', + result: { stdout: 'first result' }, + }); + expect(executions).toBe(1); + expect( + await commitJobResult(redis, target, { stdout: 'different result' }, 60), + ).toBe('already_completed'); + expect(await readCommittedJobResult(redis, target)).toEqual({ + result: { stdout: 'first result' }, + }); +}); + +test('a lost execution-claim reply never authorizes a second attempt', async () => { + const lostReply = { + eval: async (...args: Parameters) => { + await redis.eval(...args); + throw new Error('claim reply lost'); + }, + } as unknown as typeof redis; + await expect(claimJobExecution(lostReply, target, 60)).rejects.toThrow( + 'claim reply lost', + ); + await expect(claimJobExecution(redis, target, 60)).rejects.toThrow( + 'already claimed', + ); +}); + +test('cancel-before-claim and missing completion payload fail closed', async () => { + await requestJobCancellation(redis, target, 60); + await expect(claimJobExecution(redis, target, 60)).rejects.toThrow( + 'cancelled', + ); + const completedTarget = { ...target, jobId: 'missing-payload-claim' }; + await commitJobResult(redis, completedTarget, { stdout: 'done' }, 60); + await redis.del( + `${jobCancellationInternals.cancellationKey(completedTarget)}:result`, + ); + await expect(claimJobExecution(redis, completedTarget, 60)).rejects.toThrow( + 'refusing re-execution', + ); +}); + +for (const corrupt of [false, true]) + test(`invalid committed result fails immediately without Redis retries (corrupt=${corrupt})`, async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 60); + const key = jobCancellationInternals.cancellationKey(target); + if (corrupt) await redis.set(`${key}:result`, '{invalid'); + else await redis.del(`${key}:result`); + let calls = 0; + const commands = { + eval: (...args: Parameters) => { + calls += 1; + return redis.eval(...args); + }, + } as unknown as typeof redis; + await expect( + fenceJobCancellation({ + commands, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 30_000, + }), + ).rejects.toThrow(); + expect(calls).toBe(1); + }); + +test('completion retention includes the API producer across timeout configuration drift', async () => { + const ttl = jobCancellationRetentionSeconds(30_000, 430); + expect(ttl).toBe(430); + expect(jobCancellationRetentionSeconds(300_000, 430)).toBe(780); + await commitJobResult(redis, target, { stdout: 'done' }, ttl); + const key = jobCancellationInternals.cancellationKey(target); + expect(await redis.ttl(key)).toBeGreaterThanOrEqual(429); + expect(await redis.ttl(`${key}:result`)).toBeGreaterThanOrEqual(429); +}); + +test('a late Stop renews completion evidence along with its request tombstone', async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 1); + expect(await requestJobCancellation(redis, target, 60)).toBe(false); + const key = jobCancellationInternals.cancellationKey(target); + expect(await redis.ttl(key)).toBeGreaterThanOrEqual(59); + expect(await redis.ttl(`${key}:result`)).toBeGreaterThanOrEqual(59); +}); + +test('retention renewal does not lose subsecond time to rounded TTL readings', async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 60); + const key = jobCancellationInternals.cancellationKey(target); + await redis.pexpire(key, 59_900); + const expiration = async () => + Number( + await redis.eval( + ` + local now = redis.call('TIME') + return tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) + redis.call('PTTL', KEYS[1]) + `, + 1, + key, + ), + ); + const before = await expiration(); + await requestJobCancellation(redis, target, 60); + expect(await expiration()).toBeGreaterThan(before); +}); + +test('fencing returns the committed result without a vulnerable second Redis read', async () => { + const result = { stdout: 'one committed effect' }; + await commitJobResult(redis, target, result, 60); + let calls = 0; + const connectionDropsAfterDecision = { + eval: async (...args: Parameters) => { + calls += 1; + return redis.eval(...args); + }, + get: async () => { + throw new Error('connection lost after decision'); + }, + mget: async () => { + throw new Error('connection lost after decision'); + }, + } as unknown as typeof redis; + expect( + await fenceJobCancellation({ + commands: connectionDropsAfterDecision, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 1_000, + }), + ).toEqual({ status: 'completed', result }); + expect(calls).toBe(1); +}); + +test('disconnect returns a known completed result without waiting for a lost queue event', async () => { + const result = { stdout: 'done' }; + await commitJobResult(redis, target, result, 60); + const registry = new JobCancellationRegistry(redis); + const controller = new AbortController(); + controller.abort(); + const job = { + id: target.jobId, + queueName: target.queueName, + waitUntilFinished: () => new Promise(() => {}), + } as unknown as Parameters[0]['job']; + try { + expect( + await waitForJobWithCancellation({ + commands: redis, + registry, + job, + events: {} as Parameters< + typeof waitForJobWithCancellation + >[0]['events'], + timeoutMs: 1_000, + cancellationTtlSeconds: 60, + signal: controller.signal, + }), + ).toEqual(result); + } finally { + await registry.close(); + } +}); + +test('durable cancellation wins even before its subscriber notification arrives', async () => { + expect(await requestJobCancellation(redis, target, 60)).toBe(true); + expect(await commitJobResult(redis, target, { stdout: 'late' }, 60)).toBe( + 'cancelled', + ); + expect(await readCommittedJobResult(redis, target)).toBeUndefined(); +}); + +test('committed results reject late Stop and survive a lost BullMQ completion reply', async () => { + const result = { stdout: 'one mutation', files: [] }; + expect(await commitJobResult(redis, target, result, 60)).toBe('committed'); + expect(await requestJobCancellation(redis, target, 60)).toBe(false); + expect(await readCommittedJobResult(redis, target)).toEqual({ result }); + expect( + await redis.get(jobCancellationInternals.cancellationKey(target)), + ).toBe('completed'); + const registry = new JobCancellationRegistry(redis); + const controller = new AbortController(); + try { + await registry.register(target, controller); + expect(controller.signal.aborted).toBe(false); + } finally { + await registry.close(); + } +}); + +test('concurrent cancellation and completion have exactly one winner', async () => { + const [cancelled, committed] = await Promise.all([ + requestJobCancellation(redis, target, 60), + commitJobResult(redis, target, { stdout: 'result' }, 60), + ]); + expect(Number(cancelled) + Number(committed === 'committed')).toBe(1); +}); + +test('a missing committed result fails closed instead of re-executing', async () => { + await commitJobResult(redis, target, { stdout: 'already applied' }, 60); + await redis.del(`${jobCancellationInternals.cancellationKey(target)}:result`); + await expect(readCommittedJobResult(redis, target)).rejects.toThrow( + 'refusing re-execution', + ); +}); + +test('an enqueue failure can recover a result that won cancellation fencing', async () => { + const result = { stdout: 'effect already applied' }; + await commitJobResult(redis, target, result, 60); + expect( + await fenceJobCancellation({ + commands: redis, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 5_000, + }), + ).toEqual({ status: 'completed', result }); + expect(await readCommittedJobResult(redis, target)).toEqual({ result }); +}); + +test('enqueue fencing still recovers completion after the original deadline', async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 60); + expect( + await fenceJobCancellation({ + commands: redis, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() - 1_000, + }), + ).toEqual({ status: 'completed', result: { stdout: 'done' } }); +}); + +test('Redis rejects commitment when recovery happens after the producer deadline', async () => { + const delayed = { + eval: async (...args: Parameters) => { + await new Promise(resolve => setTimeout(resolve, 150)); + return redis.eval(...args); + }, + } as unknown as typeof redis; + await expect( + commitJobResult(delayed, target, { stdout: 'late' }, 60, Date.now() + 100), + ).rejects.toThrow('exceeded its deadline'); + expect(await readCommittedJobResult(redis, target)).toBeUndefined(); +}); + +test('a timely durable commit remains successful when only its acknowledgement is late', async () => { + const delayedReply = { + eval: async (...args: Parameters) => { + const value = await redis.eval(...args); + await new Promise(resolve => setTimeout(resolve, 150)); + return value; + }, + } as unknown as typeof redis; + expect( + await commitJobResult( + delayedReply, + target, + { stdout: 'committed' }, + 60, + Date.now() + 100, + ), + ).toBe('committed'); + expect(await readCommittedJobResult(redis, target)).toEqual({ + result: { stdout: 'committed' }, + }); +}); + +for (const failedStage of ['subscription', 'completion'] as const) { + test(`a lost ${failedStage} reply recovers the committed result instead of reporting failure`, async () => { + const result = { stdout: 'already applied once' }; + await commitJobResult(redis, target, result, 60); + const registry = new JobCancellationRegistry(redis); + if (failedStage === 'subscription') + registry.register = async () => { + throw new Error('lost reply'); + }; + const job = { + id: target.jobId, + queueName: target.queueName, + waitUntilFinished: () => Promise.reject(new Error('lost result event')), + } as unknown as Parameters[0]['job']; + try { + expect( + await waitForJobWithCancellation({ + commands: redis, + registry, + job, + events: {} as Parameters< + typeof waitForJobWithCancellation + >[0]['events'], + timeoutMs: 1_000, + cancellationTtlSeconds: 60, + }), + ).toEqual(result); + } finally { + await registry.close(); + } + }); +} diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts new file mode 100644 index 00000000..e23186e4 --- /dev/null +++ b/service/src/job-cancellation.test.ts @@ -0,0 +1,599 @@ +import { expect, test } from 'bun:test'; +import { EventEmitter } from 'node:events'; +import type IORedis from 'ioredis'; +import type { Job, QueueEvents } from 'bullmq'; +import { + CLIENT_DISCONNECT_REASON, + JobCancellationRegistry, + jobResultCommitFailure, + jobCancellationInternals, + removeJobIfWaiting, + requestJobCancellation, + throwIfJobAborted, + waitForJobWithCancellation, + fenceJobCancellation, +} from './job-cancellation'; + +class FakeSubscriber extends EventEmitter { + subscribed?: string; + closed = false; + subscribeFailures = 0; + + async subscribe(channel: string): Promise { + if (this.subscribeFailures > 0) { + this.subscribeFailures -= 1; + throw new Error('subscriber unavailable'); + } + this.subscribed = channel; + return 1; + } + + async quit(): Promise<'OK'> { + this.closed = true; + return 'OK'; + } + + disconnect(): void { + this.closed = true; + } +} + +class FakeTransaction { + readonly operations: unknown[][] = []; + + set(...args: unknown[]): this { + this.operations.push(['set', ...args]); + return this; + } + + publish(...args: unknown[]): this { + this.operations.push(['publish', ...args]); + return this; + } + + async exec(): Promise> { + return this.operations.map(() => [null, 'OK']); + } +} + +class FakeRedis { + readonly subscriber = new FakeSubscriber(); + duplicateCalls = 0; + readonly existing = new Set(); + readonly deleted: string[] = []; + readonly transactions: FakeTransaction[] = []; + mgetFailures = 0; + cancellationFailures = 0; + cancellationAttempts = 0; + + duplicate(): FakeSubscriber { + this.duplicateCalls += 1; + return this.subscriber; + } + + async get(key: string): Promise { + return this.existing.has(key) ? '1' : null; + } + + async eval( + _script: string, + _keys: number, + key: string, + _resultKey: string, + ttl: number, + channel: string, + payload: string, + ): Promise { + this.cancellationAttempts += 1; + if (this.cancellationFailures-- > 0) throw new Error('Redis unavailable'); + const transaction = this.multi(); + transaction.set(key, '1', 'EX', ttl); + transaction.publish(channel, payload); + await transaction.exec(); + return [1]; + } + + async mget(...keys: string[]): Promise> { + if (this.mgetFailures > 0) { + this.mgetFailures -= 1; + throw new Error('command connection unavailable'); + } + return keys.map(key => (this.existing.has(key) ? '1' : null)); + } + + async del(key: string): Promise { + this.deleted.push(key); + this.existing.delete(key); + return 1; + } + + multi(): FakeTransaction { + const transaction = new FakeTransaction(); + this.transactions.push(transaction); + return transaction; + } +} + +function redis(fake: FakeRedis): IORedis { + return fake as unknown as IORedis; +} + +test('idle registries allocate no subscriber connection', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + + await registry.close(); + + expect(fake.duplicateCalls).toBe(0); +}); + +test('shutdown disconnects a subscriber whose startup is still waiting for Redis', async () => { + const fake = new FakeRedis(); + fake.subscriber.subscribe = async () => new Promise(() => {}); + const registry = new JobCancellationRegistry(redis(fake)); + void registry + .register( + { queueName: 'other', jobId: 'shutdown-startup' }, + new AbortController(), + ) + .catch(() => undefined); + await registry.close(); + expect(fake.subscriber.closed).toBe(true); + expect(fake.subscriber.listenerCount('message')).toBe(0); +}, 1_000); + +test('failed subscription startup removes handlers before a bounded retry', async () => { + const fake = new FakeRedis(); + fake.subscriber.subscribeFailures = 1; + const registry = new JobCancellationRegistry(redis(fake)); + const first = new AbortController(); + + await expect( + registry.register({ queueName: 'other', jobId: 'job-failed-start' }, first), + ).rejects.toThrow('subscriber unavailable'); + expect(fake.subscriber.listenerCount('message')).toBe(0); + expect(fake.subscriber.listenerCount('ready')).toBe(0); + expect(fake.subscriber.listenerCount('error')).toBe(0); + + const second = new AbortController(); + await registry.register({ queueName: 'other', jobId: 'job-retry' }, second); + expect(fake.duplicateCalls).toBe(2); + expect(fake.subscriber.listenerCount('message')).toBe(1); + await registry.close(); +}); + +test('registry catches durable cancellation before subscriber registration', async () => { + const fake = new FakeRedis(); + const target = { queueName: 'other', jobId: 'job-1' }; + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + const registry = new JobCancellationRegistry(redis(fake)); + const controller = new AbortController(); + + await registry.register(target, controller); + + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + expect(fake.subscriber.subscribed).toBe(jobCancellationInternals.channel); + await registry.unregister(target); + await registry.close(); + expect(fake.subscriber.closed).toBe(true); +}); + +test('one pubsub listener cancels only the matching active job', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const first = new AbortController(); + const second = new AbortController(); + await registry.register({ queueName: 'other', jobId: 'job-1' }, first); + await registry.register({ queueName: 'other', jobId: 'job-2' }, second); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify({ queueName: 'other', jobId: 'job-2' }), + ); + + expect(first.signal.aborted).toBe(false); + expect(second.signal.aborted).toBe(true); + await registry.close(); +}); + +test('one pubsub listener wakes every local waiter for the same job', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-shared' }; + const first = new AbortController(); + const second = new AbortController(); + await registry.register(target, first); + await registry.register(target, second); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify(target), + ); + + expect(first.signal.aborted).toBe(true); + expect(second.signal.aborted).toBe(true); + expect(fake.duplicateCalls).toBe(1); + await registry.close(); +}); + +test('unregistering one local waiter preserves other waiters for the job', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-shared-unregister' }; + const first = new AbortController(); + const second = new AbortController(); + await registry.register(target, first); + await registry.register(target, second); + await registry.unregister(target, first); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify(target), + ); + + expect(first.signal.aborted).toBe(false); + expect(second.signal.aborted).toBe(true); + await registry.close(); +}); + +test('subscriber reconnect reconciles active jobs against durable markers', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-reconnect' }; + const controller = new AbortController(); + await registry.register(target, controller); + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + + fake.subscriber.emit('ready'); + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + await registry.close(); +}); + +test('subscriber reconnect retries durable-marker reconciliation', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-retry-reconcile' }; + const controller = new AbortController(); + await registry.register(target, controller); + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + fake.mgetFailures = 1; + + fake.subscriber.emit('ready'); + await new Promise(resolve => setTimeout(resolve, 150)); + + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + await registry.close(); +}); + +test('terminal subscriber disconnect rebuilds the subscription and reconciles markers', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-terminal-reconnect' }; + const controller = new AbortController(); + await registry.register(target, controller); + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + + fake.subscriber.emit('end'); + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(fake.duplicateCalls).toBe(2); + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + await registry.close(); +}); + +test('cancellation writes a durable marker before publishing', async () => { + const fake = new FakeRedis(); + const target = { queueName: 'other', jobId: 'job-3' }; + + await requestJobCancellation(redis(fake), target, 42); + + expect(fake.transactions).toHaveLength(1); + expect(fake.transactions[0]?.operations).toEqual([ + ['set', jobCancellationInternals.cancellationKey(target), '1', 'EX', 42], + ['publish', jobCancellationInternals.channel, JSON.stringify(target)], + ]); +}); + +test('result commit barrier rejects cancellation observed after execution', () => { + const controller = new AbortController(); + expect(() => throwIfJobAborted(controller.signal)).not.toThrow(); + controller.abort(CLIENT_DISCONNECT_REASON); + expect(() => throwIfJobAborted(controller.signal)).toThrow( + CLIENT_DISCONNECT_REASON, + ); +}); + +test('result cleanup maps late cancellation to stable worker failures', () => { + const disconnected = new AbortController(); + disconnected.abort(CLIENT_DISCONNECT_REASON); + expect(jobResultCommitFailure(disconnected.signal, 30_000)?.message).toBe( + 'Job cancelled after client disconnected', + ); + + const deadline = new AbortController(); + deadline.abort('deadline'); + expect(jobResultCommitFailure(deadline.signal, 30_000)?.message).toBe( + 'Job timed out after 30000ms', + ); + expect( + jobResultCommitFailure(new AbortController().signal, 30_000), + ).toBeUndefined(); +}); + +test('disconnect frees a waiting job and rejects promptly', async () => { + const fake = new FakeRedis(); + const controller = new AbortController(); + let removed = false; + const never = new Promise(() => {}); + const job = { + id: 'job-4', + queueName: 'other', + waitUntilFinished: () => never, + getState: async () => 'waiting', + remove: async () => { + removed = true; + }, + } as unknown as Job; + + const registry = new JobCancellationRegistry(redis(fake)); + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + signal: controller.signal, + }); + controller.abort(CLIENT_DISCONNECT_REASON); + + await expect(waiting).rejects.toMatchObject({ name: 'AbortError' }); + expect(removed).toBe(true); + expect(fake.cancellationAttempts).toBe(1); + expect(fake.transactions[0]?.operations[0]).toEqual([ + 'set', + jobCancellationInternals.cancellationKey({ + queueName: 'other', + jobId: 'job-4', + }), + '1', + 'EX', + 120, + ]); + await registry.close(); +}); + +test('registration failure fences and removes the already-enqueued job', async () => { + const fake = new FakeRedis(); + fake.subscriber.subscribeFailures = 1; + let removed = false; + const job = { + id: 'job-register-failure', + queueName: 'other', + waitUntilFinished: () => new Promise(() => {}), + getState: async () => 'waiting', + remove: async () => { + removed = true; + }, + } as unknown as Job; + const registry = new JobCancellationRegistry(redis(fake)); + + await expect( + waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + }), + ).rejects.toThrow('subscriber unavailable'); + + expect(removed).toBe(true); + expect(fake.transactions[0]?.operations[0]).toEqual([ + 'set', + jobCancellationInternals.cancellationKey({ + queueName: 'other', + jobId: 'job-register-failure', + }), + '1', + 'EX', + 120, + ]); + await registry.close(); +}); + +test('a result rejection is owned while subscription registration is pending', async () => { + const fake = new FakeRedis(); + let release!: () => void; + fake.subscriber.subscribe = async () => { + await new Promise(resolve => { + release = resolve; + }); + return 1; + }; + const registry = new JobCancellationRegistry(redis(fake)); + const job = { + id: 'pending-registration', + queueName: 'other', + waitUntilFinished: () => Promise.reject(new Error('completion timeout')), + getState: async () => 'active', + remove: async () => {}, + } as unknown as Job; + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 1_000, + cancellationTtlSeconds: 60, + }); + const rejection = waiting.catch((error: Error) => error); + // An unowned rejection fails the test runner on this event-loop turn. + await new Promise(resolve => setImmediate(resolve)); + release(); + expect(await rejection).toMatchObject({ message: 'completion timeout' }); + await registry.close(); +}); + +test('external cancellation frees a waiting job before rejecting its waiter', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + let removed = false; + const job = { + id: 'job-external-waiting', + queueName: 'other', + waitUntilFinished: () => new Promise(() => {}), + getState: async () => 'waiting', + remove: async () => { + removed = true; + }, + } as unknown as Job; + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + }); + await new Promise(resolve => setImmediate(resolve)); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify({ queueName: 'other', jobId: 'job-external-waiting' }), + ); + + await expect(waiting).rejects.toMatchObject({ name: 'AbortError' }); + expect(removed).toBe(true); + await registry.close(); +}); + +test('a separate cancellation request wakes the original job waiter', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const never = new Promise(() => {}); + const job = { + id: 'job-external-cancel', + queueName: 'other', + waitUntilFinished: () => never, + getState: async () => 'active', + remove: async () => undefined, + } as unknown as Job; + + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + }); + await new Promise(resolve => setImmediate(resolve)); + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify({ queueName: 'other', jobId: 'job-external-cancel' }), + ); + + await expect(waiting).rejects.toMatchObject({ name: 'AbortError' }); + expect(fake.deleted).toEqual([]); + await registry.close(); +}); + +test('queued removal never removes an active job', async () => { + let removed = false; + const job = { + getState: async () => 'active' as const, + remove: async () => { + removed = true; + }, + }; + + expect(await removeJobIfWaiting(job)).toBe(false); + expect(removed).toBe(false); +}); + +test('a failed cancellation write retains ownership until a durable retry succeeds', async () => { + const fake = new FakeRedis(); + fake.cancellationFailures = 2; + const target = { queueName: 'other', jobId: 'ambiguous-enqueue' }; + let released = false; + const fencing = fenceJobCancellation({ + commands: redis(fake), + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 500, + }).then(() => { + released = true; + }); + await new Promise(resolve => setTimeout(resolve, 10)); + expect(released).toBe(false); + await fencing; + expect(fake.cancellationAttempts).toBe(3); + expect(released).toBe(true); +}); + +test('an unavailable Redis cannot release ownership before the fixed job deadline', async () => { + const fake = new FakeRedis(); + fake.cancellationFailures = 1_000; + const startedAt = Date.now(); + await fenceJobCancellation({ + commands: redis(fake), + target: { queueName: 'other', jobId: 'offline' }, + ttlSeconds: 60, + deadlineAtMs: startedAt + 80, + }); + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(80); + expect(fake.cancellationAttempts).toBeLessThanOrEqual(4); +}); + +test('a pending Redis write allocates no retry backlog and waits until the deadline', async () => { + const fake = new FakeRedis(); + let calls = 0; + fake.eval = async () => { + calls += 1; + return new Promise(() => {}); + }; + const startedAt = Date.now(); + await fenceJobCancellation({ + commands: redis(fake), + target: { queueName: 'other', jobId: 'pending-write' }, + ttlSeconds: 60, + deadlineAtMs: startedAt + 40, + }); + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(39); + expect(calls).toBe(1); +}); + +test('queued removal frees waiting capacity and tolerates an activation race', async () => { + let removals = 0; + expect( + await removeJobIfWaiting({ + getState: async () => 'waiting', + remove: async () => { + removals += 1; + }, + }), + ).toBe(true); + expect( + await removeJobIfWaiting({ + getState: async () => 'waiting', + remove: async () => { + removals += 1; + throw new Error('job is active'); + }, + }), + ).toBe(false); + expect(removals).toBe(2); +}); diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts new file mode 100644 index 00000000..4da3253d --- /dev/null +++ b/service/src/job-cancellation.ts @@ -0,0 +1,664 @@ +import type IORedis from 'ioredis'; +import type { Job, QueueEvents } from 'bullmq'; + +const JOB_CANCELLATION_PREFIX = 'codeapi:job-cancellation:v1'; +const JOB_CANCELLATION_CHANNEL = `${JOB_CANCELLATION_PREFIX}:events`; +export const CLIENT_DISCONNECT_REASON = 'client_disconnected'; +export const JOB_CANCELLED_MESSAGE = 'Job cancelled after client disconnected'; + +interface JobTarget { + queueName: string; + jobId: string; +} + +function targetKey(target: JobTarget): string { + return `${target.queueName}:${target.jobId}`; +} + +function cancellationKey(target: JobTarget): string { + return `${JOB_CANCELLATION_PREFIX}:${encodeURIComponent( + target.queueName, + )}:${encodeURIComponent(target.jobId)}`; +} + +function parseTarget(raw: string): JobTarget | undefined { + try { + const parsed = JSON.parse(raw) as Partial; + if ( + typeof parsed.queueName !== 'string' || + parsed.queueName.length === 0 || + parsed.queueName.length > 256 || + typeof parsed.jobId !== 'string' || + parsed.jobId.length === 0 || + parsed.jobId.length > 256 + ) { + return undefined; + } + return { queueName: parsed.queueName, jobId: parsed.jobId }; + } catch { + return undefined; + } +} + +/** + * Cross-process cancellation for BullMQ work. + * + * The durable marker closes the publish-before-subscribe race while one + * process-wide pub/sub connection makes active cancellation O(events), not + * O(active jobs) Redis polling. Only explicitly cancellable replay jobs use + * this path, so ordinary queue traffic pays no extra Redis round trips. + */ +export class JobCancellationRegistry { + private subscriber?: IORedis; + private readonly controllers = new Map< + string, + { target: JobTarget; controllers: Set } + >(); + private startPromise?: Promise; + private readonly subscriberEndHandlers = new WeakMap void>(); + private reconcileTimer?: ReturnType; + private subscriberRestartTimer?: ReturnType; + private reconcileRetryMs = 100; + private closed = false; + + constructor(private readonly commands: IORedis) {} + + private readonly onSubscriberError = (): void => { + // ioredis reconnects using the shared policy. The listener prevents a + // transient subscriber outage from becoming an uncaught process error. + }; + + private readonly onSubscriberReady = (): void => { + this.scheduleReconcile(0); + }; + + private readonly onSubscriberMessage = ( + channel: string, + raw: string, + ): void => { + if (channel !== JOB_CANCELLATION_CHANNEL) return; + const target = parseTarget(raw); + if (target == null) return; + for (const controller of this.controllers.get(targetKey(target)) + ?.controllers ?? []) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + }; + + private detachSubscriber(subscriber: IORedis): void { + subscriber.removeListener('error', this.onSubscriberError); + subscriber.removeListener('ready', this.onSubscriberReady); + subscriber.removeListener('message', this.onSubscriberMessage); + const onEnd = this.subscriberEndHandlers.get(subscriber); + if (onEnd != null) subscriber.removeListener('end', onEnd); + this.subscriberEndHandlers.delete(subscriber); + } + + private restartAfterTerminalDisconnect(subscriber: IORedis): void { + if (this.closed || this.subscriber !== subscriber) return; + this.detachSubscriber(subscriber); + this.subscriber = undefined; + this.startPromise = undefined; + if (this.controllers.size === 0) return; + void this.start().then( + () => this.scheduleReconcile(0), + () => this.scheduleSubscriberRestart(), + ); + } + + private scheduleSubscriberRestart(): void { + if ( + this.closed || + this.controllers.size === 0 || + this.startPromise != null || + this.subscriberRestartTimer != null + ) + return; + const retryMs = this.reconcileRetryMs; + this.reconcileRetryMs = Math.min(2_000, retryMs * 2); + this.subscriberRestartTimer = setTimeout(() => { + this.subscriberRestartTimer = undefined; + if ( + this.closed || + this.controllers.size === 0 || + this.startPromise != null + ) + return; + void this.start().then( + () => { + this.reconcileRetryMs = 100; + this.scheduleReconcile(0); + }, + () => this.scheduleSubscriberRestart(), + ); + }, retryMs); + } + + private async reconcile(): Promise { + const entries = [...this.controllers.values()]; + if (entries.length === 0) return; + const cancelled = await this.commands.mget( + ...entries.map(({ target }) => cancellationKey(target)), + ); + cancelled.forEach((value, index) => { + if (value === '1') { + for (const controller of entries[index]?.controllers ?? []) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + } + }); + } + + private scheduleReconcile(delayMs: number): void { + if ( + this.closed || + this.controllers.size === 0 || + this.reconcileTimer != null + ) { + return; + } + this.reconcileTimer = setTimeout(() => { + this.reconcileTimer = undefined; + void this.reconcile().then( + () => { + this.reconcileRetryMs = 100; + }, + () => { + const retryMs = this.reconcileRetryMs; + this.reconcileRetryMs = Math.min(2_000, retryMs * 2); + this.scheduleReconcile(retryMs); + }, + ); + }, delayMs); + } + + private start(): Promise { + if (this.closed) { + return Promise.reject(new Error('Job cancellation registry is closed')); + } + if (this.startPromise != null) return this.startPromise; + const starting = (async (): Promise => { + const subscriber = this.commands.duplicate(); + this.subscriber = subscriber; + subscriber.on('error', this.onSubscriberError); + subscriber.on('ready', this.onSubscriberReady); + subscriber.on('message', this.onSubscriberMessage); + const onEnd = (): void => this.restartAfterTerminalDisconnect(subscriber); + this.subscriberEndHandlers.set(subscriber, onEnd); + subscriber.on('end', onEnd); + try { + await subscriber.subscribe(JOB_CANCELLATION_CHANNEL); + } catch (error) { + this.detachSubscriber(subscriber); + if (this.subscriber === subscriber) this.subscriber = undefined; + subscriber.disconnect(false); + throw error; + } + })(); + this.startPromise = starting; + void starting.catch(() => { + if (this.startPromise === starting) this.startPromise = undefined; + }); + return starting; + } + + async register( + target: JobTarget, + controller: AbortController, + ): Promise { + const key = targetKey(target); + const entry = this.controllers.get(key) ?? { + target, + controllers: new Set(), + }; + entry.controllers.add(controller); + this.controllers.set(key, entry); + try { + await this.start(); + if ((await this.commands.get(cancellationKey(target))) === '1') { + controller.abort(CLIENT_DISCONNECT_REASON); + } + } catch (error) { + entry.controllers.delete(controller); + if (entry.controllers.size === 0) this.controllers.delete(key); + throw error; + } + } + + async unregister( + target: JobTarget, + controller?: AbortController, + ): Promise { + const key = targetKey(target); + const entry = this.controllers.get(key); + if (controller == null) { + this.controllers.delete(key); + } else if (entry != null) { + entry.controllers.delete(controller); + if (entry.controllers.size === 0) this.controllers.delete(key); + } + // Markers expire by TTL. Deleting one here can erase the only evidence + // needed by another replica whose subscriber was reconnecting. + } + + async close(): Promise { + this.closed = true; + this.controllers.clear(); + if (this.reconcileTimer != null) clearTimeout(this.reconcileTimer); + this.reconcileTimer = undefined; + if (this.subscriberRestartTimer != null) + clearTimeout(this.subscriberRestartTimer); + this.subscriberRestartTimer = undefined; + const subscriber = this.subscriber; + this.subscriber = undefined; + this.startPromise = undefined; + if (subscriber == null) return; + this.detachSubscriber(subscriber); + // This socket only carries notifications. Disconnect it before awaiting + // anything: subscribe() may be queued through an indefinite Redis outage. + subscriber.disconnect(false); + } +} + +async function cancelJobInRedis( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, + includeResult: boolean, +): Promise { + // Cancellation and result publication have ONE durable winner. Pub/sub is + // only a notification; it must not decide whether Stop was accepted. + return commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == 'completed' then + -- Keep completion evidence at least as long as the requesting process + -- requires, even across API/worker config differences. Attached request + -- tombstones never renew independently of this decision. + local requestedTtlMs = tonumber(ARGV[1]) * 1000 + for i = 1, 2 do + if redis.call('PTTL', KEYS[i]) < requestedTtlMs then + redis.call('PEXPIRE', KEYS[i], requestedTtlMs) + end + end + if ARGV[4] == '1' then return {0, redis.call('GET', KEYS[2])} end + return {0} + end + if state and state ~= '1' then return {-1} end + redis.call('SET', KEYS[1], '1', 'EX', ARGV[1]) + redis.call('PUBLISH', ARGV[2], ARGV[3]) + return {1} + `, + 2, + cancellationKey(target), + `${cancellationKey(target)}:result`, + Math.max(1, ttlSeconds), + JOB_CANCELLATION_CHANNEL, + JSON.stringify(target), + includeResult ? '1' : '0', + ); +} + +export async function requestJobCancellation( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, +): Promise { + const decision = await cancelJobInRedis(commands, target, ttlSeconds, false); + if (!Array.isArray(decision) || ![0, 1].includes(decision[0])) { + throw new Error('Invalid durable cancellation decision'); + } + return decision[0] === 1; +} + +export type JobFenceOutcome = + | { status: 'cancelled' | 'expired' } + | { status: 'completed'; result: T }; + +function decodeCommittedResult(value: unknown): { result: T } { + if (typeof value !== 'string') { + throw new Error( + 'Committed programmatic result expired; refusing re-execution', + ); + } + const decoded: unknown = JSON.parse(value); + if ( + decoded == null || + typeof decoded !== 'object' || + !Object.prototype.hasOwnProperty.call(decoded, 'result') + ) { + throw new Error( + 'Invalid committed programmatic result; refusing re-execution', + ); + } + return decoded as { result: T }; +} + +export function jobCancellationRetentionSeconds( + localTimeoutMs: number, + producerTtlSeconds = 0, +): number { + return Math.max( + Math.ceil(localTimeoutMs / 1_000) * 2 + 180, + Number.isFinite(producerTtlSeconds) ? producerTtlSeconds : 0, + ); +} + +/** Retain the actual result so a BullMQ retry after a lost completion reply + * cannot repeat sandbox mutations. Keep the status small: reconnect MGETs must + * never load every active job's output into each API/worker replica. */ +export async function commitJobResult( + commands: IORedis, + target: JobTarget, + result: T, + ttlSeconds: number, + deadlineAtMs = Number.MAX_SAFE_INTEGER, +): Promise<'committed' | 'cancelled' | 'already_completed'> { + const serialized = JSON.stringify({ result }); + if (Buffer.byteLength(serialized) > 16 * 1024 * 1024) { + throw new Error('Programmatic completion exceeds the 16 MiB result limit'); + } + const decision = await commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == '1' then return 0 end + if state == 'completed' then return 2 end + if state then return -2 end + if not state then + local now = redis.call('TIME') + if tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) >= tonumber(ARGV[3]) then + return -1 + end + -- One write command, so an OOM cannot publish just half the decision. + redis.call('MSET', KEYS[1], 'completed', KEYS[2], ARGV[1]) + redis.call('EXPIRE', KEYS[1], ARGV[2]) + redis.call('EXPIRE', KEYS[2], ARGV[2]) + end + return 1 + `, + 2, + cancellationKey(target), + `${cancellationKey(target)}:result`, + serialized, + Math.max(1, ttlSeconds), + deadlineAtMs, + ); + if (decision === -1) + throw new Error('Job result commitment exceeded its deadline'); + if (decision === 1) return 'committed'; + if (decision === 0) return 'cancelled'; + if (decision === 2) return 'already_completed'; + throw new Error('Invalid durable result commitment'); +} + +/** BullMQ lock loss can redeliver a job while its first processor still runs. + * Claim once before any sandbox work, retaining the claim through the job's + * recovery horizon. An ambiguous/stalled attempt is never permission to rerun. + * Completion lookup and claim are atomic, so there is no read-then-start gap. */ +export async function claimJobExecution( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, +): Promise<{ status: 'claimed' } | { status: 'completed'; result: T }> { + const key = cancellationKey(target); + const decision = await commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == 'completed' then return {0, redis.call('GET', KEYS[2])} end + if state == '1' then return {-1} end + if state then return {-3} end + if redis.call('SET', KEYS[3], '1', 'NX', 'EX', ARGV[1]) then return {1} end + return {-2} + `, + 3, + key, + `${key}:result`, + `${key}:execution`, + Math.max(1, ttlSeconds), + ); + if (!Array.isArray(decision)) throw new Error('Invalid execution claim'); + if (decision[0] === 1) return { status: 'claimed' }; + if (decision[0] === 0) + return { + status: 'completed', + result: decodeCommittedResult(decision[1]).result, + }; + if (decision[0] === -1) throw new Error(JOB_CANCELLED_MESSAGE); + if (decision[0] === -2) + throw new Error( + 'Programmatic job already claimed; refusing duplicate execution', + ); + throw new Error('Invalid durable execution claim'); +} + +export async function readCommittedJobResult( + commands: IORedis, + target: JobTarget, +): Promise<{ result: T } | undefined> { + const [state, value] = await commands.mget( + cancellationKey(target), + `${cancellationKey(target)}:result`, + ); + if (state !== 'completed') return undefined; + return decodeCommittedResult(value); +} + +/** Do not release replay ownership on an ambiguous Redis failure. Keep one + * outstanding marker write, retry rejected writes with bounded backoff, and + * retain ownership until it succeeds or the job's ORIGINAL deadline expires. + * A delayed queue.add must carry that same timestamp into the worker. */ +export async function fenceJobCancellation(args: { + commands: IORedis; + target: JobTarget; + ttlSeconds: number; + deadlineAtMs: number; +}): Promise> { + let retryMs = 25; + let firstAttempt = true; + while (firstAttempt || Date.now() < args.deadlineAtMs) { + firstAttempt = false; + // If a lost enqueue reply arrives after the execution deadline, still + // give a healthy Redis one bounded opportunity to return completion's + // winning decision. Never translate a known committed effect to failure. + const remainingMs = args.deadlineAtMs - Date.now(); + let timer: ReturnType | undefined; + let decision: unknown; + try { + decision = await Promise.race([ + cancelJobInRedis(args.commands, args.target, args.ttlSeconds, true), + new Promise(resolve => { + timer = setTimeout( + () => resolve(undefined), + remainingMs > 0 ? remainingMs : 1_000, + ); + }), + ]); + } catch { + await new Promise(resolve => + setTimeout( + resolve, + Math.min(retryMs, Math.max(0, args.deadlineAtMs - Date.now())), + ), + ); + retryMs = Math.min(1_000, retryMs * 2); + continue; + } finally { + if (timer != null) clearTimeout(timer); + } + // Only transport failures retry. Corrupt/missing durable results are + // deterministic invariant failures, not an invitation to extend their TTL. + if (decision === undefined) return { status: 'expired' }; + if (!Array.isArray(decision)) + throw new Error('Invalid durable cancellation decision'); + if (decision[0] === 1) return { status: 'cancelled' }; + if (decision[0] === 0) + return { + status: 'completed', + result: decodeCommittedResult(decision[1]).result, + }; + throw new Error('Invalid durable cancellation decision'); + } + return { status: 'expired' }; +} + +const REMOVABLE_JOB_STATES = new Set([ + 'waiting', + 'delayed', + 'prioritized', + 'waiting-children', +]); + +/** Frees queued capacity without ever removing an active or settled job. */ +export async function removeJobIfWaiting( + job: Pick, +): Promise { + if (!REMOVABLE_JOB_STATES.has(await job.getState())) return false; + try { + await job.remove(); + return true; + } catch { + // A worker may have activated the job between getState() and remove(). + // The durable marker remains authoritative for that race. + return false; + } +} + +export function programmaticCancellationError(): Error { + return new DOMException( + 'Programmatic execution request disconnected', + 'AbortError', + ); +} + +/** Commit barrier for result-processing stages that may yield after execution. */ +export function throwIfJobAborted(signal: AbortSignal): void { + if (!signal.aborted) return; + if (signal.reason instanceof Error) throw signal.reason; + throw new DOMException( + typeof signal.reason === 'string' ? signal.reason : 'Job aborted', + 'AbortError', + ); +} + +/** Maps cancellation observed during asynchronous result cleanup to the same + * stable worker failure used by the main execution catch path. */ +export function jobResultCommitFailure( + signal: AbortSignal, + jobTimeoutMs: number, +): Error | undefined { + if (!signal.aborted) return undefined; + return new Error( + signal.reason === CLIENT_DISCONNECT_REASON + ? JOB_CANCELLED_MESSAGE + : `Job timed out after ${jobTimeoutMs}ms`, + ); +} + +export async function waitForJobWithCancellation(args: { + commands: IORedis; + registry: JobCancellationRegistry; + job: Job; + events: QueueEvents; + timeoutMs: number; + cancellationTtlSeconds: number; + deadlineAtMs?: number; + signal?: AbortSignal; +}): Promise { + const { + commands, + registry, + job, + events, + timeoutMs, + cancellationTtlSeconds, + signal, + } = args; + const completion = job.waitUntilFinished(events, timeoutMs); + // Subscription startup can itself wait for Redis recovery. Own the losing + // promise immediately, before any await, rather than after registration. + void completion.catch(() => undefined); + const target = { queueName: job.queueName, jobId: String(job.id) }; + const deadlineAtMs = args.deadlineAtMs ?? Date.now() + timeoutMs; + let fencing: Promise> | undefined; + const fence = (): Promise> => + (fencing ??= fenceJobCancellation({ + commands, + target, + ttlSeconds: cancellationTtlSeconds, + deadlineAtMs, + })); + const externalController = new AbortController(); + try { + await registry.register(target, externalController); + } catch (error) { + void completion.catch(() => undefined); + const outcome = await fence(); + if (outcome.status === 'completed') return outcome.result; + await removeJobIfWaiting(job).catch(() => false); + throw error; + } + + let removeAbortListener = (): void => {}; + const disconnected = new Promise((resolve, reject) => { + let cancelling = false; + const cancel = (): void => { + if (cancelling) return; + cancelling = true; + void fence() + .then(async outcome => { + if (outcome.status === 'completed') { + resolve(outcome.result); + return; + } + // Removing a waiting job immediately frees queue capacity. An active + // job cannot be removed; its worker observes the durable marker or + // pub/sub event and aborts the sandbox transport instead. + await removeJobIfWaiting(job).catch(() => false); + reject(programmaticCancellationError()); + }) + .catch(reject); + }; + if (signal != null) { + removeAbortListener = (): void => + signal.removeEventListener('abort', cancel); + signal.addEventListener('abort', cancel, { once: true }); + if (signal.aborted) cancel(); + } + }); + const cancelled = new Promise((_, reject) => { + const cancel = (): void => { + void removeJobIfWaiting(job).then( + () => reject(programmaticCancellationError()), + () => reject(programmaticCancellationError()), + ); + }; + externalController.signal.addEventListener('abort', cancel, { + once: true, + }); + if (externalController.signal.aborted) cancel(); + }); + + // A cancelled request stops awaiting the BullMQ result, so attach a sink to + // the losing promise before racing it to avoid an unhandled late rejection. + void completion.catch(() => undefined); + try { + return await Promise.race([completion, disconnected, cancelled]); + } catch (error) { + // Includes waitUntilFinished timeouts and registration/transport errors, + // not only explicit Stop. Replay cleanup is unsafe until this barrier. + const outcome = await fence(); + if (outcome.status === 'completed') return outcome.result; + throw error; + } finally { + removeAbortListener(); + await registry + .unregister(target, externalController) + .catch(() => undefined); + } +} + +export const jobCancellationInternals = { + channel: JOB_CANCELLATION_CHANNEL, + cancellationKey, + parseTarget, +}; diff --git a/service/src/metrics.ts b/service/src/metrics.ts index adfd9872..42c5536f 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -117,6 +117,12 @@ export const jobsFailed = new Counter({ labelNames: ['language'] as const, }); +export const jobsCancelled = new Counter({ + name: 'codeapi_jobs_cancelled_total', + help: 'Total number of jobs cancelled after the calling client disconnected', + labelNames: ['language'] as const, +}); + export const activeJobs = new Gauge({ name: 'codeapi_active_jobs', help: 'Number of jobs currently being processed', diff --git a/service/src/middleware/limits.ts b/service/src/middleware/limits.ts index 099261a1..9ffaeb9d 100644 --- a/service/src/middleware/limits.ts +++ b/service/src/middleware/limits.ts @@ -182,6 +182,19 @@ export const executionLimiter = createRateLimiter( } ); +/** Keep Stop available when execution admission is full, while independently + * bounding request-id churn in the cancellation registry. */ +export const cancellationLimiter = createRateLimiter( + 'exec-cancel', + env.EXEC_LIMIT_WINDOW, + Math.max(80, env.EXEC_MAX_REQUESTS * 4), + { + message: 'Too many CodeAPI cancellation requests.', + structuredBody: true, + logRejections: true, + } +); + export const uploadLimiter = createRateLimiter( 'upload', env.UPLOAD_LIMIT_WINDOW, diff --git a/service/src/preamble-bash.test.ts b/service/src/preamble-bash.test.ts index 815040d3..e110abdc 100644 --- a/service/src/preamble-bash.test.ts +++ b/service/src/preamble-bash.test.ts @@ -1,10 +1,19 @@ import { execFileSync } from 'child_process'; -import { mkdtempSync, rmSync, writeFileSync } from 'fs'; +import { + mkdirSync, + mkdtempSync, + readFileSync, + rmSync, + writeFileSync, +} from 'fs'; import { tmpdir } from 'os'; import { join } from 'path'; import { describe, expect, test } from 'bun:test'; import { extractPendingFromStdout, type LCTool } from './preamble'; -import { generateBashReplayPostamble, generateBashReplayPreamble } from './preamble-bash'; +import { + generateBashReplayPostamble, + generateBashReplayPreamble, +} from './preamble-bash'; interface BashRunResult { stdout: string; @@ -60,9 +69,13 @@ function assemble(userCode: string, toolSet: LCTool[] = tools): string { ].join('\n'); } -function runBash(script: string, options: number | BashRunOptions = {}): BashRunResult { - const timeoutMs = typeof options === 'number' ? options : options.timeoutMs ?? 3000; - const history = typeof options === 'number' ? {} : options.history ?? {}; +function runBash( + script: string, + options: number | BashRunOptions = {}, +): BashRunResult { + const timeoutMs = + typeof options === 'number' ? options : (options.timeoutMs ?? 3000); + const history = typeof options === 'number' ? {} : (options.history ?? {}); const dir = mkdtempSync(join(tmpdir(), 'ptc-bash-unit-')); const file = join(dir, 'main.sh'); const historyPath = join(dir, 'history.json'); @@ -99,46 +112,144 @@ function pendingNames(stdout: string): string[] { return (parsed.pending ?? []).map(call => call.tool_name).sort(); } +describe('generateBashReplayPreamble - private runtime directory', () => { + test('creates every replay tempfile beneath TMPDIR', () => { + const dir = mkdtempSync(join(tmpdir(), 'ptc-bash-private-tmp-')); + const dataDir = join(dir, 'data'); + const runtimeDir = join(dir, 'runtime'); + const file = join(dir, 'main.sh'); + const historyPath = join(dataDir, 'history.json'); + mkdirSync(dataDir, { recursive: true }); + mkdirSync(runtimeDir, { recursive: true }); + writeFileSync(historyPath, '{}'); + writeFileSync( + file, + assemble(` +printf '%s\\n' "$_PTC_PENDING_FILE" "$_PTC_ERROR_FILE" "$_PTC_COUNTER_FILE" +`), + { mode: 0o755 }, + ); + + try { + const stdout = execFileSync('bash', [file], { + env: { + ...process.env, + PTC_HISTORY_PATH: historyPath, + TMPDIR: runtimeDir, + }, + encoding: 'utf8', + }); + const paths = stdout.trim().split('\n'); + expect(paths).toHaveLength(3); + expect( + paths.every(value => value.startsWith(`${runtimeDir}/`)), + ).toBe(true); + expect( + generateBashReplayPreamble({ executionId, tools }), + ).not.toContain('mktemp -t'); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); + + test('persists pending calls through the private native control path', () => { + const dir = mkdtempSync(join(tmpdir(), 'ptc-bash-control-')); + const file = join(dir, 'main.sh'); + const historyPath = join(dir, 'history.json'); + const controlPath = join(dir, 'control.json'); + writeFileSync(file, assemble(`get_weather '{"city":"Paris"}'`), { + mode: 0o755, + }); + writeFileSync(historyPath, '{}'); + try { + execFileSync('bash', [file], { + env: { + ...process.env, + PTC_HISTORY_PATH: historyPath, + LIBRECHAT_CODE_CONTROL_PATH: controlPath, + TMPDIR: dir, + }, + encoding: 'utf8', + }); + expect(JSON.parse(readFileSync(controlPath, 'utf8'))).toMatchObject( + { + pending: [ + { + call_id: 'call_001', + tool_name: 'get_weather', + input: { city: 'Paris' }, + }, + ], + }, + ); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); + + test('uses the trusted jq path instead of resolving jq through PATH', () => { + const preamble = generateBashReplayPreamble({ executionId, tools }); + expect(preamble).toContain( + '_PTC_JQ_PATH="${LIBRECHAT_CODE_JQ_PATH:-jq}"', + ); + expect(preamble).not.toMatch(/(^|[|;(]\s*)jq\s/m); + }); +}); + describe('generateBashReplayPreamble - command substitution pending emission', () => { test('emits ClickHouse-style object input with SQL quotes from double-quoted JSON', () => { - const run = runBash(assemble(` + const run = runBash( + assemble( + ` SVC="45886e06-932b-4cff-bb49-3f7281d80717" result=$(run_select_query_mcp_ClickHouse "{\\"serviceId\\":\\"$SVC\\",\\"query\\":\\"SELECT name, round(avg(tempAvg)/10.0, 2) AS avg_temp_c FROM system.columns WHERE database='default' AND table='uk_prices_3' AND tempAvg != -9999\\"}") echo "AFTER: $result" -`, [clickHouseTool])); +`, + [clickHouseTool], + ), + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); expect(parsed.pending).toHaveLength(1); - expect(parsed.pending?.[0]?.tool_name).toBe('run_select_query_mcp_ClickHouse'); + expect(parsed.pending?.[0]?.tool_name).toBe( + 'run_select_query_mcp_ClickHouse', + ); expect(parsed.pending?.[0]?.input).toEqual({ serviceId: '45886e06-932b-4cff-bb49-3f7281d80717', - query: - "SELECT name, round(avg(tempAvg)/10.0, 2) AS avg_temp_c FROM system.columns WHERE database='default' AND table='uk_prices_3' AND tempAvg != -9999", + query: "SELECT name, round(avg(tempAvg)/10.0, 2) AS avg_temp_c FROM system.columns WHERE database='default' AND table='uk_prices_3' AND tempAvg != -9999", }); expect(parsed.stdout).not.toContain('AFTER'); }); test('emits ClickHouse-style object input with shell-escaped SQL quotes', () => { - const run = runBash(assemble(` + const run = runBash( + assemble( + ` result=$(run_select_query_mcp_ClickHouse '{"serviceId":"45886e06-932b-4cff-bb49-3f7281d80717","query":"SELECT name, type FROM system.columns WHERE database='"'"'default'"'"' AND table='"'"'uk_prices_3'"'"' ORDER BY position"}') echo "AFTER: $result" -`, [clickHouseTool])); +`, + [clickHouseTool], + ), + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); expect(parsed.pending).toHaveLength(1); - expect(parsed.pending?.[0]?.tool_name).toBe('run_select_query_mcp_ClickHouse'); + expect(parsed.pending?.[0]?.tool_name).toBe( + 'run_select_query_mcp_ClickHouse', + ); expect(parsed.pending?.[0]?.input).toEqual({ serviceId: '45886e06-932b-4cff-bb49-3f7281d80717', - query: - "SELECT name, type FROM system.columns WHERE database='default' AND table='uk_prices_3' ORDER BY position", + query: "SELECT name, type FROM system.columns WHERE database='default' AND table='uk_prices_3' ORDER BY position", }); expect(parsed.stdout).not.toContain('AFTER'); }); test('batches parallel ClickHouse-style command substitutions into one pending block', () => { - const run = runBash(assemble(` + const run = runBash( + assemble( + ` SVC="45886e06-932b-4cff-bb49-3f7281d80717" { @@ -158,7 +269,11 @@ SVC="45886e06-932b-4cff-bb49-3f7281d80717" wait echo "AFTER" -`, [clickHouseTool]), 3000); +`, + [clickHouseTool], + ), + 3000, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); @@ -168,7 +283,11 @@ echo "AFTER" 'run_select_query_mcp_ClickHouse', 'run_select_query_mcp_ClickHouse', ]); - expect(parsed.pending?.map(call => (call.input as { query: string }).query).sort()).toEqual([ + expect( + parsed.pending + ?.map(call => (call.input as { query: string }).query) + .sort(), + ).toEqual([ "SELECT name, engine, total_rows, formatReadableSize(total_bytes) AS size, sorting_key, partition_key FROM system.tables WHERE database='default' AND name IN ('uk_prices_3','weather_noaa_mt')", "SELECT name, type, comment FROM system.columns WHERE database='default' AND table='uk_prices_3' ORDER BY position", "SELECT name, type, comment FROM system.columns WHERE database='default' AND table='weather_noaa_mt' ORDER BY position", @@ -177,12 +296,14 @@ echo "AFTER" }); test('emits a command-substitution tool call before later user code while another job is running', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` sleep 0.2 & result=$(get_weather '{"city":"Madrid"}') echo "AFTER: $result" wait -`)); +`), + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -194,12 +315,15 @@ wait }); test('batches background and command-substitution tool calls before command-substitution side effects', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Oslo"}' & result=$(calculate '{"expression":"2+3"}') echo "SIDE_EFFECT: $result" wait -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -210,10 +334,13 @@ wait }); test('waits for background compound commands that invoke tools later', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` (sleep 0.2; get_weather '{"city":"Paris"}') & echo "AFTER LAUNCH" -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -225,10 +352,13 @@ echo "AFTER LAUNCH" }); test('does not wait for unrelated background commands with tool names as arguments', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` bash -c 'sleep 2' get_weather & echo "DONE" -`), 700); +`), + 700, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -238,14 +368,17 @@ echo "DONE" }); test('does not treat arithmetic expansion as command substitution while batching background tools', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Oslo"}' & sleep 0.1 x=$((1+1)) calculate '{"expression":"2+3"}' & wait echo "DONE $x" -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -256,12 +389,15 @@ echo "DONE $x" }); test('handles backtick command substitution without waiting for unrelated background jobs', () => { - const run = runBash(assemble(` + const run = runBash( + assemble(` sleep 5 & result=\`get_weather '{"city":"Porto"}'\` echo "AFTER: $result" wait -`), 1500); +`), + 1500, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -281,7 +417,10 @@ wait echo "DONE" `; const firstRun = runBash(assemble(userCode)); - const firstParsed = extractPendingFromStdout(firstRun.stdout, executionId); + const firstParsed = extractPendingFromStdout( + firstRun.stdout, + executionId, + ); expect(firstRun.exitCode).toBe(0); expect(firstParsed.pending).toHaveLength(2); @@ -302,7 +441,10 @@ echo "DONE" }), ); const replayRun = runBash(assemble(userCode), { history }); - const replayParsed = extractPendingFromStdout(replayRun.stdout, executionId); + const replayParsed = extractPendingFromStdout( + replayRun.stdout, + executionId, + ); expect(replayRun.exitCode).toBe(0); expect(replayParsed.pending).toBeNull(); expect(replayParsed.stdout).toContain('"slot":"first"'); @@ -326,12 +468,15 @@ echo "DONE" }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' get_weather '{"city":"Paris"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); @@ -352,12 +497,15 @@ printf '\\nDONE\\n' }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' calculate '{"expression":"2+3"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.signal).not.toBe('SIGTERM'); @@ -385,12 +533,15 @@ printf '\\nDONE\\n' }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' calculate '{"expression":"2+3"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); @@ -418,12 +569,15 @@ printf '\\nDONE\\n' }, }; - const run = runBash(assemble(` + const run = runBash( + assemble(` get_weather '{"city":"Paris"}' printf '\\n' calculate '{"expression":"2+3"}' printf '\\nDONE\\n' -`), { history }); +`), + { history }, + ); const parsed = extractPendingFromStdout(run.stdout, executionId); expect(run.exitCode).toBe(0); diff --git a/service/src/preamble-bash.ts b/service/src/preamble-bash.ts index 4ec3fc05..31ed4c6b 100644 --- a/service/src/preamble-bash.ts +++ b/service/src/preamble-bash.ts @@ -1,8 +1,5 @@ import type { LCTool } from './preamble'; -import { - buildScopedSentinel, - PTC_HISTORY_SANDBOX_PATH, -} from './ptc-constants'; +import { buildScopedSentinel, PTC_HISTORY_SANDBOX_PATH } from './ptc-constants'; export interface BashReplayPreambleConfig { executionId: string; @@ -42,12 +39,55 @@ export class BashToolNameCollisionError extends Error { } const BASH_RESERVED = new Set([ - 'if', 'then', 'else', 'elif', 'fi', 'case', 'esac', 'for', 'select', - 'while', 'until', 'do', 'done', 'in', 'function', 'time', 'coproc', - 'return', 'exit', 'break', 'continue', 'shift', 'export', 'readonly', - 'local', 'declare', 'typeset', 'unset', 'alias', 'unalias', 'source', - 'echo', 'printf', 'read', 'cd', 'pwd', 'kill', 'trap', 'wait', 'eval', - 'exec', 'jobs', 'bg', 'fg', 'set', 'let', 'test', 'true', 'false', + 'if', + 'then', + 'else', + 'elif', + 'fi', + 'case', + 'esac', + 'for', + 'select', + 'while', + 'until', + 'do', + 'done', + 'in', + 'function', + 'time', + 'coproc', + 'return', + 'exit', + 'break', + 'continue', + 'shift', + 'export', + 'readonly', + 'local', + 'declare', + 'typeset', + 'unset', + 'alias', + 'unalias', + 'source', + 'echo', + 'printf', + 'read', + 'cd', + 'pwd', + 'kill', + 'trap', + 'wait', + 'eval', + 'exec', + 'jobs', + 'bg', + 'fg', + 'set', + 'let', + 'test', + 'true', + 'false', ]); function normalizeBashFunctionName(name: string): string { @@ -60,10 +100,7 @@ function normalizeBashFunctionName(name: string): string { * the end-of-preamble `readonly -f` lockdown runs. Compared case- * insensitively because the `_PTC_` prefix is used for variables and * `_ptc_` for functions, and both live in the same identifier space. */ - if ( - BASH_RESERVED.has(normalized) || - /^_ptc_/i.test(normalized) - ) { + if (BASH_RESERVED.has(normalized) || /^_ptc_/i.test(normalized)) { normalized = normalized + '_tool'; } if (normalized === '') normalized = 'tool'; @@ -95,9 +132,12 @@ function escapeForBashEre(s: string): string { * Users capture results via command substitution; input is passed as a single * JSON object string argument (validated by jq). */ -export function generateBashReplayPreamble(config: BashReplayPreambleConfig): string { +export function generateBashReplayPreamble( + config: BashReplayPreambleConfig, +): string { const { executionId, tools } = config; - const { start: scopedStart, end: scopedEnd } = buildScopedSentinel(executionId); + const { start: scopedStart, end: scopedEnd } = + buildScopedSentinel(executionId); let preamble = `#!/bin/bash # ============================================================================ @@ -109,20 +149,26 @@ _PTC_EXECUTION_ID="${executionId}" _PTC_SENTINEL_START="${scopedStart}" _PTC_SENTINEL_END="${scopedEnd}" _PTC_HISTORY_PATH="\${PTC_HISTORY_PATH:-${PTC_HISTORY_SANDBOX_PATH}}" -_PTC_PENDING_FILE="$(mktemp -t _ptc_pending.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_pending.XXXXXX)" -_PTC_ERROR_FILE="$(mktemp -t _ptc_error.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_error.XXXXXX)" -_PTC_CONSUMED_FILE="$(mktemp -t _ptc_consumed.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_consumed.XXXXXX)" -_PTC_SAW_BARE_TOOL_FILE="$(mktemp -t _ptc_saw_tool.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_saw_tool.XXXXXX)" -_PTC_PRE_TOOL_JOBS_FILE="$(mktemp -t _ptc_pre_tool_jobs.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_pre_tool_jobs.XXXXXX)" -_PTC_PRE_TOOL_JOBS_READY_FILE="$(mktemp -t _ptc_pre_tool_jobs_ready.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_pre_tool_jobs_ready.XXXXXX)" -_PTC_TOOL_JOBS_FILE="$(mktemp -t _ptc_tool_jobs.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_tool_jobs.XXXXXX)" -_PTC_WAIT_RAN_FILE="$(mktemp -t _ptc_wait_ran.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_wait_ran.XXXXXX)" -_PTC_SUPPRESS_SUBSHELL_TOOL_FILE="$(mktemp -t _ptc_suppress_subshell_tool.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_suppress_subshell_tool.XXXXXX)" -_PTC_SUPPRESS_SUBSHELL_TOOL_CLEAR_FILE="$(mktemp -t _ptc_suppress_subshell_tool_clear.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_suppress_subshell_tool_clear.XXXXXX)" +_PTC_CONTROL_PATH="\${LIBRECHAT_CODE_CONTROL_PATH:-}" +_PTC_JQ_PATH="\${LIBRECHAT_CODE_JQ_PATH:-jq}" +_PTC_RUNTIME_DIR="\${TMPDIR:-/tmp}" +_ptc_mktemp() { + mktemp "\${_PTC_RUNTIME_DIR%/}/$1.XXXXXX" +} +_PTC_PENDING_FILE="$(_ptc_mktemp _ptc_pending)" +_PTC_ERROR_FILE="$(_ptc_mktemp _ptc_error)" +_PTC_CONSUMED_FILE="$(_ptc_mktemp _ptc_consumed)" +_PTC_SAW_BARE_TOOL_FILE="$(_ptc_mktemp _ptc_saw_tool)" +_PTC_PRE_TOOL_JOBS_FILE="$(_ptc_mktemp _ptc_pre_tool_jobs)" +_PTC_PRE_TOOL_JOBS_READY_FILE="$(_ptc_mktemp _ptc_pre_tool_jobs_ready)" +_PTC_TOOL_JOBS_FILE="$(_ptc_mktemp _ptc_tool_jobs)" +_PTC_WAIT_RAN_FILE="$(_ptc_mktemp _ptc_wait_ran)" +_PTC_SUPPRESS_SUBSHELL_TOOL_FILE="$(_ptc_mktemp _ptc_suppress_subshell_tool)" +_PTC_SUPPRESS_SUBSHELL_TOOL_CLEAR_FILE="$(_ptc_mktemp _ptc_suppress_subshell_tool_clear)" # Counter must persist across subshells (command substitution) so call_ids # stay deterministic across cached/uncached calls. Bash variables set in a # subshell don't propagate back, so we use a file. -_PTC_COUNTER_FILE="$(mktemp -t _ptc_counter.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_counter.XXXXXX)" +_PTC_COUNTER_FILE="$(_ptc_mktemp _ptc_counter)" _PTC_LOCK_DIR="\${_PTC_PENDING_FILE}.lock" printf '0' > "$_PTC_COUNTER_FILE" : > "$_PTC_CONSUMED_FILE" @@ -177,7 +223,7 @@ _ptc_sha256() { _ptc_hash_input() { local _ptc_canonical - _ptc_canonical=$(printf '%s' "$1" | jq -cS . 2>/dev/null) || return 1 + _ptc_canonical=$(printf '%s' "$1" | "$_PTC_JQ_PATH" -cS . 2>/dev/null) || return 1 printf '%s' "$_ptc_canonical" | _ptc_sha256 } @@ -243,7 +289,7 @@ _ptc_prune_finished_tool_jobs() { return 0 fi local _ptc_tmp_file - _ptc_tmp_file="$(mktemp -t _ptc_tool_jobs_live.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_tool_jobs_live.XXXXXX)" + _ptc_tmp_file="$(_ptc_mktemp _ptc_tool_jobs_live)" while IFS= read -r _ptc_pid; do [ -n "$_ptc_pid" ] || continue if kill -0 "$_ptc_pid" 2>/dev/null; then @@ -324,12 +370,23 @@ _ptc_maybe_emit_pending() { return 0 fi local _ptc_payload - if ! _ptc_payload=$(jq -c -s '{pending:.}' "$_PTC_PENDING_FILE" 2>/dev/null); then + if ! _ptc_payload=$("$_PTC_JQ_PATH" -c -s '{pending:.}' "$_PTC_PENDING_FILE" 2>/dev/null); then printf 'failed to serialize pending PTC tool calls\\n' >&2 _ptc_cleanup_tempfiles trap - DEBUG EXIT exit 1 fi + # Native BYOM workers use this private execution-scoped control file so a + # large stdout stream cannot truncate away the replay frame. Other + # backends continue to consume the stdout sentinel below. + if [ -n "$_PTC_CONTROL_PATH" ]; then + printf '%s' "$_ptc_payload" > "$_PTC_CONTROL_PATH" || { + printf 'failed to persist pending PTC tool calls\n' >&2 + _ptc_cleanup_tempfiles + trap - DEBUG EXIT + exit 1 + } + fi if [ "\${BASH_SUBSHELL:-0}" -eq 1 ]; then trap - DEBUG EXIT exit 0 @@ -424,7 +481,7 @@ _ptc_history_matches_by_signature() { return 0 fi # Path, not inline: large input can exceed ARG_MAX via --argjson. - jq -c \\ + "$_PTC_JQ_PATH" -c \\ --arg nm "$_ptc_name" \\ --arg site "$_ptc_call_site" \\ --arg hash "$_ptc_input_hash" \\ @@ -448,7 +505,7 @@ _ptc_first_unconsumed_history_match() { local _ptc_key while IFS= read -r _ptc_match; do [ -n "$_ptc_match" ] || continue - _ptc_key=$(printf '%s' "$_ptc_match" | jq -r '.key // empty' 2>/dev/null) + _ptc_key=$(printf '%s' "$_ptc_match" | "$_PTC_JQ_PATH" -r '.key // empty' 2>/dev/null) if [ -n "$_ptc_key" ] && ! grep -Fxq "$_ptc_key" "$_PTC_CONSUMED_FILE" 2>/dev/null; then printf '%s' "$_ptc_match" return 0 @@ -460,15 +517,15 @@ _ptc_first_unconsumed_history_match() { _ptc_print_history_entry() { local _ptc_entry="$1" local _ptc_is_err - _ptc_is_err=$(printf '%s' "$_ptc_entry" | jq -r 'if type == "object" then (.is_error // false) else false end' 2>/dev/null) + _ptc_is_err=$(printf '%s' "$_ptc_entry" | "$_PTC_JQ_PATH" -r 'if type == "object" then (.is_error // false) else false end' 2>/dev/null) if [ "$_ptc_is_err" = "true" ]; then local _ptc_msg - _ptc_msg=$(printf '%s' "$_ptc_entry" | jq -r '.error_message // "tool execution failed"' 2>/dev/null) + _ptc_msg=$(printf '%s' "$_ptc_entry" | "$_PTC_JQ_PATH" -r '.error_message // "tool execution failed"' 2>/dev/null) _ptc_write_error "$_ptc_msg" exit 1 fi local _ptc_result - _ptc_result=$(printf '%s' "$_ptc_entry" | jq -c 'if type == "object" and has("result") then .result else . end' 2>/dev/null || printf 'null') + _ptc_result=$(printf '%s' "$_ptc_entry" | "$_PTC_JQ_PATH" -c 'if type == "object" and has("result") then .result else . end' 2>/dev/null || printf 'null') printf '%s' "$_ptc_result" return 0 } @@ -479,7 +536,7 @@ _ptc_history_entry_matches_current_call() { local _ptc_input_file="$3" local _ptc_input_hash="$4" # Path, same ARG_MAX reason as above. - printf '%s' "$_ptc_entry" | jq -e \\ + printf '%s' "$_ptc_entry" | "$_PTC_JQ_PATH" -e \\ --arg nm "$_ptc_name" \\ --arg hash "$_ptc_input_hash" \\ --slurpfile inp_arr "$_ptc_input_file" \\ @@ -499,7 +556,7 @@ _ptc_call_tool() { local _ptc_call_site="\${BASH_LINENO[1]:-\${BASH_LINENO[0]:-0}}" # Reject extra trailing JSON values instead of silently dropping them. - if ! printf '%s' "$_ptc_input" | jq -e -n '[inputs] as $docs | ($docs | length) == 1 and ($docs[0] | type) == "object"' >/dev/null 2>&1; then + if ! printf '%s' "$_ptc_input" | "$_PTC_JQ_PATH" -e -n '[inputs] as $docs | ($docs | length) == 1 and ($docs[0] | type) == "object"' >/dev/null 2>&1; then _ptc_write_error "tool input for $_ptc_name must be a single JSON object, got: $_ptc_input" exit 1 fi @@ -512,7 +569,7 @@ _ptc_call_tool() { # Large input can exceed ARG_MAX via --argjson; write once, reuse path below. local _ptc_input_tmp - _ptc_input_tmp="$(mktemp -t _ptc_input.XXXXXX 2>/dev/null || mktemp /tmp/_ptc_input.XXXXXX)" + _ptc_input_tmp="$(_ptc_mktemp _ptc_input)" printf '%s' "$_ptc_input" > "$_ptc_input_tmp" local _ptc_matches @@ -527,8 +584,8 @@ _ptc_call_tool() { if [ -n "$_ptc_match" ] && [ "$_ptc_match" != "null" ]; then local _ptc_matched_call_id local _ptc_matched_entry - _ptc_matched_call_id=$(printf '%s' "$_ptc_match" | jq -r '.key' 2>/dev/null) - _ptc_matched_entry=$(printf '%s' "$_ptc_match" | jq -c '.value' 2>/dev/null) + _ptc_matched_call_id=$(printf '%s' "$_ptc_match" | "$_PTC_JQ_PATH" -r '.key' 2>/dev/null) + _ptc_matched_entry=$(printf '%s' "$_ptc_match" | "$_PTC_JQ_PATH" -c '.value' 2>/dev/null) printf '%s\\n' "$_ptc_matched_call_id" >> "$_PTC_CONSUMED_FILE" _ptc_mark_counter_at_least "$_ptc_matched_call_id" _ptc_release_lock @@ -542,7 +599,7 @@ _ptc_call_tool() { while :; do _ptc_call_id=$(_ptc_next_call_id) if [ -r "$_PTC_HISTORY_PATH" ]; then - _ptc_entry=$(jq -c --arg id "$_ptc_call_id" '.[$id] // empty' "$_PTC_HISTORY_PATH" 2>/dev/null || printf '') + _ptc_entry=$("$_PTC_JQ_PATH" -c --arg id "$_ptc_call_id" '.[$id] // empty' "$_PTC_HISTORY_PATH" 2>/dev/null || printf '') else _ptc_entry="" fi @@ -558,7 +615,7 @@ _ptc_call_tool() { fi done - if ! printf '%s' "$_ptc_input" | jq -c -n \\ + if ! printf '%s' "$_ptc_input" | "$_PTC_JQ_PATH" -c -n \\ --arg cid "$_ptc_call_id" \\ --arg nm "$_ptc_name" \\ --arg hash "$_ptc_input_hash" \\ @@ -668,8 +725,12 @@ exit $_ptc_user_exit_code function generateBashToolStub(tool: LCTool): string { const fnName = normalizeBashFunctionName(tool.name); - const desc = (tool.description ?? '').split('\n').map(l => `# ${l}`).join('\n'); - const nameComment = fnName !== tool.name ? `# Original tool name: ${tool.name}\n` : ''; + const desc = (tool.description ?? '') + .split('\n') + .map(l => `# ${l}`) + .join('\n'); + const nameComment = + fnName !== tool.name ? `# Original tool name: ${tool.name}\n` : ''; const escapedToolName = escapeForBashDoubleQuote(tool.name); return `${nameComment}${desc ? desc + '\n' : ''}${fnName}() { local _default_input='{}' @@ -683,7 +744,8 @@ function generateBashToolStub(tool: LCTool): string { } function generateBashPendingDeferHelper(tools: readonly LCTool[]): string { - const toolNamesPattern = tools + const toolNamesPattern = + tools .map(tool => normalizeBashFunctionName(tool.name)) .map(escapeForBashEre) .join('|') || 'a^'; diff --git a/service/src/preamble.test.ts b/service/src/preamble.test.ts index bf50e3ec..803e1839 100644 --- a/service/src/preamble.test.ts +++ b/service/src/preamble.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from 'bun:test'; import { buildScopedSentinel, createProgrammaticPayload, + extractPendingFromControlPayload, extractPendingFromStdout, generatePreamble, } from './preamble'; @@ -33,7 +34,9 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { expect(preamble).toMatch(/AF_UNIX/); expect(preamble).toMatch(/\.connect\(_TOOL_CALL_SOCKET\)/); /* Regression guard against reintroducing the user-spoofable check. */ - expect(preamble).not.toMatch(/if\s+os\.path\.exists\(_TOOL_CALL_SOCKET\)/); + expect(preamble).not.toMatch( + /if\s+os\.path\.exists\(_TOOL_CALL_SOCKET\)/, + ); }); test('caches the probe verdict at module load (before user code can plant a spoof)', () => { @@ -41,10 +44,14 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { /* The probe call must appear at top level of the preamble, NOT * inside _do_request. Otherwise a user could plant a regular file * at the path between calls and flip the gate per-request. */ - const probeCallIdx = preamble.indexOf('_USE_TOOL_CALL_SOCKET = _probe_tool_call_socket()'); + const probeCallIdx = preamble.indexOf( + '_USE_TOOL_CALL_SOCKET = _probe_tool_call_socket()', + ); expect(probeCallIdx).toBeGreaterThan(-1); /* _do_request must consult the cached verdict, not re-probe. */ - const doReqMatch = preamble.match(/def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/); + const doReqMatch = preamble.match( + /def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/, + ); expect(doReqMatch).not.toBeNull(); expect(doReqMatch![0]).toContain('_USE_TOOL_CALL_SOCKET'); expect(doReqMatch![0]).not.toContain('_probe_tool_call_socket('); @@ -55,7 +62,9 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { /* The fallback must still construct the URL from _CALLBACK_URL and * delegate to _tcp_request. Without this, runners without the * proxy bind-mount would have no way to reach the orchestrator. */ - const doReqMatch = preamble.match(/def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/); + const doReqMatch = preamble.match( + /def\s+_do_request[\s\S]*?(?=\ndef\s|\nclass\s|\Z)/, + ); expect(doReqMatch).not.toBeNull(); expect(doReqMatch![0]).toContain('_CALLBACK_URL + path'); expect(doReqMatch![0]).toContain('_tcp_request('); @@ -67,24 +76,47 @@ describe('generatePreamble — Unix-vs-TCP transport gate', () => { * exported. The path must remain hardcoded so the preamble does * not depend on env-var injection. */ expect(preamble).toContain('_TOOL_CALL_SOCKET = "/tmp/tcs.sock"'); - expect(preamble).not.toMatch(/os\.environ\.get\(['"]TOOL_CALL_SOCKET['"]/); + expect(preamble).not.toMatch( + /os\.environ\.get\(['"]TOOL_CALL_SOCKET['"]/, + ); expect(preamble).not.toMatch(/os\.environ\[['"]TOOL_CALL_SOCKET['"]\]/); }); }); describe('extractPendingFromStdout — input hash metadata', () => { + test('normalizes native control payload hashes instead of trusting the sandbox', () => { + const forgedHash = hashToolInput({ resource: 'B' }); + const expectedHash = hashToolInput({ resource: 'A' }); + const pending = extractPendingFromControlPayload( + JSON.stringify({ + pending: [ + { + call_id: 'call_001', + tool_name: 'authorize', + input: { resource: 'A' }, + input_hash: forgedHash, + }, + ], + }), + ); + expect(pending?.[0]?.input_hash).toBe(expectedHash); + expect(pending?.[0]?.input_hash).not.toBe(forgedHash); + }); + test('ignores sandbox-supplied input_hash and uses the parsed input hash', () => { const executionId = 'exec_hash_guard'; const { start, end } = buildScopedSentinel(executionId); const forgedHash = hashToolInput({ resource: 'B' }); const expectedHash = hashToolInput({ resource: 'A' }); const payload = { - pending: [{ + pending: [ + { call_id: 'call_001', tool_name: 'authorize', input: { resource: 'A' }, input_hash: forgedHash, - }], + }, + ], }; const parsed = extractPendingFromStdout( diff --git a/service/src/preamble.ts b/service/src/preamble.ts index 9ff75dc7..94685c6c 100644 --- a/service/src/preamble.ts +++ b/service/src/preamble.ts @@ -2,7 +2,10 @@ import fs from 'fs'; import path from 'path'; import type * as t from './types'; import { planLimits } from './config'; -import { generateBashReplayPreamble, generateBashReplayPostamble } from './preamble-bash'; +import { + generateBashReplayPreamble, + generateBashReplayPostamble, +} from './preamble-bash'; import { PTC_HISTORY_FILENAME, PTC_HISTORY_SANDBOX_PATH, @@ -11,10 +14,16 @@ import { buildScopedSentinel, isReservedPtcFilename, } from './ptc-constants'; -import { hashToolInput, pendingInputHashesFromRawPayload } from './tool-input-signature'; +import { + hashToolInput, + pendingInputHashesFromRawPayload, +} from './tool-input-signature'; // Load async matplotlib template for programmatic tool calling -const templateCodeAsync = fs.readFileSync(path.join(__dirname, 'matplotlib-async.py'), 'utf8'); +const templateCodeAsync = fs.readFileSync( + path.join(__dirname, 'matplotlib-async.py'), + 'utf8', +); // ============================================================================= // Programmatic Tool Calling Types & Preamble Generation @@ -88,11 +97,41 @@ function normalizePythonFunctionName(name: string): string { // Python keywords to avoid const pythonKeywords = new Set([ - 'False', 'None', 'True', 'and', 'as', 'assert', 'async', 'await', - 'break', 'class', 'continue', 'def', 'del', 'elif', 'else', 'except', - 'finally', 'for', 'from', 'global', 'if', 'import', 'in', 'is', - 'lambda', 'nonlocal', 'not', 'or', 'pass', 'raise', 'return', - 'try', 'while', 'with', 'yield' + 'False', + 'None', + 'True', + 'and', + 'as', + 'assert', + 'async', + 'await', + 'break', + 'class', + 'continue', + 'def', + 'del', + 'elif', + 'else', + 'except', + 'finally', + 'for', + 'from', + 'global', + 'if', + 'import', + 'in', + 'is', + 'lambda', + 'nonlocal', + 'not', + 'or', + 'pass', + 'raise', + 'return', + 'try', + 'while', + 'with', + 'yield', ]); if (pythonKeywords.has(normalized)) { @@ -137,11 +176,14 @@ function jsonSchemaToPythonType(schema: JsonSchemaProperty): string { * Sort property names so required parameters come before optional ones. * Uses a Set for O(1) lookups instead of repeated array includes() calls. */ -function getSortedPropertyNames(propertyNames: string[], required: string[]): string[] { +function getSortedPropertyNames( + propertyNames: string[], + required: string[], +): string[] { const requiredSet = new Set(required); return [ ...propertyNames.filter(name => requiredSet.has(name)), - ...propertyNames.filter(name => !requiredSet.has(name)) + ...propertyNames.filter(name => !requiredSet.has(name)), ]; } @@ -155,7 +197,10 @@ function schemaToParams(schema?: JsonSchema): string { const required = schema.required ?? []; const requiredSet = new Set(required); - const sortedNames = getSortedPropertyNames(Object.keys(schema.properties), required); + const sortedNames = getSortedPropertyNames( + Object.keys(schema.properties), + required, + ); const params: string[] = []; @@ -192,7 +237,11 @@ function inferReturnType(description?: string): string { const desc = description.toLowerCase(); - if (desc.includes('returns list') || desc.includes('returns array') || desc.includes('list of')) { + if ( + desc.includes('returns list') || + desc.includes('returns array') || + desc.includes('list of') + ) { return 'List[Dict[str, Any]]'; } if (desc.includes('returns dict') || desc.includes('returns object')) { @@ -225,7 +274,10 @@ function generateDocstring(tool: LCTool): string { doc += '\n\n Parameters:'; const required = tool.parameters.required ?? []; const requiredSet = new Set(required); - const sortedNames = getSortedPropertyNames(Object.keys(tool.parameters.properties), required); + const sortedNames = getSortedPropertyNames( + Object.keys(tool.parameters.properties), + required, + ); for (const name of sortedNames) { const propSchema = tool.parameters.properties[name]; @@ -253,7 +305,8 @@ function generateToolStub(tool: LCTool): string { const pythonFunctionName = normalizePythonFunctionName(tool.name); // If name was changed, add a comment - const nameComment = pythonFunctionName !== tool.name + const nameComment = + pythonFunctionName !== tool.name ? ` # Original tool name: ${tool.name}\n` : ''; @@ -451,7 +504,8 @@ async def _execute_tool_internal_async(tool_name: str, tool_input: Dict[str, Any */ export function generateReplayPreamble(config: ReplayPreambleConfig): string { const { executionId, tools } = config; - const { start: scopedStart, end: scopedEnd } = buildScopedSentinel(executionId); + const { start: scopedStart, end: scopedEnd } = + buildScopedSentinel(executionId); let preamble = ` # ============================================================================ @@ -564,6 +618,63 @@ export interface ExtractPendingResult { }> | null; } +export function extractPendingFromControlPayload( + rawPayload: string, +): ExtractPendingResult['pending'] { + let parsed: { pending?: unknown } | null = null; + try { + parsed = JSON.parse(rawPayload) as { pending?: unknown }; + } catch { + return null; + } + + const pendingField = parsed?.pending; + if (!Array.isArray(pendingField)) return null; + + const rawInputHashes = pendingInputHashesFromRawPayload(rawPayload); + type PendingWithIndex = { + c: { call_id: string; tool_name: string; input: unknown }; + index: number; + }; + const isPendingWithIndex = (entry: { + c: unknown; + index: number; + }): entry is PendingWithIndex => { + const { c } = entry; + return ( + c != null && + typeof c === 'object' && + typeof (c as { call_id?: unknown }).call_id === 'string' && + typeof (c as { tool_name?: unknown }).tool_name === 'string' + ); + }; + return pendingField + .map((c, index) => ({ c, index })) + .filter(isPendingWithIndex) + .map(({ c, index }) => { + const callSite = (c as { call_site?: unknown }).call_site; + const rawInputHash = rawInputHashes[index]; + const hasObjectInput = + c.input != null && typeof c.input === 'object'; + const input = (hasObjectInput ? c.input : {}) as Record< + string, + unknown + >; + return { + call_id: c.call_id, + tool_name: c.tool_name, + input, + input_hash: + hasObjectInput && typeof rawInputHash === 'string' + ? rawInputHash + : hashToolInput(input), + ...(typeof callSite === 'string' + ? { call_site: callSite } + : {}), + }; + }); +} + /** * Locate the last line whose trimmed content exactly equals `marker`. * Using full-line anchoring prevents user-provided tool payloads that happen @@ -581,7 +692,8 @@ function findSentinelLine( for (let i = lines.length - 1; i >= searchFromLine; i--) { if (lines[i].trim() === marker) { const startOffset = lineStartOffsets[i]; - const endOffset = i + 1 < lineStartOffsets.length + const endOffset = + i + 1 < lineStartOffsets.length ? lineStartOffsets[i + 1] - 1 : startOffset + lines[i].length; return { line: i, startOffset, endOffset }; @@ -619,48 +731,8 @@ export function extractPendingFromStdout( const payloadLines = lines.slice(startLine.line + 1, endLine.line); const rawPayload = payloadLines.join('\n').trim(); - let parsed: { pending?: unknown } | null = null; - try { - parsed = JSON.parse(rawPayload) as { pending?: unknown }; - } catch { - return { stdout, pending: null }; - } - - const pendingField = parsed?.pending; - if (!Array.isArray(pendingField)) return { stdout, pending: null }; - - const rawInputHashes = pendingInputHashesFromRawPayload(rawPayload); - type PendingWithIndex = { - c: { call_id: string; tool_name: string; input: unknown }; - index: number; - }; - const isPendingWithIndex = (entry: { c: unknown; index: number }): entry is PendingWithIndex => { - const { c } = entry; - return ( - c != null && - typeof c === 'object' && - typeof (c as { call_id?: unknown }).call_id === 'string' && - typeof (c as { tool_name?: unknown }).tool_name === 'string' - ); - }; - const pending = pendingField - .map((c, index) => ({ c, index })) - .filter(isPendingWithIndex) - .map(({ c, index }) => { - const callSite = (c as { call_site?: unknown }).call_site; - const rawInputHash = rawInputHashes[index]; - const hasObjectInput = c.input != null && typeof c.input === 'object'; - const input = (hasObjectInput ? c.input : {}) as Record; - return { - call_id: c.call_id, - tool_name: c.tool_name, - input, - input_hash: hasObjectInput && typeof rawInputHash === 'string' - ? rawInputHash - : hashToolInput(input), - ...(typeof callSite === 'string' ? { call_site: callSite } : {}), - }; - }); + const pending = extractPendingFromControlPayload(rawPayload); + if (pending == null) return { stdout, pending: null }; /** Strip only the sentinel block and leave every other byte of user * stdout untouched. Both the Python and bash preambles defensively @@ -674,7 +746,10 @@ export function extractPendingFromStdout( * emission, and anything else that depends on byte-accurate stdout. */ const rawHead = stdout.slice(0, startLine.startOffset); const head = rawHead.endsWith('\n') ? rawHead.slice(0, -1) : rawHead; - const tailStart = endLine.endOffset < stdout.length ? endLine.endOffset + 1 : stdout.length; + const tailStart = + endLine.endOffset < stdout.length + ? endLine.endOffset + 1 + : stdout.length; const tail = stdout.slice(tailStart); const cleaned = head + tail; @@ -688,11 +763,14 @@ export function extractPendingFromStdout( function wrapUserCodeInAsync(userCode: string): string { const lines = userCode.split('\n'); - let wrapped = '# ============================================================================\n'; + let wrapped = + '# ============================================================================\n'; wrapped += '# USER CODE BEGINS BELOW\n'; - wrapped += '# ============================================================================\n\n'; + wrapped += + '# ============================================================================\n\n'; wrapped += 'async def __user_main__():\n'; - wrapped += ' """Auto-generated wrapper for user code to support top-level await"""\n'; + wrapped += + ' """Auto-generated wrapper for user code to support top-level await"""\n'; // Indent all user code for (const line of lines) { @@ -743,10 +821,21 @@ const PROGRAMMATIC_RUN_TIMEOUT = 300000; // 5 minutes wall time * Create a payload for programmatic tool calling execution * Combines the tool preamble with user code */ -export function createProgrammaticPayload(options: CreateProgrammaticPayloadOptions): t.PayloadBody { +export function createProgrammaticPayload( + options: CreateProgrammaticPayloadOptions, +): t.PayloadBody { const { - req, session_id, execution_id, callbackUrl, callbackToken, tools, timeout, - mode = 'blocking', history, codeOverride, filesOverride, + req, + session_id, + execution_id, + callbackUrl, + callbackToken, + tools, + timeout, + mode = 'blocking', + history, + codeOverride, + filesOverride, language = 'python', } = options; const body = req.body as t.ProgrammaticRequestBody; @@ -762,7 +851,14 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti throw new Error('bash PTC is only supported in replay mode'); } return buildBashPayload({ - req, execution_id, session_id, tools, userCode, files, history, timeout, + req, + execution_id, + session_id, + tools, + userCode, + files, + history, + timeout, }); } @@ -771,7 +867,9 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti preamble = generateReplayPreamble({ executionId: execution_id, tools }); } else { if (!callbackUrl || !callbackToken) { - throw new Error('blocking PTC mode requires callbackUrl and callbackToken'); + throw new Error( + 'blocking PTC mode requires callbackUrl and callbackToken', + ); } preamble = generatePreamble({ callbackUrl, @@ -781,15 +879,21 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti }); } - const isPyPlot = userCode.includes('import matplotlib') || userCode.includes('import seaborn'); + const isPyPlot = + userCode.includes('import matplotlib') || + userCode.includes('import seaborn'); let finalCode: string; if (isPyPlot) { - const indentedUserCode = userCode.trim().split('\n').map(line => ` ${line}`).join('\n'); + const indentedUserCode = userCode + .trim() + .split('\n') + .map(line => ` ${line}`) + .join('\n'); const wrappedUserCode = templateCodeAsync.replace( /# BEGIN USER CODE\n[\s\S]*?# END USER CODE/, - `# BEGIN USER CODE\n${indentedUserCode}\n # END USER CODE` + `# BEGIN USER CODE\n${indentedUserCode}\n # END USER CODE`, ); finalCode = preamble + '\n' + wrappedUserCode; } else { @@ -797,7 +901,9 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti finalCode = preamble + wrappedUserCode; } - const run_memory_limit = planLimits[req.planId ?? '']?.run_memory_limit ?? planLimits.default.run_memory_limit; + const run_memory_limit = + planLimits[req.planId ?? '']?.run_memory_limit ?? + planLimits.default.run_memory_limit; const run_timeout = timeout ?? PROGRAMMATIC_RUN_TIMEOUT; const payload: t.PayloadBody = { @@ -809,8 +915,8 @@ export function createProgrammaticPayload(options: CreateProgrammaticPayloadOpti files: [ { name: 'main.py', - content: finalCode - } + content: finalCode, + }, ], session_id, }; @@ -851,13 +957,27 @@ function buildBashPayload(args: { history?: Record; timeout?: number; }): t.PayloadBody { - const { req, execution_id, session_id, tools, userCode, files, history, timeout } = args; - - const preamble = generateBashReplayPreamble({ executionId: execution_id, tools }); + const { + req, + execution_id, + session_id, + tools, + userCode, + files, + history, + timeout, + } = args; + + const preamble = generateBashReplayPreamble({ + executionId: execution_id, + tools, + }); const postamble = generateBashReplayPostamble(); const finalCode = preamble + userCode + '\n' + postamble; - const run_memory_limit = planLimits[req.planId ?? '']?.run_memory_limit ?? planLimits.default.run_memory_limit; + const run_memory_limit = + planLimits[req.planId ?? '']?.run_memory_limit ?? + planLimits.default.run_memory_limit; const run_timeout = timeout ?? PROGRAMMATIC_RUN_TIMEOUT; const payload: t.PayloadBody = { @@ -865,6 +985,8 @@ function buildBashPayload(args: { run_timeout, language: 'bash', version: '5.2.0', + execution_id, + replay_tool_count: tools.length, files: [ { name: 'main.sh', diff --git a/service/src/programmatic-cancellation.test.ts b/service/src/programmatic-cancellation.test.ts new file mode 100644 index 00000000..53a4340e --- /dev/null +++ b/service/src/programmatic-cancellation.test.ts @@ -0,0 +1,234 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test'; +import type IORedis from 'ioredis'; +import { startTestRedis } from './test/redis'; +import { commitJobResult, requestJobCancellation } from './job-cancellation'; +import { + attachProgrammaticCancellationTarget, + cancelProgrammaticRequest, + normalizeProgrammaticRequestId, + programmaticCancellationInternals, + releaseProgrammaticCancellation, + reserveProgrammaticCancellation, +} from './programmatic-cancellation'; + +let redis: IORedis & { closeTestServer(): Promise }; + +beforeEach(async () => { + redis = await startTestRedis(); +}); + +afterEach(async () => { + await redis.closeTestServer(); +}); + +test('normalizes only bounded opaque request IDs', () => { + expect(normalizeProgrammaticRequestId('request_123456789')).toBe( + 'request_123456789', + ); + expect(normalizeProgrammaticRequestId(' short ')).toBeUndefined(); + expect( + normalizeProgrammaticRequestId('../request_123456789'), + ).toBeUndefined(); + expect(normalizeProgrammaticRequestId('a'.repeat(129))).toBeUndefined(); +}); + +test('Stop cannot extend an attached tombstone before the outcome command succeeds', async () => { + const requestId = 'request_no_split_renewal'; + const owner = 'owner-a'; + const target = { queueName: 'other', jobId: 'split-renewal' }; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }); + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner, + target, + ttlSeconds: 60, + }); + await commitJobResult(redis, target, { stdout: 'done' }, 60); + const key = programmaticCancellationInternals.requestKey(requestId); + await redis.pexpire(key, 5_000); + const before = await redis.pttl(key); + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 600, + }), + ).toEqual({ status: 'accepted', target }); + // Simulate losing Redis before requestJobCancellation: no second command. + expect(await redis.pttl(key)).toBeLessThanOrEqual(before); + expect(await requestJobCancellation(redis, target, 600)).toBe(false); +}); + +test('cancellation before queue attachment is retained atomically', async () => { + const requestId = 'request_early_cancel_123'; + const owner = 'owner-a'; + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toEqual({ status: 'accepted' }); + + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('cancelled'); + expect( + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner, + target: { queueName: 'other', jobId: '42' }, + ttlSeconds: 60, + }), + ).toBe('cancelled'); +}); + +test('cancellation after attachment returns the exact queue target', async () => { + const requestId = 'request_attached_cancel_1'; + const owner = 'owner-a'; + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('active'); + expect( + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner, + target: { queueName: 'other', jobId: '43' }, + ttlSeconds: 60, + }), + ).toBe('active'); + + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toEqual({ + status: 'accepted', + target: { queueName: 'other', jobId: '43' }, + }); +}); + +test('overlapping requests from the same owner cannot share cancellation state', async () => { + const requestId = 'request_duplicate_owner_1'; + const owner = 'owner-a'; + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('active'); + + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('duplicate'); +}); + +test('a different principal cannot reserve, attach, cancel, or release a request', async () => { + const requestId = 'request_owned_cancel_123'; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-b', + ttlSeconds: 60, + }), + ).toBe('forbidden'); + expect( + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner: 'owner-b', + target: { queueName: 'other', jobId: '44' }, + ttlSeconds: 60, + }), + ).toBe('forbidden'); + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner: 'owner-b', + ttlSeconds: 60, + }), + ).toEqual({ status: 'forbidden' }); + await releaseProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-b', + }); + expect( + await redis.exists(programmaticCancellationInternals.requestKey(requestId)), + ).toBe(1); +}); + +test('settlement retains a bounded target tombstone for late Stop classification', async () => { + const requestId = 'request_release_cancel_1'; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + const target = { queueName: 'other', jobId: 'settled-job' }; + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner: 'owner-a', + target, + ttlSeconds: 60, + }); + await commitJobResult(redis, target, { stdout: 'done' }, 60); + await releaseProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + }); + const key = programmaticCancellationInternals.requestKey(requestId); + expect(await redis.exists(key)).toBe(1); + expect(await redis.ttl(key)).toBeGreaterThan(0); + expect(await redis.ttl(key)).toBeLessThanOrEqual(60); + const cancelled = await cancelProgrammaticRequest({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + expect(cancelled).toEqual({ status: 'accepted', target }); + expect(await requestJobCancellation(redis, target, 60)).toBe(false); +}); diff --git a/service/src/programmatic-cancellation.ts b/service/src/programmatic-cancellation.ts new file mode 100644 index 00000000..12a6694d --- /dev/null +++ b/service/src/programmatic-cancellation.ts @@ -0,0 +1,182 @@ +import { createHash } from 'node:crypto'; +import type IORedis from 'ioredis'; +import type { AuthenticatedRequest } from './types'; +import { getCredentialId } from './auth/principal'; +import { getExecutionIdentity } from './execution-identity'; + +export const CODEAPI_PROGRAMMATIC_REQUEST_HEADER = + 'X-LibreChat-Code-Request-ID'; +const REQUEST_PREFIX = 'codeapi:programmatic-cancellation:v1'; +const REQUEST_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; + +interface CancellationTarget { + queueName: string; + jobId: string; +} + +export type CancellationRequestResult = + | { status: 'accepted'; target?: CancellationTarget } + | { status: 'forbidden' }; + +function requestKey(requestId: string): string { + return `${REQUEST_PREFIX}:${requestId}`; +} + +export function normalizeProgrammaticRequestId( + value: unknown, +): string | undefined { + if (typeof value !== 'string') return undefined; + const trimmed = value.trim(); + return REQUEST_ID_PATTERN.test(trimmed) ? trimmed : undefined; +} + +export function programmaticCancellationOwner( + req: AuthenticatedRequest, + userId: string, +): string { + const identity = getExecutionIdentity(req, userId); + return createHash('sha256') + .update( + JSON.stringify([ + identity.storageNamespace, + identity.canonicalUserId, + getCredentialId(req), + identity.authContextHash ?? '', + ]), + ) + .digest('hex'); +} + +const RESERVE_SCRIPT = ` +local key = KEYS[1] +local owner = ARGV[1] +local ttl = tonumber(ARGV[2]) +local existing = redis.call('HGET', key, 'owner') +if existing and existing ~= owner then return -1 end +if redis.call('HGET', key, 'reserved') == '1' then return -2 end +if not existing then + redis.call('HSET', key, 'owner', owner, 'cancelled', '0') +end +redis.call('HSET', key, 'reserved', '1') +redis.call('EXPIRE', key, ttl) +return tonumber(redis.call('HGET', key, 'cancelled') or '0') +`; + +const ATTACH_SCRIPT = ` +local key = KEYS[1] +local owner = ARGV[1] +local queueName = ARGV[2] +local jobId = ARGV[3] +local ttl = tonumber(ARGV[4]) +if redis.call('HGET', key, 'owner') ~= owner then return -1 end +redis.call('HSET', key, 'queueName', queueName, 'jobId', jobId) +redis.call('EXPIRE', key, ttl) +return tonumber(redis.call('HGET', key, 'cancelled') or '0') +`; + +const CANCEL_SCRIPT = ` +local key = KEYS[1] +local owner = ARGV[1] +local ttl = tonumber(ARGV[2]) +local existing = redis.call('HGET', key, 'owner') +if existing and existing ~= owner then return {-1} end +if not existing then redis.call('HSET', key, 'owner', owner) end +redis.call('HSET', key, 'cancelled', '1') +local queueName = redis.call('HGET', key, 'queueName') +local jobId = redis.call('HGET', key, 'jobId') +-- Once attached, never extend this mapping independently of the job decision. +-- Its original admission TTL already covers execution and late cancellation. +if queueName and jobId then return {1, queueName, jobId} end +redis.call('EXPIRE', key, ttl) +return {1} +`; + +const RELEASE_SCRIPT = ` +if redis.call('HGET', KEYS[1], 'owner') == ARGV[1] then + -- Keep the owner/target tombstone through its existing bounded TTL. A Stop + -- racing response delivery must still reach the job's completion decision. + return redis.call('HSET', KEYS[1], 'finished', '1') +end +return 0 +`; + +export async function reserveProgrammaticCancellation(args: { + redis: IORedis; + requestId: string; + owner: string; + ttlSeconds: number; +}): Promise<'active' | 'cancelled' | 'duplicate' | 'forbidden'> { + const result = Number( + await args.redis.eval( + RESERVE_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + Math.max(1, args.ttlSeconds), + ), + ); + if (result === -1) return 'forbidden'; + if (result === -2) return 'duplicate'; + return result === 1 ? 'cancelled' : 'active'; +} + +export async function attachProgrammaticCancellationTarget(args: { + redis: IORedis; + requestId: string; + owner: string; + target: CancellationTarget; + ttlSeconds: number; +}): Promise<'active' | 'cancelled' | 'forbidden'> { + const result = Number( + await args.redis.eval( + ATTACH_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + args.target.queueName, + args.target.jobId, + Math.max(1, args.ttlSeconds), + ), + ); + if (result < 0) return 'forbidden'; + return result === 1 ? 'cancelled' : 'active'; +} + +export async function cancelProgrammaticRequest(args: { + redis: IORedis; + requestId: string; + owner: string; + ttlSeconds: number; +}): Promise { + const raw = await args.redis.eval( + CANCEL_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + Math.max(1, args.ttlSeconds), + ); + const result = Array.isArray(raw) ? raw.map(String) : []; + if (result[0] === '-1') return { status: 'forbidden' }; + if (result.length >= 3) { + return { + status: 'accepted', + target: { queueName: result[1]!, jobId: result[2]! }, + }; + } + return { status: 'accepted' }; +} + +export async function releaseProgrammaticCancellation(args: { + redis: IORedis; + requestId: string; + owner: string; +}): Promise { + await args.redis.eval( + RELEASE_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + ); +} + +export const programmaticCancellationInternals = { requestKey }; diff --git a/service/src/ptc-constants.test.ts b/service/src/ptc-constants.test.ts new file mode 100644 index 00000000..5028e385 --- /dev/null +++ b/service/src/ptc-constants.test.ts @@ -0,0 +1,9 @@ +import { expect, test } from 'bun:test'; +import { isReservedPtcFilename } from './ptc-constants'; + +test('reserves replay inputs and output control channels after normalization', () => { + for (const name of ['_ptc_history.json', '_ptc_pending_result.json', '_PTC_PENDING_RESULT.JSON', 'sub/../_ptc_pending_result.json', 'sub\\_ptc_pending_result.json']) { + expect(isReservedPtcFilename(name)).toBe(true); + } + expect(isReservedPtcFilename('_ptc_data.csv')).toBe(false); +}); diff --git a/service/src/ptc-constants.ts b/service/src/ptc-constants.ts index 5009c3c3..b7b3999b 100644 --- a/service/src/ptc-constants.ts +++ b/service/src/ptc-constants.ts @@ -15,8 +15,9 @@ export const PTC_HISTORY_SANDBOX_PATH = `/mnt/data/${PTC_HISTORY_FILENAME}`; * Returns `true` for any filename the submission layer must refuse. * * Two things make a name "reserved": - * 1. Its post-normalization basename is `_ptc_history.json` — the single - * runtime fixture the replay preamble injects into the submission dir. + * 1. Its post-normalization basename is `_ptc_history.json` or + * `_ptc_pending_result.json`, compared case-insensitively for macOS. + * These are the replay input and output control channels. * Any user-supplied file with that exact basename would shadow our * injected history and silently corrupt replay correctness, so we * reject it on the request path. The bash preamble's `_ptc_pending.*` @@ -61,7 +62,7 @@ export function isReservedPtcFilename(name: string): boolean { } if (escapes) return true; const basename = segments.length > 0 ? segments[segments.length - 1] : ''; - return basename === PTC_HISTORY_FILENAME; + return [PTC_HISTORY_FILENAME, '_ptc_pending_result.json'].includes(basename.toLowerCase()); } /** diff --git a/service/src/queue.ts b/service/src/queue.ts index 54fea308..72c2fbb1 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -1,6 +1,7 @@ // src/queue.ts import IORedis from 'ioredis'; import { Queue, QueueEvents } from 'bullmq'; +import type { Job } from 'bullmq'; import { setMaxListeners } from 'events'; import type { CommonRedisOptions } from 'ioredis'; import type * as tls from 'tls'; @@ -17,19 +18,13 @@ import type { SandboxBackendName, } from './execution-profile'; import logger from './logger'; -import { redisKeepAliveOptions } from './redis-options'; +import { redisKeepAliveOptions, redisReconnectDelay } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; - -const MAX_RECONNECT_ATTEMPTS = 5; -const RECONNECT_DELAY = 2000; +import { JobCancellationRegistry } from './job-cancellation'; const retryStrategy: CommonRedisOptions['retryStrategy'] = (times) => { - if (times > MAX_RECONNECT_ATTEMPTS) { - logger.error(`Failed to connect to Redis after ${times} attempts`); - return null; - } logger.warn(`Retrying Redis connection attempt ${times}`); - return RECONNECT_DELAY; + return redisReconnectDelay(times); }; const reconnectOnError: CommonRedisOptions['reconnectOnError'] = (err) => { @@ -60,6 +55,7 @@ const connection = new IORedis({ ? { dnsLookup: (address: string, callback: (err: Error | null, addr: string) => void): void => callback(null, address) } : {}) }); +const jobCancellationRegistry = new JobCancellationRegistry(connection); // Global queues - no INSTANCE_ID prefix // This enables horizontal scaling where any worker can process any job @@ -110,6 +106,19 @@ export function getExecutionQueueBinding( return { ...getQueueResources(name), language }; } +/** + * Resolve a job only from this deployment's already-open queue set. Every + * homogeneous API replica opens both execution queues at startup, so this + * supports cross-replica cancellation without allocating attacker-shaped + * QueueEvents connections for arbitrary names recovered from Redis. + */ +export async function getExistingExecutionJob( + queueName: string, + jobId: string, +): Promise | undefined> { + return queueResources.get(queueName)?.queue.getJob(jobId); +} + const { queue: pyQueue, events: pyQueueEvents } = getQueueResources(queueNames.python); const { queue: otherQueue, events: otherQueueEvents } = getQueueResources(queueNames.other); @@ -163,8 +172,16 @@ export async function closeQueueConnections(): Promise { [...queueResources.values()].flatMap(({ queue, events }) => [ queue.close(), events.close(), - ]), + ]).concat(jobCancellationRegistry.close()), ); } -export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, queueNames, connection }; +export { + pyQueue, + otherQueue, + pyQueueEvents, + otherQueueEvents, + queueNames, + connection, + jobCancellationRegistry, +}; diff --git a/service/src/redis-options.test.ts b/service/src/redis-options.test.ts index 1bb942e6..9795e279 100644 --- a/service/src/redis-options.test.ts +++ b/service/src/redis-options.test.ts @@ -1,5 +1,15 @@ import { afterEach, describe, expect, test } from 'bun:test'; -import { redisKeepAliveMs, redisKeepAliveOptions } from './redis-options'; +import { + redisKeepAliveMs, + redisKeepAliveOptions, + redisReconnectDelay, +} from './redis-options'; + +test('long-lived command connections keep recovering with a bounded retry delay', () => { + expect(redisReconnectDelay(1)).toBe(100); + expect(redisReconnectDelay(6)).toBe(600); + expect(redisReconnectDelay(1_000)).toBe(2_000); +}); describe('Redis keepalive options', () => { afterEach(() => { diff --git a/service/src/redis-options.ts b/service/src/redis-options.ts index 98d3f5b1..1f2c8e42 100644 --- a/service/src/redis-options.ts +++ b/service/src/redis-options.ts @@ -1,5 +1,11 @@ import type { CommonRedisOptions } from 'ioredis'; +/** Long-lived queue/cancellation command and subscriber connections must both + * recover after an outage. Never leave a live process with a terminal client. */ +export function redisReconnectDelay(attempt: number): number { + return Math.min(2_000, 100 * Math.max(1, attempt)); +} + export function redisKeepAliveMs(): number { const raw = process.env.REDIS_KEEP_ALIVE_MS; const trimmed = raw?.trim(); diff --git a/service/src/request-disconnect.test.ts b/service/src/request-disconnect.test.ts new file mode 100644 index 00000000..15660aca --- /dev/null +++ b/service/src/request-disconnect.test.ts @@ -0,0 +1,60 @@ +import { expect, test } from 'bun:test'; +import { EventEmitter } from 'node:events'; +import type { Response } from 'express'; +import type { AuthenticatedRequest } from './types'; +import { CLIENT_DISCONNECT_REASON } from './job-cancellation'; +import { observeRequestDisconnect } from './request-disconnect'; + +function requestAndResponse(options: { + requestAborted?: boolean; + requestDestroyed?: boolean; + responseDestroyed?: boolean; +} = {}): { + req: AuthenticatedRequest & EventEmitter; + res: Response & EventEmitter; +} { + const req = Object.assign(new EventEmitter(), { + aborted: options.requestAborted ?? false, + destroyed: options.requestDestroyed ?? false, + }) as AuthenticatedRequest & EventEmitter; + const res = Object.assign(new EventEmitter(), { + destroyed: options.responseDestroyed ?? false, + writableFinished: false, + }) as Response & EventEmitter; + return { req, res }; +} + +test('a consumed Bun request stream is not mistaken for a disconnect', () => { + const { req, res } = requestAndResponse({ requestDestroyed: true }); + const observer = observeRequestDisconnect(req, res); + + expect(observer.isDisconnected()).toBe(false); + expect(observer.signal.aborted).toBe(false); + observer.dispose(); +}); + +test('current and future transport abandonment abort exactly once', () => { + const current = requestAndResponse({ requestAborted: true }); + const currentObserver = observeRequestDisconnect(current.req, current.res); + expect(currentObserver.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + + const future = requestAndResponse(); + const futureObserver = observeRequestDisconnect(future.req, future.res); + future.res.emit('close'); + future.req.emit('aborted'); + expect(futureObserver.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + expect(future.req.listenerCount('aborted')).toBe(0); + expect(future.res.listenerCount('close')).toBe(0); +}); + +test('a completed response disposes listeners without aborting', () => { + const { req, res } = requestAndResponse(); + const observer = observeRequestDisconnect(req, res); + (res as unknown as { writableFinished: boolean }).writableFinished = true; + res.emit('finish'); + res.emit('close'); + + expect(observer.isDisconnected()).toBe(false); + expect(req.listenerCount('aborted')).toBe(0); + expect(res.listenerCount('close')).toBe(0); +}); diff --git a/service/src/request-disconnect.ts b/service/src/request-disconnect.ts new file mode 100644 index 00000000..2a4d7a9d --- /dev/null +++ b/service/src/request-disconnect.ts @@ -0,0 +1,49 @@ +import type { Response } from 'express'; +import type { AuthenticatedRequest } from './types'; +import { CLIENT_DISCONNECT_REASON } from './job-cancellation'; + +export interface RequestDisconnectObserver { + signal: AbortSignal; + isDisconnected(): boolean; + dispose(): void; +} + +/** + * Observe a genuinely abandoned HTTP response across Node and Bun. + * + * Bun may mark the consumed IncomingMessage stream as `destroyed` while the + * response remains healthy, so request stream destruction is deliberately not + * treated as a disconnect. Express' `aborted` event and ServerResponse's + * pre-finish `close` event are the portable abandonment signals. + */ +export function observeRequestDisconnect( + req: AuthenticatedRequest, + res: Response, +): RequestDisconnectObserver { + const controller = new AbortController(); + let disposed = false; + const dispose = (): void => { + if (disposed) return; + disposed = true; + req.removeListener('aborted', disconnect); + res.removeListener('close', disconnect); + res.removeListener('finish', dispose); + }; + const disconnect = (): void => { + if (!res.writableFinished && !controller.signal.aborted) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + dispose(); + }; + + req.once('aborted', disconnect); + res.once('close', disconnect); + res.once('finish', dispose); + if (req.aborted || res.destroyed) disconnect(); + + return { + signal: controller.signal, + isDisconnected: () => controller.signal.aborted, + dispose, + }; +} diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts index a65bc89e..697271ee 100644 --- a/service/src/sandbox-backend/remote-bridge.test.ts +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -62,6 +62,34 @@ describe('RemoteBridgeSandboxBackend', () => { }); }); + test('preserves an authenticated selected workspace on remote dispatch', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { session_id: 'session-1', language: 'bash', version: '5.2', files: [] }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend(store, 'default-vm'); + + await backend.execute(request(), { ...context(), workspaceId: 'project-a' }); + + expect(dispatched).toMatchObject({ + workerId: 'user-vm', + workspaceId: 'project-a', + requireTenantBinding: true, + }); + }); + test('maps tenant authorization rejection to a bridge backend error', async () => { const store = { dispatch: async (): ReturnType => { diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index 0bee0038..6e06eda8 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -60,6 +60,7 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { (this.dynamicWorkers || ctx.bridgeWorkerId !== this.workerId), body: req.body, headers: req.headers, + ...(ctx.workspaceId != null ? { workspaceId: ctx.workspaceId } : {}), runtimeSessionId: ctx.runtimeSessionId, deadlineAtMs: ctx.deadlineAtMs ?? Date.now() + env.JOB_TIMEOUT, signal: ctx.signal, diff --git a/service/src/sandbox-backend/types.ts b/service/src/sandbox-backend/types.ts index fbaa2d20..e21982d9 100644 --- a/service/src/sandbox-backend/types.ts +++ b/service/src/sandbox-backend/types.ts @@ -39,6 +39,8 @@ export interface SandboxExecuteContext { canonicalUserId?: string; /** Trusted API-selected outbound worker. Presence requires a tenant-bound credential. */ bridgeWorkerId?: string; + /** Trusted selected workspace for native replay-mode PTC. */ + workspaceId?: string; /** Stable identifier for this queued iteration, used to derive an idempotent * stateless launch token. PTC replay reuses one executionId across every * iteration, so the executionId alone cannot separate them; the request body @@ -63,6 +65,7 @@ export type SandboxRawResponse = t.ExecuteResponse & { session_id: string; files?: t.FileRefs; run?: t.ExecuteResponse['run']; + pending_tool_calls_payload?: string; }; export interface SandboxBackend { diff --git a/service/src/sandbox-dispatch.test.ts b/service/src/sandbox-dispatch.test.ts index 435be9aa..5b1ef749 100644 --- a/service/src/sandbox-dispatch.test.ts +++ b/service/src/sandbox-dispatch.test.ts @@ -12,8 +12,10 @@ import { } from './execution-manifest'; const SECRET = 'test-secret'; -const PRIVATE_KEY = 'MC4CAQAwBQYDK2VwBCIEIBoxzSJjQ5jTVyuohHtlD+uDGqv/tZ6hQS2CmxuOg2Wn'; -const PUBLIC_KEY = 'MCowBQYDK2VwAyEAeY3PRoTS3adfU6E3gQUB5hSZdrdMSw6OrKkH4UhYh0U='; +const PRIVATE_KEY = + 'MC4CAQAwBQYDK2VwBCIEIBoxzSJjQ5jTVyuohHtlD+uDGqv/tZ6hQS2CmxuOg2Wn'; +const PUBLIC_KEY = + 'MCowBQYDK2VwAyEAeY3PRoTS3adfU6E3gQUB5hSZdrdMSw6OrKkH4UhYh0U='; function payload(overrides: Partial = {}): t.PayloadBody { return { @@ -25,14 +27,22 @@ function payload(overrides: Partial = {}): t.PayloadBody { }; } -function claims(overrides: Partial = {}): ExecutionManifestClaims { +function claims( + overrides: Partial = {}, +): ExecutionManifestClaims { return { v: EXECUTION_MANIFEST_VERSION, exec_id: 'exec_123', tenant_id: 'tenant_abc', user_id: 'user_123', session_key: 'tenant:tenant_abc:user:user_123', - input_files: [{ id: 'file_123', session_id: 'sess_input', name: 'inputs/data.csv' }], + input_files: [ + { + id: 'file_123', + session_id: 'sess_input', + name: 'inputs/data.csv', + }, + ], read_sessions: ['sess_input'], output_session_id: 'sess_output', max_upload_bytes: 1024, @@ -46,6 +56,20 @@ function claims(overrides: Partial = {}): ExecutionMani } describe('sandbox execute request dispatch', () => { + test('budgets every input and output batch before signing the request', () => { + const request = buildSandboxExecuteRequest({ + payload: payload({ files: Array.from({ length: 9 }, (_, index) => ({ name: `${index}.txt`, id: `file_${index}`, storage_session_id: 'input' })) }), + programmaticTransferReserveMs: 60_000, + executionManifestClaims: claims({ max_output_files: 10 }), + executionManifestSecret: SECRET, + executionManifestTtlSeconds: 300, + nowSeconds: 1_000, + }); + // Three download batches plus three upload batches share one reserve. + expect(request.body.transfer_timeout_ms).toBe(10_000); + const verified = verifyExecutionManifest(request.body.execution_manifest!, SECRET, { nowSeconds: 1_000 }); + expect(verified.execute_body_sha256).toBe(executionManifestBodySha256(request.body)); + }); test('keeps large egress grants out of HTTP headers', () => { const largeGrant = `ceg1.${'a'.repeat(24_000)}`; const request = buildSandboxExecuteRequest({ @@ -64,18 +88,27 @@ describe('sandbox execute request dispatch', () => { const request = buildSandboxExecuteRequest({ payload: payload(), executionManifestClaims: claims(), + maxOutputFileBytes: 1_000, executionManifestSecret: SECRET, executionManifestTtlSeconds: 300, nowSeconds: 1_000, }); expect(request.headers[EXECUTION_MANIFEST_HEADER]).toBeUndefined(); + expect(request.body.max_output_files).toBe(10); + expect(request.body.max_output_file_bytes).toBe(1_000); expect(request.body.execution_manifest).toEqual(expect.any(String)); - expect(verifyExecutionManifest(request.body.execution_manifest!, SECRET, { nowSeconds: 1_100 })).toEqual(claims({ + expect( + verifyExecutionManifest(request.body.execution_manifest!, SECRET, { + nowSeconds: 1_100, + }), + ).toEqual( + claims({ execute_body_sha256: executionManifestBodySha256(request.body), iat: 1_000, exp: 1_300, - })); + }), + ); }); test('signs execution manifests with a private key when configured', () => { @@ -88,11 +121,19 @@ describe('sandbox execute request dispatch', () => { nowSeconds: 1_000, }); - expect(verifyExecutionManifestWithPublicKey(request.body.execution_manifest!, PUBLIC_KEY, { nowSeconds: 1_100 })).toEqual(claims({ + expect( + verifyExecutionManifestWithPublicKey( + request.body.execution_manifest!, + PUBLIC_KEY, + { nowSeconds: 1_100 }, + ), + ).toEqual( + claims({ execute_body_sha256: executionManifestBodySha256(request.body), iat: 1_000, exp: 1_300, - })); + }), + ); }); test('binds body-carried egress grants into signed execution manifests', () => { @@ -106,10 +147,16 @@ describe('sandbox execute request dispatch', () => { }); expect(request.body.egress_grant).toBe('ceg1.sealed-grant'); - expect(verifyExecutionManifest(request.body.execution_manifest!, SECRET, { nowSeconds: 1_100 })).toEqual(claims({ + expect( + verifyExecutionManifest(request.body.execution_manifest!, SECRET, { + nowSeconds: 1_100, + }), + ).toEqual( + claims({ execute_body_sha256: executionManifestBodySha256(request.body), iat: 1_000, exp: 1_300, - })); + }), + ); }); }); diff --git a/service/src/sandbox-dispatch.ts b/service/src/sandbox-dispatch.ts index e3066905..340834b5 100644 --- a/service/src/sandbox-dispatch.ts +++ b/service/src/sandbox-dispatch.ts @@ -1,5 +1,14 @@ import type * as t from './types'; -import { executionManifestBodySha256, signExecutionManifestWithKey, type ExecutionManifestClaims } from './execution-manifest'; +import { + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES, + BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY, + BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS, +} from '../../packages/code/src/protocol'; +import { + executionManifestBodySha256, + signExecutionManifestWithKey, + type ExecutionManifestClaims, +} from './execution-manifest'; interface BuildSandboxExecuteRequestArgs { payload: t.PayloadBody; @@ -9,6 +18,8 @@ interface BuildSandboxExecuteRequestArgs { executionManifestSecret: string; executionManifestTtlSeconds: number; nowSeconds?: number; + maxOutputFileBytes?: number; + programmaticTransferReserveMs?: number; } interface SandboxExecuteRequest { @@ -21,15 +32,31 @@ interface SandboxExecuteRequest { * ride in the JSON body instead of HTTP headers. Otherwise skill-heavy jobs can * fail with 431 before sandbox-runner reaches capability validation. */ -export function buildSandboxExecuteRequest(args: BuildSandboxExecuteRequestArgs): SandboxExecuteRequest { +export function buildSandboxExecuteRequest( + args: BuildSandboxExecuteRequestArgs, +): SandboxExecuteRequest { const body: t.PayloadBody = { ...args.payload }; - const headers: Record = { 'Content-Type': 'application/json' }; + const headers: Record = { + 'Content-Type': 'application/json', + }; if (args.egressGrantToken) { body.egress_grant = args.egressGrantToken; } + if (args.maxOutputFileBytes != null) { + body.max_output_file_bytes = args.maxOutputFileBytes; + } + if (args.programmaticTransferReserveMs != null) { + const batches = Math.max(1, + Math.ceil(body.files.filter(file => 'id' in file).length / BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY) + + Math.ceil((args.executionManifestClaims?.max_output_files ?? BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILES) / BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_CONCURRENCY), + ); + body.transfer_timeout_ms = Math.min(BRIDGE_WORKSPACE_PROGRAMMATIC_TRANSFER_TIMEOUT_MS, + Math.max(1, Math.floor(args.programmaticTransferReserveMs / batches))); + } if (args.executionManifestClaims) { + body.max_output_files = args.executionManifestClaims.max_output_files; const nowSeconds = args.nowSeconds ?? Math.floor(Date.now() / 1000); body.execution_manifest = signExecutionManifestWithKey( { diff --git a/service/src/sandbox-egress.ts b/service/src/sandbox-egress.ts index f6549c09..e47806a2 100644 --- a/service/src/sandbox-egress.ts +++ b/service/src/sandbox-egress.ts @@ -3,6 +3,7 @@ import { env } from './config'; import { createGatewayPtcCallbackToken } from './egress-gateway-client'; import type { ExecutionManifestClaims } from './execution-manifest'; import type * as t from './types'; +import { programmaticTransferReserveMs } from '../../packages/code/src/protocol'; export type SandboxJobSecurity = { payload: t.PayloadBody; @@ -65,6 +66,9 @@ export function timeoutMsToGrantSeconds(timeoutMs: number): number { } const DEFAULT_PROGRAMMATIC_TIMEOUT_MS = 300000; +const SELECTED_WORKSPACE_REPLAY_PASSES = 2; +const SELECTED_WORKSPACE_SETTLEMENT_RESERVE_MS = 5_000; +const SELECTED_WORKSPACE_MAX_QUEUE_RESERVE_MS = 30_000; export function normalizeProgrammaticTimeoutMs( rawTimeout: unknown, @@ -80,6 +84,31 @@ export function normalizeProgrammaticTimeoutMs( return Math.min(Math.ceil(rawTimeout), maxTimeout); } +/** + * Selected-workspace Bash replay may run one read-only probe and one commit + * pass in its final iteration. Bound each pass so both plus settlement reserve + * fit inside the worker-owned JOB_TIMEOUT instead of advertising a duration + * the assignment cannot complete. + */ +export function normalizeSelectedWorkspaceProgrammaticTimeoutMs( + rawTimeout: unknown, + jobTimeoutMs = env.JOB_TIMEOUT, +): number { + const totalBudget = Math.max(1, Math.floor(jobTimeoutMs)); + const queueReserve = Math.min( + SELECTED_WORKSPACE_MAX_QUEUE_RESERVE_MS, + Math.floor(totalBudget / 5), + ); + const executionBudget = Math.max( + 1, + totalBudget - queueReserve - SELECTED_WORKSPACE_SETTLEMENT_RESERVE_MS - programmaticTransferReserveMs(totalBudget), + ); + return normalizeProgrammaticTimeoutMs( + rawTimeout, + Math.max(1, Math.floor(executionBudget / SELECTED_WORKSPACE_REPLAY_PASSES)), + ); +} + export async function sealPtcCallbackTokenForGateway(args: { executionId: string; sessionId: string; diff --git a/service/src/service/blocking-poll.test.ts b/service/src/service/blocking-poll.test.ts index 11399f24..96903d1f 100644 --- a/service/src/service/blocking-poll.test.ts +++ b/service/src/service/blocking-poll.test.ts @@ -4,9 +4,13 @@ import type * as t from '../types'; const result: t.ExecuteResult = { session_id: 'session', stdout: 'successful code', stderr: '', files: [], + deleted_files: ['removed.txt'], artifact_delivery: { code: 'artifact_delivery_failed', status: 'failed', attempted: 1, delivered: 0, failed: 1, }, + artifact_truncation: { + code: 'artifact_truncated', reasons: { size: 1 }, skipped: ['large.csv'], skipped_count: 1, + }, }; function fixture(): BlockingPollDependencies { @@ -26,7 +30,9 @@ describe('blocking worker settlement', () => { const deps = fixture(); expect(await pollBlockingExecution('exec', 5, deps)).toEqual({ status: 'completed', stdout: result.stdout, stderr: '', files: [], + deleted_files: result.deleted_files, artifact_delivery: result.artifact_delivery, + artifact_truncation: result.artifact_truncation, }); expect(deps.now()).toBe(2); }); diff --git a/service/src/service/blocking-poll.ts b/service/src/service/blocking-poll.ts index 8ce07fa8..1505f818 100644 --- a/service/src/service/blocking-poll.ts +++ b/service/src/service/blocking-poll.ts @@ -33,7 +33,9 @@ export async function pollBlockingExecution( stdout?: string; stderr?: string; files?: t.FileRefs; + deleted_files?: string[]; artifact_delivery?: t.ArtifactDeliveryFailure; + artifact_truncation?: t.ArtifactTruncation; }> { const start = deps.now(); while (deps.now() - start < timeout) { @@ -47,7 +49,9 @@ export async function pollBlockingExecution( stdout: result.stdout, stderr: result.stderr, files: result.files, + deleted_files: result.deleted_files, artifact_delivery: result.artifact_delivery, + artifact_truncation: result.artifact_truncation, }; } } diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 9896518c..1063fbfe 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -4,19 +4,45 @@ import { Router } from 'express'; import type { Response } from 'express'; import type * as t from '../types'; import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle'; -import { executionLimiter } from '../middleware/limits'; +import { cancellationLimiter, executionLimiter } from '../middleware/limits'; import { pyQueue, pyQueueEvents, connection, + jobCancellationRegistry, getExecutionQueueBinding, + getExistingExecutionJob, } from '../queue'; -import { createProgrammaticPayload, extractPendingFromStdout } from '../preamble'; +import { + JOB_CANCELLED_MESSAGE, + programmaticCancellationError, + removeJobIfWaiting, + requestJobCancellation, + fenceJobCancellation, + waitForJobWithCancellation, +} from '../job-cancellation'; +import { + CODEAPI_PROGRAMMATIC_REQUEST_HEADER, + attachProgrammaticCancellationTarget, + cancelProgrammaticRequest, + normalizeProgrammaticRequestId, + programmaticCancellationOwner, + releaseProgrammaticCancellation, + reserveProgrammaticCancellation, +} from '../programmatic-cancellation'; +import { + createProgrammaticPayload, + extractPendingFromControlPayload, + extractPendingFromStdout, +} from '../preamble'; import { findBashToolNameCollision } from '../preamble-bash'; import type { LCTool } from '../preamble'; import { isReservedPtcFilename } from '../ptc-constants'; import { internalServiceHeaders } from '../internal-service-auth'; -import { resolveOutputBucketSessionKey, SessionKeyResolutionError } from '../session-key'; +import { + resolveOutputBucketSessionKey, + SessionKeyResolutionError, +} from '../session-key'; import { getCredentialId, getPrincipalOrReject } from '../auth/principal'; import { getExecutionIdentity } from '../execution-identity'; import { PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION } from '../runtime-session/job-policy'; @@ -31,18 +57,29 @@ import { Jobs } from '../enum'; import { env, jobCompletionWaitTimeoutMs } from '../config'; import { resolveQueuedSandboxBackend } from '../execution-profile'; import { publicExecutionFailure } from '../utils'; +import { observeRequestDisconnect } from '../request-disconnect'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, + normalizeSelectedWorkspaceProgrammaticTimeoutMs, prepareSandboxJobSecurity, sealPtcCallbackTokenForGateway, timeoutMsToGrantSeconds, } from '../sandbox-egress'; import { findUnregisteredToolCall } from '../tool-scope'; import { summarizeRequestedFiles } from '../execution-log'; -import { pollBlockingExecution, type BlockingPendingState } from './blocking-poll'; -import { clearSessionOwnership, recordSessionOwnership } from '../session-ownership'; -import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; +import { + pollBlockingExecution, + type BlockingPendingState, +} from './blocking-poll'; +import { + clearSessionOwnership, + recordSessionOwnership, +} from '../session-ownership'; +import { + FileRefAuthorizationError, + authorizeRequestedFiles, +} from './file-authorization'; import { buildReplayExecutionState, resolveReplayStateSandboxBackend, @@ -50,8 +87,10 @@ import { import { BridgeWorkerSelectionError, CODEAPI_BRIDGE_WORKER_HEADER, + CODEAPI_BRIDGE_WORKSPACE_HEADER, resolveBridgeWorkerSelection, } from '../bridge/selection'; +import { isValidBridgeWorkerId, BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES } from '../../../packages/code/src/protocol'; import logger from '../logger'; import { type ExecutionState, @@ -88,6 +127,18 @@ const JOB_COMPLETION_WAIT_TIMEOUT_MS = jobCompletionWaitTimeoutMs( env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS, ); +const PROGRAMMATIC_CANCELLATION_TTL_SECONDS = + Math.ceil(JOB_COMPLETION_WAIT_TIMEOUT_MS / 1000) + 60; + +interface ReplayRequestCancellation { + signal: AbortSignal; + isDisconnected(): boolean; + request?: { + requestId: string; + owner: string; + cancelledBeforeStart: boolean; + }; +} const router = Router(); @@ -145,26 +196,41 @@ async function retryToolCallServerRequest( ): Promise { let lastError: Error | undefined; - for (let attempt = 1; attempt <= TOOL_CALL_SERVER_RETRY_ATTEMPTS; attempt++) { + for ( + let attempt = 1; + attempt <= TOOL_CALL_SERVER_RETRY_ATTEMPTS; + attempt++ + ) { try { return await requestFn(); } catch (error) { lastError = error as Error; if (axios.isAxiosError(error)) { - if (error.response && error.response.status >= 400 && error.response.status < 500) { + if ( + error.response && + error.response.status >= 400 && + error.response.status < 500 + ) { throw error; } } if (attempt < TOOL_CALL_SERVER_RETRY_ATTEMPTS) { - logger.warn(`${context} failed (attempt ${attempt}/${TOOL_CALL_SERVER_RETRY_ATTEMPTS}), retrying...`, { + logger.warn( + `${context} failed (attempt ${attempt}/${TOOL_CALL_SERVER_RETRY_ATTEMPTS}), retrying...`, + { error: lastError.message, - }); - await new Promise(resolve => setTimeout(resolve, TOOL_CALL_SERVER_RETRY_DELAY * attempt)); + }, + ); + await new Promise(resolve => + setTimeout(resolve, TOOL_CALL_SERVER_RETRY_DELAY * attempt), + ); } } } - logger.error(`${context} failed after ${TOOL_CALL_SERVER_RETRY_ATTEMPTS} attempts`); + logger.error( + `${context} failed after ${TOOL_CALL_SERVER_RETRY_ATTEMPTS} attempts`, + ); throw lastError; } @@ -179,7 +245,9 @@ setInterval(() => { }, STALE_CLEANUP_INTERVAL_MS); function generateContinuationToken(execution_id: string): string { - return Buffer.from(JSON.stringify({ execution_id, ts: Date.now() })).toString('base64'); + return Buffer.from( + JSON.stringify({ execution_id, ts: Date.now() }), + ).toString('base64'); } /** Map a replay-continuation HTTP status to its operational outcome @@ -200,14 +268,21 @@ function classifyContinuationOutcome(statusCode: number): string { * timestamp is older than the execution-state TTL — without this, the * `ts` field was dead data and a client could replay an ancient token * against a freshly-reused-execution-id window. */ -function decodeContinuationToken(token: string): { execution_id: string } | null { +function decodeContinuationToken( + token: string, +): { execution_id: string } | null { try { - const parsed: unknown = JSON.parse(Buffer.from(token, 'base64').toString('utf-8')); + const parsed: unknown = JSON.parse( + Buffer.from(token, 'base64').toString('utf-8'), + ); if (parsed === null || typeof parsed !== 'object') { return null; } const candidate = parsed as { execution_id?: unknown; ts?: unknown }; - if (typeof candidate.execution_id !== 'string' || candidate.execution_id.length === 0) { + if ( + typeof candidate.execution_id !== 'string' || + candidate.execution_id.length === 0 + ) { return null; } if (typeof candidate.ts === 'number' && Number.isFinite(candidate.ts)) { @@ -226,13 +301,17 @@ function decodeContinuationToken(token: string): { execution_id: string } | null // Blocking mode (legacy path) // --------------------------------------------------------------------------- -function waitForExecutionState(execution_id: string, timeout: number): ReturnType { +function waitForExecutionState( + execution_id: string, + timeout: number, +): ReturnType { return pollBlockingExecution(execution_id, timeout, { getExecutionState, getBlockingResult, - getPending: async (id) => { + getPending: async id => { const response = await retryToolCallServerRequest( - () => axios.get( + () => + axios.get( `${env.TOOL_CALL_SERVER_URL}/sessions/${id}/pending`, { headers: internalServiceHeaders() }, ), @@ -240,7 +319,8 @@ function waitForExecutionState(execution_id: string, timeout: number): ReturnTyp ); return response.data; }, - isNotFound: (error) => axios.isAxiosError(error) && error.response?.status === 404, + isNotFound: error => + axios.isAxiosError(error) && error.response?.status === 404, sleep: () => new Promise(resolve => setTimeout(resolve, POLL_INTERVAL)), now: Date.now, }); @@ -274,7 +354,10 @@ async function runReplayIteration( state: ExecutionState, apiKeyId: string, userId: string, + signal?: AbortSignal, + cancellation?: { requestId: string; owner: string }, ): Promise { + if (signal?.aborted) throw programmaticCancellationError(); const history = await loadToolHistory(state.execution_id); const rawPayload = buildReplayPayload(req, state, history); const sessionKey = state.sessionKey ?? state.userId; @@ -295,7 +378,8 @@ async function runReplayIteration( }); if (DEBUG_MODE) { - const firstFile = rawPayload.files[0] as { content?: string } | undefined; + const firstFile = rawPayload.files[0] as + { content?: string } | undefined; logger.debug('Replay enqueue details', { execution_id: state.execution_id, historySize: Object.keys(history).length, @@ -320,32 +404,83 @@ async function runReplayIteration( state.executionProfile ?? env.EXECUTION_PROFILE, state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, ); - const job = await queue.add(Jobs.execute, { - code: state.userCode ?? '', - userId, - payload: sandboxSecurity.payload, - apiKeyId, - isPyPlot: state.isPyPlot ?? false, - principalSource: state.principalSource, - executionId: state.execution_id, - tenantId: state.tenantId, - canonicalUserId: state.canonicalUserId, - executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, - sandboxBackend: replayBackend, - ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), - runtimeSessionMode: 'stateless', - runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, - executionManifestClaims: sandboxSecurity.executionManifestClaims, - egressGrantClaims: sandboxSecurity.egressGrantClaims, - egressGrantToken: sandboxSecurity.egressGrantToken, - }, { - removeOnComplete: { age: 60, count: 1 }, - removeOnFail: { age: 180, count: 1 }, - attempts: 1, - }); + if (signal?.aborted) throw programmaticCancellationError(); + const cancellationTarget = { queueName: queue.name, jobId: nanoid() }; + if (cancellation != null) { + const attachment = await attachProgrammaticCancellationTarget({ + redis: connection, + requestId: cancellation.requestId, + owner: cancellation.owner, + target: cancellationTarget, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + if (attachment === 'forbidden') { + throw new Error('Programmatic cancellation request ownership changed'); + } + if (attachment === 'cancelled') { + throw programmaticCancellationError(); + } + } + const submittedAtMs = Date.now(); + const deadlineAtMs = submittedAtMs + env.JOB_TIMEOUT; + let job: Awaited>; + try { + job = await queue.add( + Jobs.execute, + { + code: state.userCode ?? '', + userId, + payload: sandboxSecurity.payload, + apiKeyId, + isPyPlot: state.isPyPlot ?? false, + principalSource: state.principalSource, + executionId: state.execution_id, + tenantId: state.tenantId, + canonicalUserId: state.canonicalUserId, + executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, + sandboxBackend: replayBackend, + ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), + ...(state.workspaceId != null ? { workspaceId: state.workspaceId } : {}), + cancellable: true, + deadlineAtMs, + cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + runtimeSessionMode: 'stateless', + runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, + executionManifestClaims: sandboxSecurity.executionManifestClaims, + egressGrantClaims: sandboxSecurity.egressGrantClaims, + egressGrantToken: sandboxSecurity.egressGrantToken, + }, + { + removeOnComplete: { age: 60, count: 1 }, + removeOnFail: { age: 180, count: 1 }, + attempts: 1, + jobId: cancellationTarget.jobId, + timestamp: submittedAtMs, + }, + ); + } catch (error) { + // Redis may have enqueued the job even though its reply was lost. + // Preserve replay ownership until cancellation is durable or the job's + // fixed worker deadline prevents a late admission from executing. + const outcome = await fenceJobCancellation({ + commands: connection, target: cancellationTarget, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, deadlineAtMs, + }); + if (outcome.status === 'completed') return outcome.result; + throw error; + } jobsSubmitted.inc({ language }); - return job.waitUntilFinished(events, JOB_COMPLETION_WAIT_TIMEOUT_MS); + return waitForJobWithCancellation({ + commands: connection, + registry: jobCancellationRegistry, + job, + events, + timeoutMs: JOB_COMPLETION_WAIT_TIMEOUT_MS, + cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + deadlineAtMs, + signal, + }); } function isSandboxRunSuccess(result: t.ExecuteResult): boolean { @@ -365,18 +500,25 @@ async function handleReplayInitial( apiKeyId: string; userId: string; bridgeWorkerId?: string; + workspaceId?: string; }, + cancellation: ReplayRequestCancellation, ): Promise { - const { apiKeyId, userId, bridgeWorkerId } = params; - const { - code, - tools, - user_id, - files, - } = req.body as t.ProgrammaticRequestBody; + const { apiKeyId, userId, bridgeWorkerId, workspaceId } = params; + const { code, tools, user_id, files } = + req.body as t.ProgrammaticRequestBody; let timeout: number; try { - timeout = normalizeProgrammaticTimeoutMs((req.body as t.ProgrammaticRequestBody).timeout); + if (workspaceId != null && Array.isArray(files) && files.length > BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES) { + throw new Error(`Selected-workspace execution allows at most ${BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_INPUT_FILES} input files; main and replay history occupy two reserved slots`); + } + timeout = workspaceId != null + ? normalizeSelectedWorkspaceProgrammaticTimeoutMs( + (req.body as t.ProgrammaticRequestBody).timeout, + ) + : normalizeProgrammaticTimeoutMs( + (req.body as t.ProgrammaticRequestBody).timeout, + ); } catch (error) { res.status(400).json({ error: (error as Error).message }); return; @@ -400,14 +542,23 @@ async function handleReplayInitial( }); return; } - const language: 'python' | 'bash' = requestedLanguage === 'bash' ? 'bash' : 'python'; + const language: 'python' | 'bash' = + requestedLanguage === 'bash' ? 'bash' : 'python'; + if (workspaceId != null && language !== 'bash') { + res.status(400).json({ + error: 'Selected-workspace programmatic execution supports bash only', + }); + return; + } if (!code) { res.status(400).json({ error: 'Missing required field: code' }); return; } - if (!tools || !Array.isArray(tools) || tools.length === 0) { - res.status(400).json({ error: 'Missing required field: tools (must be a non-empty array)' }); + if (!Array.isArray(tools) || (tools.length === 0 && workspaceId == null)) { + res.status(400).json({ + error: 'Missing required field: tools (must be non-empty unless a selected workspace executes bash)', + }); return; } if (tools.length > MAX_TOOLS_PER_REQUEST) { @@ -442,7 +593,8 @@ async function handleReplayInitial( files, store: connection, }); - (req.body as t.ProgrammaticRequestBody).files = authorizedFiles.length > 0 ? authorizedFiles : undefined; + (req.body as t.ProgrammaticRequestBody).files = + authorizedFiles.length > 0 ? authorizedFiles : undefined; } catch (error) { if (sendFileRefAuthorizationError(error, res, req)) return; logger.error('Error authorizing replay file refs:', error); @@ -456,19 +608,39 @@ async function handleReplayInitial( try { sessionKey = resolveOutputBucketSessionKey(req); } catch (error) { - if (sendSessionKeyResolutionError(error, res, req, 'programmatic /exec: resolveOutputBucketSessionKey')) { + if ( + sendSessionKeyResolutionError( + error, + res, + req, + 'programmatic /exec: resolveOutputBucketSessionKey', + ) + ) { return; } throw error; } + if ( + cancellation.signal.aborted || + cancellation.request?.cancelledBeforeStart === true + ) { + if (!cancellation.isDisconnected()) { + res.status(200).json({ + status: 'error', + error: 'Programmatic execution request cancelled', + }); + } + return; + } + const session_id = nanoid(); const execution_id = nanoid(); const authContext = req.codeApiAuthContext; const identity = getExecutionIdentity(req, userId); - const isPyPlot = language === 'python' && ( - code.includes('import matplotlib') || code.includes('import seaborn') - ); + const isPyPlot = + language === 'python' && + (code.includes('import matplotlib') || code.includes('import seaborn')); await recordSessionOwnership(connection, session_id, sessionKey); @@ -487,6 +659,7 @@ async function handleReplayInitial( timeout, language, bridgeWorkerId, + workspaceId, executionProfile: env.EXECUTION_PROFILE, executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: resolveReplayStateSandboxBackend({ @@ -506,13 +679,16 @@ async function handleReplayInitial( await setExecutionState(state); } catch (err) { if (err instanceof ExecutionStateTooLargeError) { - logger.warn('Rejecting replay request: ExecutionState exceeds Redis cap', { + logger.warn( + 'Rejecting replay request: ExecutionState exceeds Redis cap', + { execution_id, userId, apiKeyId, bytes: err.bytes, cap: err.cap, - }); + }, + ); await clearSessionOwnership(connection, session_id).catch(() => {}); ptcReplayStateOversize.inc(); res.status(413).json({ @@ -537,7 +713,7 @@ async function handleReplayInitial( timeout, }); - await runAndRespond(req, res, state, apiKeyId, userId); + await runAndRespond(req, res, state, apiKeyId, userId, cancellation); } async function handleReplayContinuation( @@ -549,6 +725,7 @@ async function handleReplayContinuation( decoded: { execution_id: string }; tool_results: NonNullable; }, + cancellation: ReplayRequestCancellation, ): Promise { const { apiKeyId, userId, decoded, tool_results } = params; @@ -568,9 +745,15 @@ async function handleReplayContinuation( * adds outcome plumbing through `runAndRespond`. */ const startMs = performance.now(); res.once('finish', () => { - const labels = { mode: 'replay' as const, outcome: classifyContinuationOutcome(res.statusCode) }; + const labels = { + mode: 'replay' as const, + outcome: classifyContinuationOutcome(res.statusCode), + }; ptcReplayContinuations.inc(labels); - ptcReplayContinuationDuration.observe(labels, (performance.now() - startMs) / 1000); + ptcReplayContinuationDuration.observe( + labels, + (performance.now() - startMs) / 1000, + ); }); /** Reject oversized batches before we spend any CPU on per-entry @@ -601,6 +784,20 @@ async function handleReplayContinuation( res.status(404).json({ error: 'Execution not found or expired' }); return; } + if ( + cancellation.signal.aborted || + cancellation.request?.cancelledBeforeStart === true + ) { + await cleanupExecution(state.execution_id, 'replay'); + if (!cancellation.isDisconnected()) { + res.status(200).json({ + status: 'error', + error: 'Programmatic execution request cancelled', + session_id: state.session_id, + }); + } + return; + } /** Compute the delta against already-persisted history first so the * cap checks see the real impact of this batch (new call_ids only * advance `callCount`; overwrites may shrink or grow `historyBytes` @@ -621,9 +818,14 @@ async function handleReplayContinuation( call_site: emitted.call_site, }; }); - const deltaOrError = await computeToolHistoryDelta(state.execution_id, enrichedResults); + const deltaOrError = await computeToolHistoryDelta( + state.execution_id, + enrichedResults, + ); if ('error' in deltaOrError) { - res.status(deltaOrError.status ?? 400).json({ error: deltaOrError.error }); + res.status(deltaOrError.status ?? 400).json({ + error: deltaOrError.error, + }); return; } const delta = deltaOrError; @@ -642,7 +844,9 @@ async function handleReplayContinuation( }); if (!pre.ok) { if (pre.status === 403) { - logger.warn('Unauthorized replay continuation request rejected', { + logger.warn( + 'Unauthorized replay continuation request rejected', + { execution_id: state.execution_id, requestUserId: userId, requestApiKeyId: apiKeyId, @@ -650,7 +854,8 @@ async function handleReplayContinuation( executionUserId: state.userId, executionApiKeyId: state.apiKeyId, executionTenantId: state.tenantId, - }); + }, + ); } if (pre.cleanupOnReject === true) { await cleanupExecution(state.execution_id, 'replay'); @@ -672,7 +877,10 @@ async function handleReplayContinuation( * Redis MULTI/EXEC so counters and the hash can't drift out of sync * on a partial failure. */ state.callCount = (state.callCount ?? 0) + delta.newCallIds.length; - state.historyBytes = Math.max(0, (state.historyBytes ?? 0) + delta.bytesDelta); + state.historyBytes = Math.max( + 0, + (state.historyBytes ?? 0) + delta.bytesDelta, + ); state.lastActivity = Date.now(); try { await commitToolHistoryAndState(state, delta); @@ -687,14 +895,19 @@ async function handleReplayContinuation( * forward is a fresh execution with smaller inputs. Reap the * old execution to free the lock and Redis keys, then return * an actionable 413 instead of a generic 500. */ - logger.warn('Replay continuation rejected: ExecutionState exceeds Redis cap', { + logger.warn( + 'Replay continuation rejected: ExecutionState exceeds Redis cap', + { execution_id: state.execution_id, bytes: err.bytes, cap: err.cap, callCount: state.callCount, historyBytes: state.historyBytes, - }); - await cleanupExecution(state.execution_id, 'replay').catch(() => {}); + }, + ); + await cleanupExecution(state.execution_id, 'replay').catch( + () => {}, + ); ptcReplayStateOversize.inc(); res.status(413).json({ status: 'error', @@ -715,10 +928,13 @@ async function handleReplayContinuation( * the throw bubble to the top-level catch and become an opaque * 500 — clients (and load balancers) treat 5xx classes very * differently for retry policy. */ - logger.error('Failed to commit replay continuation; returning retryable 503', { + logger.error( + 'Failed to commit replay continuation; returning retryable 503', + { execution_id: state.execution_id, err: (err as Error).message, - }); + }, + ); res.status(503).json({ status: 'error', error: 'Failed to persist replay continuation; please retry the same request', @@ -735,7 +951,14 @@ async function handleReplayContinuation( }); } - await runAndRespond(req, res, state, apiKeyId, userId); + await runAndRespond( + req, + res, + state, + apiKeyId, + userId, + cancellation, + ); } finally { await releaseExecutionLock(decoded.execution_id, lockToken); } @@ -747,28 +970,32 @@ async function runAndRespond( state: ExecutionState, apiKeyId: string, userId: string, + cancellation: ReplayRequestCancellation, ): Promise { - /** Read disconnect state through `isDisconnected()` rather than a - * direct boolean. The `req.on('close', ...)` handler flips the flag - * during awaits, but `@typescript-eslint/no-unnecessary-condition` - * (correctly per TS semantics) narrows a directly-mutated `let`/object - * member to its literal value after an early-return `if (...) return`, - * even across awaits. A function call is opaque to that narrowing. */ - let disconnected = false; - const isDisconnected = (): boolean => disconnected; - req.on('close', () => { - if (!res.writableEnded) disconnected = true; - }); - let result: t.ExecuteResult; try { - result = await runReplayIteration(req, state, apiKeyId, userId); + result = await runReplayIteration( + req, + state, + apiKeyId, + userId, + cancellation.signal, + cancellation.request, + ); } catch (err) { - logger.error('Replay iteration failed', { execution_id: state.execution_id, err }); + const cancelled = + (err as Error).name === 'AbortError' || + (err as Error).message === JOB_CANCELLED_MESSAGE; + logger.log(cancelled ? 'info' : 'error', 'Replay iteration failed', { + execution_id: state.execution_id, + cancelled, + err, + }); await cleanupExecution(state.execution_id, 'replay'); - if (!isDisconnected()) { + if (!cancellation.isDisconnected()) { const publicFailure = publicExecutionFailure(err); - const message = publicFailure?.body.message ?? (err as Error).message; + const message = + publicFailure?.body.message ?? (err as Error).message; res.status(200).json({ status: 'error', error: message !== '' ? message : 'Sandbox execution failed', @@ -778,7 +1005,7 @@ async function runAndRespond( return; } - if (isDisconnected()) { + if (cancellation.isDisconnected()) { logger.info('Client disconnected during replay; cleaning up', { execution_id: state.execution_id, }); @@ -786,10 +1013,19 @@ async function runAndRespond( return; } - const { stdout: cleanStdout, pending } = extractPendingFromStdout( + const extracted = extractPendingFromStdout( result.stdout, state.execution_id, ); + const cleanStdout = extracted.stdout; + const controlPayload = result.pending_tool_calls_payload; + const hasControlPayload = typeof controlPayload === 'string'; + const controlPending = hasControlPayload + ? extractPendingFromControlPayload(controlPayload) + : null; + const pending = hasControlPayload + ? (controlPending ?? []) + : extracted.pending; if (pending != null) { if (pending.length === 0) { @@ -806,7 +1042,10 @@ async function runAndRespond( }); return; } - const unregisteredToolCall = findUnregisteredToolCall(pending, state.tools); + const unregisteredToolCall = findUnregisteredToolCall( + pending, + state.tools, + ); if (unregisteredToolCall != null) { logger.warn('Sandbox requested unregistered replay tool call', { execution_id: state.execution_id, @@ -866,12 +1105,17 @@ async function runAndRespond( await setExecutionState(state); await refreshExecutionTtl(state.execution_id); } catch (err) { - logger.error('Failed to persist execution state before continuation; aborting', { + logger.error( + 'Failed to persist execution state before continuation; aborting', + { execution_id: state.execution_id, err: (err as Error).message, - }); - await cleanupExecution(state.execution_id, 'replay').catch(() => {}); - if (!isDisconnected()) { + }, + ); + await cleanupExecution(state.execution_id, 'replay').catch( + () => {}, + ); + if (!cancellation.isDisconnected()) { if (err instanceof ExecutionStateTooLargeError) { /** A continuation that pushes `emittedCallIds` past the * `MAX_EXECUTION_STATE_BYTES` cap is a client-input sizing @@ -887,8 +1131,7 @@ async function runAndRespond( } else { res.status(503).json({ status: 'error', - error: - 'Failed to persist replay state; please retry the request from scratch', + error: 'Failed to persist replay state; please retry the request from scratch', session_id: state.session_id, }); } @@ -912,7 +1155,8 @@ async function runAndRespond( if (!isSandboxRunSuccess(result)) { await cleanupExecution(state.execution_id, 'replay'); - const errorMessage = result.message != null && result.message !== '' + const errorMessage = + result.message != null && result.message !== '' ? result.message : `Sandbox exited with code ${result.code ?? 'unknown'}`; res.status(200).json({ @@ -920,6 +1164,10 @@ async function runAndRespond( error: errorMessage, stdout: cleanStdout, stderr: result.stderr, + files: result.files, + deleted_files: result.deleted_files, + artifact_delivery: result.artifact_delivery, + artifact_truncation: result.artifact_truncation, session_id: state.session_id, }); return; @@ -932,7 +1180,9 @@ async function runAndRespond( stdout: cleanStdout, stderr: result.stderr, files: result.files, + deleted_files: result.deleted_files, artifact_delivery: result.artifact_delivery, + artifact_truncation: result.artifact_truncation, session_id: state.session_id, }); } @@ -941,7 +1191,71 @@ async function runAndRespond( // Request entrypoint // --------------------------------------------------------------------------- -router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedRequest, res) => { +router.post( + '/exec/programmatic/cancel', + cancellationLimiter, + async (req: t.AuthenticatedRequest, res) => { + const principal = getPrincipalOrReject(req, res); + if (!principal) return; + const requestId = normalizeProgrammaticRequestId( + (req.body as Record)?.request_id, + ); + if (requestId == null) { + res.status(400).json({ error: 'Invalid or missing request_id' }); + return; + } + try { + const owner = programmaticCancellationOwner(req, principal.userId); + const cancellation = await cancelProgrammaticRequest({ + redis: connection, + requestId, + owner, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + if (cancellation.status === 'forbidden') { + res.status(403).json({ error: 'Programmatic request belongs to another principal' }); + return; + } + if (cancellation.target != null) { + const accepted = await requestJobCancellation( + connection, + cancellation.target, + PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + ); + if (!accepted) { + res.status(200).json({ status: 'already_completed' }); + return; + } + try { + const queuedJob = await getExistingExecutionJob( + cancellation.target.queueName, + cancellation.target.jobId, + ); + if (queuedJob != null) await removeJobIfWaiting(queuedJob); + } catch (error) { + logger.warn('Failed to remove cancelled waiting execution', { + requestId, + queueName: cancellation.target?.queueName, + jobId: cancellation.target?.jobId, + error: (error as Error).message, + }); + } + } + res.status(202).json({ status: 'cancellation_requested' }); + } catch (error) { + logger.error('Failed to request programmatic execution cancellation', { + requestId, + error: (error as Error).message, + }); + res.status(503).json({ error: 'Cancellation service unavailable' }); + } + }, +); + +router.post( + '/exec/programmatic', + executionLimiter, + async (req: t.AuthenticatedRequest, res) => { const principal = getPrincipalOrReject(req, res); if (!principal) return; const apiKeyId = getCredentialId(req); @@ -954,13 +1268,17 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR return res.status(503).json({ error: 'Service is starting up' }); } - const { - continuation_token, - tool_results, - } = req.body as t.ProgrammaticRequestBody; + const { continuation_token, tool_results } = + req.body as t.ProgrammaticRequestBody; const rawBody = req.body as Record; + const rawRequestId = req.header(CODEAPI_PROGRAMMATIC_REQUEST_HEADER); + const requestId = normalizeProgrammaticRequestId(rawRequestId); + if (rawRequestId != null && requestId == null) { + return res.status(400).json({ error: 'Invalid programmatic request ID' }); + } const requestedLanguage: unknown = rawBody.language ?? rawBody.lang; let bridgeWorkerId: string | undefined; + let workspaceId: string | undefined; if (continuation_token == null || continuation_token === '') { try { const bridgeSelection = resolveBridgeWorkerSelection({ @@ -970,12 +1288,35 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), trustedWorkerId: principal.codeWorkerId, }); - bridgeWorkerId = bridgeSelection?.explicit === true - ? bridgeSelection.workerId - : undefined; + bridgeWorkerId = + bridgeSelection?.explicit === true || + (bridgeSelection != null && !env.BRIDGE_DYNAMIC_WORKERS) + ? bridgeSelection.workerId + : undefined; + const requestedWorkspaceId = req + .header(CODEAPI_BRIDGE_WORKSPACE_HEADER) + ?.trim(); + if ( + requestedWorkspaceId != null && + requestedWorkspaceId !== '' + ) { + if (bridgeWorkerId == null) { + return res.status(400).json({ + error: 'Workspace selection requires an authenticated bridge worker', + }); + } + if (!isValidBridgeWorkerId(requestedWorkspaceId)) { + return res + .status(400) + .json({ error: 'Invalid code workspace ID' }); + } + workspaceId = requestedWorkspaceId; + } } catch (error) { if (error instanceof BridgeWorkerSelectionError) { - return res.status(error.status).json({ error: error.message }); + return res + .status(error.status) + .json({ error: error.message }); } throw error; } @@ -991,7 +1332,51 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR }); } + const disconnectObserver = observeRequestDisconnect(req, res); + + const cancellation: ReplayRequestCancellation = { + signal: disconnectObserver.signal, + isDisconnected: disconnectObserver.isDisconnected, + }; + let reservedCancellation: { requestId: string; owner: string } | undefined; + try { + if (requestId != null) { + const owner = programmaticCancellationOwner(req, userId); + let reservation: Awaited>; + try { + reservation = await reserveProgrammaticCancellation({ + redis: connection, + requestId, + owner, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + } catch (error) { + logger.error('Failed to reserve programmatic cancellation request', { + requestId, + error: (error as Error).message, + }); + if (!cancellation.isDisconnected()) { + return res.status(503).json({ error: 'Cancellation service unavailable' }); + } + return; + } + if (reservation === 'forbidden' || reservation === 'duplicate') { + if (!cancellation.isDisconnected()) { + return res.status(409).json({ + error: 'Programmatic request ID is already in use', + }); + } + return; + } + reservedCancellation = { requestId, owner }; + cancellation.request = { + requestId, + owner, + cancelledBeforeStart: reservation === 'cancelled', + }; + } + /** For continuations, peek at the stored execution to route by the * mode it was started in rather than the current process default. * Without this, a replay-mode execution resumed via an instance @@ -1013,7 +1398,9 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR } const decoded = decodeContinuationToken(continuation_token); if (!decoded) { - return res.status(400).json({ error: 'Invalid continuation token' }); + return res + .status(400) + .json({ error: 'Invalid continuation token' }); } const existing = await getExecutionState(decoded.execution_id); if (existing?.mode === 'replay') { @@ -1022,7 +1409,7 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR userId, decoded, tool_results, - }); + }, cancellation); } return await handleBlocking(req, res, { apiKeyId, userId }); } @@ -1036,17 +1423,46 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR }); } if (env.PTC_MODE === 'replay') { - return await handleReplayInitial(req, res, { apiKeyId, userId, bridgeWorkerId }); + return await handleReplayInitial(req, res, { + apiKeyId, + userId, + bridgeWorkerId, + workspaceId, + }, cancellation); + } + if (workspaceId != null) { + return res.status(400).json({ + error: 'Selected-workspace programmatic execution requires replay mode', + }); } - return await handleBlocking(req, res, { apiKeyId, userId, bridgeWorkerId }); + return await handleBlocking(req, res, { + apiKeyId, + userId, + bridgeWorkerId, + }); } catch (err) { logger.error(`[${INSTANCE_ID}] Programmatic routing error:`, err); if (!res.headersSent) { return res.status(500).json({ error: 'Internal server error' }); } return; + } finally { + disconnectObserver.dispose(); + if (reservedCancellation != null) { + await releaseProgrammaticCancellation({ + redis: connection, + requestId: reservedCancellation.requestId, + owner: reservedCancellation.owner, + }).catch(error => { + logger.warn('Failed to release programmatic cancellation request', { + requestId: reservedCancellation?.requestId, + error: (error as Error).message, + }); + }); + } } -}); + }, +); // --------------------------------------------------------------------------- // Blocking-mode handler (extracted from the original implementation). @@ -1059,46 +1475,47 @@ async function handleBlocking( params: { apiKeyId: string; userId: string; bridgeWorkerId?: string }, ): Promise> { const { apiKeyId, userId, bridgeWorkerId } = params; - const { - code, - tools, - user_id, - files, - continuation_token, - tool_results, - } = req.body as t.ProgrammaticRequestBody; + const { code, tools, user_id, files, continuation_token, tool_results } = + req.body as t.ProgrammaticRequestBody; let timeout: number; try { - timeout = normalizeProgrammaticTimeoutMs((req.body as t.ProgrammaticRequestBody).timeout); + timeout = normalizeProgrammaticTimeoutMs( + (req.body as t.ProgrammaticRequestBody).timeout, + ); } catch (error) { return res.status(400).json({ error: (error as Error).message }); } // CASE 1: Continuation - if (continuation_token != null && continuation_token !== '' && tool_results) { + if ( + continuation_token != null && + continuation_token !== '' && + tool_results + ) { const decoded = decodeContinuationToken(continuation_token); if (!decoded) { - return res.status(400).json({ error: 'Invalid continuation token' }); + return res + .status(400) + .json({ error: 'Invalid continuation token' }); } const { execution_id } = decoded; const execution = await getExecutionState(execution_id); if (!execution) { - return res.status(404).json({ error: 'Execution not found or expired' }); + return res + .status(404) + .json({ error: 'Execution not found or expired' }); } const identity = getExecutionIdentity(req, userId); if ( execution.userId !== userId || (execution.apiKeyId != null && execution.apiKeyId !== apiKeyId) || - ( - execution.tenantId != null && - execution.tenantId !== identity.storageNamespace - ) || - ( - execution.authContextHash != null && - execution.authContextHash !== req.codeApiAuthContext?.authContextHash - ) + (execution.tenantId != null && + execution.tenantId !== identity.storageNamespace) || + (execution.authContextHash != null && + execution.authContextHash !== + req.codeApiAuthContext?.authContextHash) ) { logger.warn('Unauthorized blocking continuation request rejected', { execution_id, @@ -1122,14 +1539,19 @@ async function handleBlocking( try { await retryToolCallServerRequest( - () => axios.post(`${env.TOOL_CALL_SERVER_URL}/sessions/${execution_id}/results`, { + () => + axios.post( + `${env.TOOL_CALL_SERVER_URL}/sessions/${execution_id}/results`, + { results: tool_results.map(r => ({ call_id: r.call_id, result: r.result, is_error: r.is_error ?? false, error_message: r.error_message, })), - }, { headers: internalServiceHeaders() }), + }, + { headers: internalServiceHeaders() }, + ), 'Submit tool results', ); @@ -1151,7 +1573,9 @@ async function handleBlocking( stdout: state.stdout ?? '', stderr: state.stderr ?? '', files: state.files ?? [], + deleted_files: state.deleted_files, artifact_delivery: state.artifact_delivery, + artifact_truncation: state.artifact_truncation, session_id: execution.session_id, }); } @@ -1174,14 +1598,21 @@ async function handleBlocking( return res.status(400).json({ error: 'Missing required field: code' }); } if (!tools || !Array.isArray(tools) || tools.length === 0) { - return res.status(400).json({ error: 'Missing required field: tools (must be a non-empty array)' }); + return res + .status(400) + .json({ + error: 'Missing required field: tools (must be a non-empty array)', + }); } if (tools.length > MAX_TOOLS_PER_REQUEST) { - logger.warn(`Too many tools provided: ${tools.length}, limit is ${MAX_TOOLS_PER_REQUEST}`, { + logger.warn( + `Too many tools provided: ${tools.length}, limit is ${MAX_TOOLS_PER_REQUEST}`, + { execution_id: 'pre-creation', userId, toolCount: tools.length, - }); + }, + ); return res.status(400).json({ error: `Too many tools provided (${tools.length}). Maximum is ${MAX_TOOLS_PER_REQUEST}.`, }); @@ -1202,7 +1633,8 @@ async function handleBlocking( files, store: connection, }); - (req.body as t.ProgrammaticRequestBody).files = authorizedFiles.length > 0 ? authorizedFiles : undefined; + (req.body as t.ProgrammaticRequestBody).files = + authorizedFiles.length > 0 ? authorizedFiles : undefined; } catch (error) { if (sendFileRefAuthorizationError(error, res, req)) return; logger.error('Error authorizing programmatic file refs:', error); @@ -1215,7 +1647,14 @@ async function handleBlocking( try { sessionKey = resolveOutputBucketSessionKey(req); } catch (error) { - if (sendSessionKeyResolutionError(error, res, req, 'programmatic /exec-blocking: resolveOutputBucketSessionKey')) { + if ( + sendSessionKeyResolutionError( + error, + res, + req, + 'programmatic /exec-blocking: resolveOutputBucketSessionKey', + ) + ) { return; } throw error; @@ -1270,24 +1709,34 @@ async function handleBlocking( try { callbackUrl = normalizeEgressGatewayUrl(env.EGRESS_GATEWAY_URL); } catch (error) { - logger.error('Blocking PTC requires egress gateway callback URL:', error); + logger.error( + 'Blocking PTC requires egress gateway callback URL:', + error, + ); await cleanupExecution(execution_id, 'blocking'); - return res.status(503).json({ error: 'Egress gateway unavailable' }); + return res + .status(503) + .json({ error: 'Egress gateway unavailable' }); } let callbackToken: string; try { const toolCallResponse = await retryToolCallServerRequest( - () => axios.post<{ + () => + axios.post<{ success: boolean; callback_token: string; - }>(`${env.TOOL_CALL_SERVER_URL}/sessions`, { + }>( + `${env.TOOL_CALL_SERVER_URL}/sessions`, + { execution_id, session_id, timeout, tools, - }, { headers: internalServiceHeaders() }), + }, + { headers: internalServiceHeaders() }, + ), 'Create Tool Call Server session', ); @@ -1299,9 +1748,14 @@ async function handleBlocking( allowedToolNames: tools.map(tool => tool.name), }); } catch (error) { - logger.error('Failed to create Tool Call Server session or callback token:', error); + logger.error( + 'Failed to create Tool Call Server session or callback token:', + error, + ); await cleanupExecution(execution_id, 'blocking'); - return res.status(503).json({ error: 'Tool Call Server unavailable' }); + return res + .status(503) + .json({ error: 'Tool Call Server unavailable' }); } let rawPayload: t.PayloadBody; @@ -1316,10 +1770,15 @@ async function handleBlocking( timeout, }); } catch (error) { - logger.error('Failed to create payload', { execution_id, error: (error as Error).message }); + logger.error('Failed to create payload', { + execution_id, + error: (error as Error).message, + }); await cleanupExecution(execution_id, 'blocking'); return res.status(400).json({ - error: (error as Error).message || 'Failed to generate code payload', + error: + (error as Error).message || + 'Failed to generate code payload', }); } const sandboxSecurity = prepareSandboxJobSecurity({ @@ -1331,7 +1790,9 @@ async function handleBlocking( payload: rawPayload, }); - const job = await pyQueue.add(Jobs.execute, { + const job = await pyQueue.add( + Jobs.execute, + { code, userId, payload: sandboxSecurity.payload, @@ -1350,18 +1811,24 @@ async function handleBlocking( ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, - executionManifestClaims: sandboxSecurity.executionManifestClaims, + executionManifestClaims: + sandboxSecurity.executionManifestClaims, egressGrantClaims: sandboxSecurity.egressGrantClaims, egressGrantToken: sandboxSecurity.egressGrantToken, - }, { + }, + { removeOnComplete: { age: 60, count: 1 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, jobId: session_id, - }); + }, + ); jobsSubmitted.inc({ language: 'python' }); - logger.info('Job queued, polling for tool calls', { execution_id, session_id }); + logger.info('Job queued, polling for tool calls', { + execution_id, + session_id, + }); let clientDisconnected = false; req.on('close', async () => { @@ -1372,21 +1839,27 @@ async function handleBlocking( await job.remove(); await cleanupExecution(execution_id, 'blocking'); } catch (error) { - logger.error('Error cleaning up after client disconnect:', error); + logger.error( + 'Error cleaning up after client disconnect:', + error, + ); } }); job.waitUntilFinished(pyQueueEvents, JOB_COMPLETION_WAIT_TIMEOUT_MS) - .then(async (result) => { + .then(async result => { if (clientDisconnected) return; await setExecutionResult(execution_id, result); }) - .catch(async (error) => { + .catch(async error => { if (clientDisconnected) return; await setExecutionError(execution_id, error); }); - const state = await waitForExecutionState(execution_id, Math.min(timeout, MAX_POLL_TIME)); + const state = await waitForExecutionState( + execution_id, + Math.min(timeout, MAX_POLL_TIME), + ); if (state.status === 'waiting' && state.pending_calls) { return res.status(200).json({ @@ -1404,7 +1877,9 @@ async function handleBlocking( stdout: state.stdout ?? '', stderr: state.stderr ?? '', files: state.files ?? [], + deleted_files: state.deleted_files, artifact_delivery: state.artifact_delivery, + artifact_truncation: state.artifact_truncation, session_id, }); } @@ -1416,7 +1891,10 @@ async function handleBlocking( session_id, }); } catch (error) { - logger.error(`[${INSTANCE_ID}] Session ID: ${session_id} | Execution ID: ${execution_id} | Error:`, error); + logger.error( + `[${INSTANCE_ID}] Session ID: ${session_id} | Execution ID: ${execution_id} | Error:`, + error, + ); await cleanupExecution(execution_id, 'blocking'); return res.status(500).json({ error: 'Internal server error' }); } diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index fc84d8f8..81405021 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -83,6 +83,7 @@ describe('buildReplayExecutionState', () => { const state = build({ authContext, bridgeWorkerId: 'code-user_123', + workspaceId: 'project-a', sandboxBackend: 'remote-bridge', executionProfile: 'stateful', executionProfileSource: 'explicit', @@ -102,6 +103,7 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', apiKeyId: 'key_legacy', bridgeWorkerId: 'code-user_123', + workspaceId: 'project-a', sandboxBackend: 'remote-bridge', executionProfile: 'stateful', executionProfileSource: 'explicit', diff --git a/service/src/service/programmatic-state.ts b/service/src/service/programmatic-state.ts index 25571fed..e6bda59a 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -38,6 +38,7 @@ export interface BuildReplayExecutionStateParams { timeout: number; language: 'python' | 'bash'; bridgeWorkerId?: string; + workspaceId?: string; sandboxBackend?: SandboxBackendName; executionProfile: ExecutionProfile; executionProfileSource: ExecutionProfileSource; @@ -66,6 +67,7 @@ export function buildReplayExecutionState( authContextHash: identity.authContextHash, apiKeyId: params.apiKeyId, bridgeWorkerId: params.bridgeWorkerId, + workspaceId: params.workspaceId, sandboxBackend: params.sandboxBackend, executionProfile: params.executionProfile, executionProfileSource: params.executionProfileSource, diff --git a/service/src/service/replay-state.ts b/service/src/service/replay-state.ts index 2254ee21..3b65cedf 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -116,6 +116,8 @@ export interface ExecutionState { apiKeyId?: string; /** Authenticated worker selection retained across every replay iteration. */ bridgeWorkerId?: string; + /** Selected workspace retained and bound across every replay iteration. */ + workspaceId?: string; /** Original queue/backend target retained across replay continuations. */ sandboxBackend?: SandboxBackendName; /** Original producer profile retained so continuations use the same queue. */ diff --git a/service/src/types/service.ts b/service/src/types/service.ts index 2a90eac7..0404ad16 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -3,7 +3,10 @@ import type { Request } from 'express'; import type { ExecutionManifestClaims } from '../execution-manifest'; import type { ExecutionIdentity } from '../execution-identity'; import type { CodeApiPrincipal } from '../auth/principal'; -import type { ExecutionProfile, SandboxBackendName } from '../execution-profile'; +import type { + ExecutionProfile, + SandboxBackendName, +} from '../execution-profile'; import { Jobs } from '@/enum/service'; /** @@ -111,6 +114,15 @@ export interface ArtifactDeliveryFailure { failed: number; } +export type ArtifactTruncationReason = 'max_files' | 'depth' | 'size' | 'path' | 'unreadable'; + +export interface ArtifactTruncation { + code: 'artifact_truncated'; + reasons: Partial>; + skipped: string[]; + skipped_count: number; +} + export type ExecuteResponse = { run?: { stdout: string; @@ -129,7 +141,9 @@ export type ExecuteResponse = { /** Top-level execution session id (one sandbox `/exec` invocation). */ session_id: string; files: FileRefs; + deleted_files?: string[]; artifact_delivery?: ArtifactDeliveryFailure; + artifact_truncation?: ArtifactTruncation; }; export interface RequestBody { @@ -149,7 +163,11 @@ export interface RequestBody { runtime_session_hint?: string; } -export type CreatePayload = { req: AuthenticatedRequest, session_id: string; isPyPlot?: boolean }; +export type CreatePayload = { + req: AuthenticatedRequest; + session_id: string; + isPyPlot?: boolean; +}; export interface FileObject { name: string; id: string; @@ -160,10 +178,12 @@ export interface FileObject { size?: number; lastModified?: string; etag?: string; - metadata?: { + metadata?: + | { 'content-type': string; 'original-filename': string; - } | undefined; + } + | undefined; versionId?: string | null; contentType?: string; } @@ -184,6 +204,14 @@ export type PayloadFileRef = { export interface PayloadBody { language: string; version: string; + /** Stable identity shared by all replay iterations of one execution. */ + execution_id?: string; + replay_tool_count?: number; + /** Manifest-bound upload ceiling exposed to remote workers. */ + max_output_files?: number; + /** Effective per-file ceiling after manifest and gateway policy intersect. */ + max_output_file_bytes?: number; + transfer_timeout_ms?: number; run_memory_limit?: number; run_timeout?: number; run_cpu_time?: number; @@ -232,12 +260,16 @@ export type ExecuteResult = { stdout: string; stderr: string; files: FileRefs; + deleted_files?: string[]; artifact_delivery?: ArtifactDeliveryFailure; + artifact_truncation?: ArtifactTruncation; code?: number | null; signal?: string | null; message?: string | null; status?: string | null; wall_time?: number | null; + /** Trusted worker control channel; avoids losing replay calls to stdout truncation. */ + pending_tool_calls_payload?: string; }; export interface LanguageConfig { @@ -265,6 +297,14 @@ export type JobData = { canonicalUserId?: string; /** Trusted dynamic outbound worker selection. */ bridgeWorkerId?: string; + /** Trusted selected workspace for native replay-mode PTC. */ + workspaceId?: string; + /** Opts replay jobs into durable client-disconnect cancellation. */ + cancellable?: boolean; + /** Absolute producer budget; queue-worker configuration may only tighten it. */ + deadlineAtMs?: number; + /** Producer request tombstones must never outlive the completion decision. */ + cancellationTtlSeconds?: number; /** Producer deployment identity. Optional only for pre-profile queued jobs. */ executionProfile?: ExecutionProfile; /** Required sandbox transport. Optional only for jobs queued before fencing. */ @@ -370,7 +410,9 @@ export interface ProgrammaticResponse { stdout?: string; stderr?: string; files?: FileRefs; + deleted_files?: string[]; artifact_delivery?: ArtifactDeliveryFailure; + artifact_truncation?: ArtifactTruncation; /** Top-level execution session id (one sandbox PTC invocation). */ session_id?: string; tool_calls_made?: number; diff --git a/service/src/workers.ts b/service/src/workers.ts index ad20fd0f..dbfd544b 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -1,68 +1,148 @@ import axios from 'axios'; import { Worker } from 'bullmq'; import type * as t from './types'; -import { filterSystemLogs, applySystemReplacements, getAxiosErrorDetails, sandboxErrorMessageFromAxios } from './utils'; -import { jobProcessingDuration, jobsCompleted, jobsFailed, activeJobs, workerRunning } from './metrics'; -import { connection, queueNames } from './queue'; +import { + filterSystemLogs, + applySystemReplacements, + getAxiosErrorDetails, + sandboxErrorMessageFromAxios, +} from './utils'; +import { + jobProcessingDuration, + jobsCancelled, + jobsCompleted, + jobsFailed, + activeJobs, + workerRunning, +} from './metrics'; +import { connection, jobCancellationRegistry, queueNames } from './queue'; import { env, jobDeadlineAtMs } from './config'; import { summarizeSandboxResponse, summarizeText } from './execution-log'; -import { createGatewayEgressGrant, restoreGatewaySandboxResult, revokeGatewayEgressGrant } from './egress-gateway-client'; +import { + createGatewayEgressGrant, + restoreGatewaySandboxResult, + revokeGatewayEgressGrant, +} from './egress-gateway-client'; import { refreshEgressGrantClaims } from './sandbox-egress'; import { buildSandboxExecuteRequest } from './sandbox-dispatch'; import { prepareInputDelivery } from './runtime-session/input-delivery'; import { SessionFilesError } from './runtime-session/files'; import { resolveRuntimeSessionForJob } from './runtime-session/job-policy'; -import { getSandboxBackend, SandboxBackendError, type SandboxRawResponse } from './sandbox-backend'; +import { + getSandboxBackend, + SandboxBackendError, + type SandboxRawResponse, +} from './sandbox-backend'; import { isSyntheticPrincipalSource } from './auth/synthetic'; import { withSpan, withTraceContext } from './telemetry'; import { workerDeadlineFailure } from './worker-error'; +import { + CLIENT_DISCONNECT_REASON, + JOB_CANCELLED_MESSAGE, + jobResultCommitFailure, + commitJobResult, + claimJobExecution, + jobCancellationRetentionSeconds, + throwIfJobAborted, +} from './job-cancellation'; import logger from './logger'; import { validateQueuedExecutionProfile, validateQueuedSandboxBackend, } from './execution-profile'; +import { + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + programmaticTransferReserveMs, +} from '../../packages/code/src/protocol'; const { INSTANCE_ID } = env; const WORKER_ID = `${INSTANCE_ID}-${process.pid}`; function isAbortError(error: unknown): boolean { - return axios.isAxiosError(error) && (error.name === 'AbortError' || error.code === 'ERR_CANCELED'); + return ( + axios.isAxiosError(error) && + (error.name === 'AbortError' || error.code === 'ERR_CANCELED') + ); } async function processJob(job: t.ExecuteJob): Promise { - return withTraceContext(job.data._otel, () => withSpan('codeapi.job.process', { - 'messaging.system': 'bullmq', - 'messaging.operation.name': 'process', - 'messaging.message.id': typeof job.id === 'string' ? job.id : String(job.id ?? ''), - 'codeapi.language': job.data.payload?.language ?? 'unknown', - 'codeapi.execution_profile': job.data.executionProfile ?? 'legacy', - 'codeapi.worker_execution_profile': env.EXECUTION_PROFILE, - }, () => processJobInner(job), 'CONSUMER')); + return withTraceContext(job.data._otel, () => + withSpan( + 'codeapi.job.process', + { + 'messaging.system': 'bullmq', + 'messaging.operation.name': 'process', + 'messaging.message.id': + typeof job.id === 'string' ? job.id : String(job.id ?? ''), + 'codeapi.language': job.data.payload?.language ?? 'unknown', + 'codeapi.execution_profile': job.data.executionProfile ?? 'legacy', + 'codeapi.worker_execution_profile': env.EXECUTION_PROFILE, + }, + () => processJobInner(job), + 'CONSUMER', + ), + ); } async function processJobInner(job: t.ExecuteJob): Promise { const { payload, isPyPlot } = job.data; - const isSyntheticJob = job.data.isSynthetic === true || isSyntheticPrincipalSource(job.data.principalSource); + const isSyntheticJob = + job.data.isSynthetic === true || + isSyntheticPrincipalSource(job.data.principalSource); const language = payload?.language ?? 'unknown'; const endTimer = jobProcessingDuration.startTimer({ language }); activeJobs.inc({ language }); const controller = new AbortController(); - const deadlineAtMs = jobDeadlineAtMs(job.timestamp, env.JOB_TIMEOUT); + const cancellationTarget = + job.data.cancellable === true && job.id != null + ? { queueName: job.queueName, jobId: String(job.id) } + : undefined; + let cancellationRegistered = false; + const deadlineAtMs = jobDeadlineAtMs( + job.timestamp, + env.JOB_TIMEOUT, + Date.now(), + job.data.deadlineAtMs, + ); const remainingBudgetMs = Math.max(0, deadlineAtMs - Date.now()); - const timer = remainingBudgetMs > 0 - ? setTimeout(() => controller.abort(), remainingBudgetMs) - : undefined; - if (remainingBudgetMs === 0) controller.abort(); + const timer = + remainingBudgetMs > 0 + ? setTimeout(() => controller.abort('deadline'), remainingBudgetMs) + : undefined; + if (remainingBudgetMs === 0) controller.abort('deadline'); let egressGrantId: string | undefined; let egressGrantTokenForRestore: string | undefined; let revokeReason = 'completed'; + let completedResult = false; + let resultToCommit: t.ExecuteResult | undefined; + let resultCommittedAtHandoff = false; + const commitAtHandoff = + cancellationTarget != null && + job.data.workspaceId != null && + env.SANDBOX_BACKEND === 'remote-bridge'; try { + if (cancellationTarget != null) { + await jobCancellationRegistry.register(cancellationTarget, controller); + cancellationRegistered = true; + const claim = await claimJobExecution( + connection, + cancellationTarget, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + ); + if (claim.status === 'completed') return claim.result; + } if (controller.signal.aborted) { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); } - validateQueuedExecutionProfile(job.data.executionProfile, env.EXECUTION_PROFILE); + validateQueuedExecutionProfile( + job.data.executionProfile, + env.EXECUTION_PROFILE, + ); validateQueuedSandboxBackend( job.data.sandboxBackend, env.SANDBOX_BACKEND, @@ -76,7 +156,10 @@ async function processJobInner(job: t.ExecuteJob): Promise { const nowSeconds = Math.floor(Date.now() / 1000); const prepared = await createGatewayEgressGrant({ payload, - claims: refreshEgressGrantClaims(job.data.egressGrantClaims, nowSeconds), + claims: refreshEgressGrantClaims( + job.data.egressGrantClaims, + nowSeconds, + ), isSynthetic: isSyntheticJob, signal: controller.signal, }); @@ -84,16 +167,30 @@ async function processJobInner(job: t.ExecuteJob): Promise { sandboxPayload = prepared.payload; egressGrantToken = prepared.egressGrantToken; egressGrantTokenForRestore = prepared.egressGrantToken; - executionManifestClaims = (env.EXECUTION_MANIFEST_PRIVATE_KEY || env.EXECUTION_MANIFEST_SECRET) - ? prepared.executionManifestClaims - : undefined; + executionManifestClaims = + env.EXECUTION_MANIFEST_PRIVATE_KEY || env.EXECUTION_MANIFEST_SECRET + ? prepared.executionManifestClaims + : undefined; } const delivery = prepareInputDelivery(payload, sandboxPayload); const sandboxRequest = buildSandboxExecuteRequest({ + ...(job.data.workspaceId == null + ? {} + : { + programmaticTransferReserveMs: programmaticTransferReserveMs( + env.JOB_TIMEOUT, + ), + }), payload: delivery.payload, egressGrantToken, executionManifestClaims, + maxOutputFileBytes: Math.min( + executionManifestClaims?.max_upload_bytes ?? + env.EGRESS_GATEWAY_MAX_FILE_BYTES, + env.EGRESS_GATEWAY_MAX_FILE_BYTES, + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + ), executionManifestPrivateKey: env.EXECUTION_MANIFEST_PRIVATE_KEY, executionManifestSecret: env.EXECUTION_MANIFEST_SECRET, executionManifestTtlSeconds: env.EXECUTION_MANIFEST_TTL_SECONDS, @@ -114,21 +211,41 @@ async function processJobInner(job: t.ExecuteJob): Promise { * the transformed object makes that second call an idempotent no-op. */ const resultRestoreToken = egressGrantTokenForRestore; const finalizedSandboxResults = new WeakSet(); - const finalizeSandboxResult = async (result: SandboxRawResponse): Promise => { - if ( - resultRestoreToken === undefined || - resultRestoreToken.length === 0 || - finalizedSandboxResults.has(result) - ) { - return result; + const finalizeSandboxResult = async ( + result: SandboxRawResponse, + ): Promise => { + if (finalizedSandboxResults.has(result)) return result; + const restored = + resultRestoreToken == null || resultRestoreToken.length === 0 + ? result + : await restoreGatewaySandboxResult({ + grantId: egressGrantId, + egressGrantToken: resultRestoreToken, + result, + isSynthetic: isSyntheticJob, + signal: controller.signal, + }); + if (commitAtHandoff && cancellationTarget != null) { + // The bridge still owns its mutation fence here. A failed/ambiguous + // commit quarantines that root before it can serve a caller retry. + throwIfJobAborted(controller.signal); + const mapped = mapSandboxResult(restored); + const committed = await commitJobResult( + connection, + cancellationTarget, + mapped, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + deadlineAtMs, + ); + if (committed === 'cancelled') throw new Error(JOB_CANCELLED_MESSAGE); + if (committed === 'already_completed') + throw new Error('Duplicate mutation handoff; quarantining workspace'); + resultToCommit = mapped; + resultCommittedAtHandoff = true; } - const restored = await restoreGatewaySandboxResult({ - grantId: egressGrantId, - egressGrantToken: resultRestoreToken, - result, - isSynthetic: isSyntheticJob, - signal: controller.signal, - }); finalizedSandboxResults.add(restored); return restored; }; @@ -149,75 +266,124 @@ async function processJobInner(job: t.ExecuteJob): Promise { tenantId: job.data.tenantId, canonicalUserId: job.data.canonicalUserId, bridgeWorkerId: job.data.bridgeWorkerId, + workspaceId: job.data.workspaceId, runtimeSessionId: runtimeSession.runtimeSessionId, runtimeSessionMode: runtimeSession.runtimeSessionMode, /* Stateful backends run this as a commit barrier after user code but * before checkpointing/reusing the mutated workspace. Stateless/HTTP * paths retain the worker-owned fallback immediately below. */ - sessionResultFinalizer: resultRestoreToken !== undefined && resultRestoreToken.length > 0 - ? finalizeSandboxResult - : undefined, + sessionResultFinalizer: + commitAtHandoff || + (resultRestoreToken !== undefined && resultRestoreToken.length > 0) + ? finalizeSandboxResult + : undefined, }, ); const responseData = await finalizeSandboxResult(responseRaw); + // Cancellation can arrive after sandbox exit while artifact restoration + // yields. Do not let BullMQ commit a success after Stop was acknowledged. + if (!resultCommittedAtHandoff) throwIfJobAborted(controller.signal); - if (!isSyntheticJob) { - logger.info('Sandbox response', summarizeSandboxResponse(responseData)); - } + function mapSandboxResult( + responseData: SandboxRawResponse, + ): t.ExecuteResult { + if (!isSyntheticJob) { + logger.info('Sandbox response', summarizeSandboxResponse(responseData)); + } - const { files } = responseData; - const run = responseData.run; - const stdout = applySystemReplacements(run?.stdout ?? ''); - const stderr = filterSystemLogs(run?.stderr ?? '', isPyPlot); + const { files } = responseData; + const run = responseData.run; + const stdout = applySystemReplacements(run?.stdout ?? ''); + const stderr = filterSystemLogs(run?.stderr ?? '', isPyPlot); - const result: t.ExecuteResult = { - session_id: responseData.session_id, - /* `files` is optional on the sandbox response (e.g. dry-run - * execute with no outputs); the public `ExecuteResult.files` is - * required and downstream callers always iterate it. Default to - * `[]` so the strictened response type from Phase B doesn't - * surface a regression that wasn't there before. */ - files: files ?? [], - ...(responseData.artifact_delivery != null - ? { artifact_delivery: responseData.artifact_delivery } - : {}), - stdout, - stderr, - }; + const result: t.ExecuteResult = { + session_id: responseData.session_id, + /* `files` is optional on the sandbox response (e.g. dry-run + * execute with no outputs); the public `ExecuteResult.files` is + * required and downstream callers always iterate it. Default to + * `[]` so the strictened response type from Phase B doesn't + * surface a regression that wasn't there before. */ + files: files ?? [], + ...(responseData.deleted_files != null + ? { deleted_files: responseData.deleted_files } + : {}), + ...(responseData.artifact_delivery != null + ? { artifact_delivery: responseData.artifact_delivery } + : {}), + ...(responseData.artifact_truncation != null + ? { artifact_truncation: responseData.artifact_truncation } + : {}), + stdout, + stderr, + ...(responseData.pending_tool_calls_payload != null + ? { + pending_tool_calls_payload: + responseData.pending_tool_calls_payload, + } + : {}), + }; - if (run) { - result.code = run.code ?? null; - result.signal = run.signal != null ? String(run.signal) : null; - result.message = run.message ?? null; - result.status = run.status ?? null; - result.wall_time = (run as Record).wall_time as number | null ?? null; - } + if (run) { + result.code = run.code ?? null; + result.signal = run.signal != null ? String(run.signal) : null; + result.message = run.message ?? null; + result.status = run.status ?? null; + result.wall_time = + ((run as Record).wall_time as number | null) ?? null; + } - if (result.message || result.signal) { - logger.warn('Sandbox execution error metadata', { - session_id: responseData.session_id, - code: result.code, - signal: result.signal, - message: summarizeText(result.message), - status: result.status, - wall_time: result.wall_time, - }); + if (result.message || result.signal) { + logger.warn('Sandbox execution error metadata', { + session_id: responseData.session_id, + code: result.code, + signal: result.signal, + message: summarizeText(result.message), + status: result.status, + wall_time: result.wall_time, + }); + } + + return result; } + const result = resultToCommit ?? mapSandboxResult(responseData); + completedResult = true; + resultToCommit = result; return result; } catch (error) { - revokeReason = controller.signal.aborted || isAbortError(error) ? 'timeout' : 'failed'; + // Bridge fence cleanup can fail after the outcome was durably committed. + // Preserve the winning result; the bridge retains/quarantines its fence. + if (resultCommittedAtHandoff && resultToCommit != null) + return resultToCommit; + const clientDisconnected = + controller.signal.aborted && + controller.signal.reason === CLIENT_DISCONNECT_REASON; + revokeReason = clientDisconnected + ? 'cancelled' + : controller.signal.aborted || isAbortError(error) + ? 'timeout' + : 'failed'; const errorDetails = getAxiosErrorDetails(error); - logger.error('Error processing job', errorDetails); + if (clientDisconnected) { + logger.info('Job cancelled after client disconnected', { + queueName: job.queueName, + jobId: job.id, + executionId: job.data.executionId, + }); + } else { + logger.error('Error processing job', errorDetails); + } const deadlineFailure = workerDeadlineFailure( error, - controller.signal.aborted, + controller.signal.aborted && !clientDisconnected, env.JOB_TIMEOUT, ); if (deadlineFailure) { throw deadlineFailure; + } else if (clientDisconnected) { + throw new Error(JOB_CANCELLED_MESSAGE); } else if (error instanceof SandboxBackendError) { throw new Error(`${error.code}: ${error.message}`); } else if (error instanceof SessionFilesError) { @@ -237,17 +403,67 @@ async function processJobInner(job: t.ExecuteJob): Promise { if (egressGrantId || egressGrantTokenForRestore) { await revokeGatewayEgressGrant({ grantId: egressGrantId, - egressGrantToken: egressGrantId ? undefined : egressGrantTokenForRestore, + egressGrantToken: egressGrantId + ? undefined + : egressGrantTokenForRestore, isSynthetic: isSyntheticJob, reason: revokeReason, timeoutMs: env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS, }).catch(error => { - logger.error('Failed to revoke egress grant', { grantId: egressGrantId, error: getAxiosErrorDetails(error) }); + logger.error('Failed to revoke egress grant', { + grantId: egressGrantId, + error: getAxiosErrorDetails(error), + }); }); } + let lateCommitFailure = + completedResult && !resultCommittedAtHandoff + ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) + : undefined; + if ( + completedResult && + !resultCommittedAtHandoff && + cancellationTarget != null && + lateCommitFailure == null + ) { + try { + const committed = await commitJobResult( + connection, + cancellationTarget, + resultToCommit, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + deadlineAtMs, + ); + if (committed === 'cancelled') { + lateCommitFailure = new Error(JOB_CANCELLED_MESSAGE); + } else if (committed === 'already_completed') { + lateCommitFailure = new Error( + 'Duplicate result handoff; refusing replacement', + ); + } + } catch (error) { + lateCommitFailure = + error instanceof Error ? error : new Error('Result commit failed'); + } + } if (timer) clearTimeout(timer); + if (cancellationTarget != null && cancellationRegistered) { + await jobCancellationRegistry + .unregister(cancellationTarget, controller) + .catch(error => { + logger.warn('Failed to clear queued execution cancellation state', { + queueName: cancellationTarget.queueName, + jobId: cancellationTarget.jobId, + error: getAxiosErrorDetails(error), + }); + }); + } endTimer(); activeJobs.dec({ language }); + if (lateCommitFailure != null) throw lateCommitFailure; } } @@ -290,21 +506,31 @@ otherWorker.on('completed', job => { }); pyWorker.on('failed', (job, err) => { + if (err.message === JOB_CANCELLED_MESSAGE) { + logger.info(`[${WORKER_ID}] Python job ${job?.id} cancelled`); + jobsCancelled.inc({ language: 'python' }); + return; + } logger.error(`[${WORKER_ID}] Python job ${job?.id} failed`, err); jobsFailed.inc({ language: 'python' }); }); otherWorker.on('failed', (job, err) => { + if (err.message === JOB_CANCELLED_MESSAGE) { + logger.info(`[${WORKER_ID}] Other job ${job?.id} cancelled`); + jobsCancelled.inc({ language: 'other' }); + return; + } logger.error(`[${WORKER_ID}] Other job ${job?.id} failed`, err); jobsFailed.inc({ language: 'other' }); }); -pyWorker.on('error', (err) => { +pyWorker.on('error', err => { logger.error(`[${WORKER_ID}] Python worker error`, err); workerRunning.set({ worker_type: 'python' }, 0); }); -otherWorker.on('error', (err) => { +otherWorker.on('error', err => { logger.error(`[${WORKER_ID}] Other worker error`, err); workerRunning.set({ worker_type: 'other' }, 0); });