diff --git a/sdk/typescript/_bundled_plugin/.codex-plugin/plugin.json b/sdk/typescript/_bundled_plugin/.codex-plugin/plugin.json index bfeb6cd23..e7efa8d68 100644 --- a/sdk/typescript/_bundled_plugin/.codex-plugin/plugin.json +++ b/sdk/typescript/_bundled_plugin/.codex-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "codex-security", - "version": "0.1.60", + "version": "0.1.88", "description": "Codex Security workflows for security scans, analysis, and investigation.", "author": { "name": "OpenAI" diff --git a/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-000 b/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-000 index 8a861abee..d8c349686 100644 Binary files a/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-000 and b/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-000 differ diff --git a/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-001 b/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-001 index 774602a86..7e81673cc 100644 Binary files a/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-001 and b/sdk/typescript/_bundled_plugin/mcp/server.mjs.br.part-001 differ diff --git a/sdk/typescript/_bundled_plugin/scripts/workbench_scan_usage.py b/sdk/typescript/_bundled_plugin/scripts/workbench_scan_usage.py index db987d20c..8e3a63739 100644 --- a/sdk/typescript/_bundled_plugin/scripts/workbench_scan_usage.py +++ b/sdk/typescript/_bundled_plugin/scripts/workbench_scan_usage.py @@ -471,8 +471,9 @@ def _read_rollout_usage( if timestamp is None or snapshot is None: warnings.add("token_record_invalid") continue + reset = snapshot["totalTokens"] < previous["totalTokens"] delta = { - key: value - previous[key] if value >= previous[key] else value + key: value if reset or value < previous[key] else value - previous[key] for key, value in snapshot.items() } previous = snapshot @@ -480,8 +481,19 @@ def _read_rollout_usage( continue if completed_at is not None and timestamp > completed_at: continue + delta["totalTokens"] = delta["inputTokens"] + delta["outputTokens"] if delta["totalTokens"] <= 0: continue + delta["cacheWriteInputTokens"] = min( + delta["cacheWriteInputTokens"], delta["inputTokens"] + ) + delta["cachedInputTokens"] = min( + delta["cachedInputTokens"], + delta["inputTokens"] - delta["cacheWriteInputTokens"], + ) + delta["reasoningOutputTokens"] = min( + delta["reasoningOutputTokens"], delta["outputTokens"] + ) _add_token_usage(total, delta) if not boundary_reached: @@ -516,19 +528,19 @@ def _is_owned_task_start( turn_id = payload.get("turn_id") if not isinstance(turn_id, str) or not turn_id: return False - thread_timestamp = _uuid7_timestamp(thread_id) - turn_timestamp = _uuid7_timestamp(turn_id) - if thread_timestamp is None: + thread_order = _uuid7_order(thread_id) + turn_order = _uuid7_order(turn_id) + if thread_order is None: return True - return turn_timestamp is not None and turn_timestamp >= thread_timestamp + return turn_order is not None and turn_order >= thread_order -def _uuid7_timestamp(value: str) -> int | None: +def _uuid7_order(value: str) -> int | None: try: parsed = uuid.UUID(value) except ValueError: return None - return parsed.int >> 80 if parsed.version == 7 else None + return parsed.int if parsed.version == 7 else None def _token_snapshot(payload: Mapping[str, Any]) -> dict[str, int] | None: diff --git a/sdk/typescript/src/api.ts b/sdk/typescript/src/api.ts index 426452c9b..e6f6f820a 100644 --- a/sdk/typescript/src/api.ts +++ b/sdk/typescript/src/api.ts @@ -125,6 +125,7 @@ import { resolveCodexCommand, resolvePluginPath, resolvePluginPython, + resolveScanSessionPaths, runWorkbench, setCodexSecurityCredentialLogout, type CodexCommand, @@ -367,6 +368,7 @@ interface ClientDependencies { signal?: AbortSignal, ) => Promise; resolvePluginPython?: typeof resolvePluginPython; + resolveScanSessionPaths?: typeof resolveScanSessionPaths; prepareOutputDir?: typeof prepareOutputDir; repositoryRevision?: typeof repositoryRevision; resolveCodexCommand?: () => CodexCommand; @@ -629,6 +631,7 @@ export class CodexSecurity { let scanFailure = false; let customValidationComplete = false; let completionCost: ScanCost | null = null; + let signaledCostUsage: unknown; let budgetRecovery: { expectation: ScanExpectation; pluginRoot: string; @@ -850,6 +853,21 @@ export class CodexSecurity { repository: repo, scanDirectory: scanDir, maxCostUsd: options.maxCostUsd, + resolveOwnedSessionPaths: + options.maxCostUsd === undefined + ? undefined + : async (threadId) => { + const scan = activeScan; + if (scan === null) { + throw new CodexSecurityError( + "The scan session ownership could not be verified.", + ); + } + return await ( + this.#dependencies.resolveScanSessionPaths ?? + resolveScanSessionPaths + )(scan.options, scan.id, threadId); + }, onActivity: options.onActivity === undefined ? undefined @@ -875,7 +893,7 @@ export class CodexSecurity { onCost: options.onCost === undefined && options.maxCostUsd === undefined ? undefined - : (cost) => { + : (cost, usage) => { notifyObserver( "onCost", options.onCost, @@ -884,8 +902,10 @@ export class CodexSecurity { ); if ( options.maxCostUsd !== undefined && - cost.estimatedUsd > options.maxCostUsd + cost.estimatedUsd > options.maxCostUsd && + !costAbortController.signal.aborted ) { + signaledCostUsage = usage; costAbortController.abort( new ScanCostLimitExceededError( options.maxCostUsd, @@ -1185,6 +1205,8 @@ export class CodexSecurity { } }, onFinalize: async (usage) => { + // Validation threads are accounted separately from the discovery turn. + const discoveryUsage = usage; if (options.validationPrompt !== undefined) { await runCustomValidation({ repository: repo, @@ -1231,17 +1253,29 @@ export class CodexSecurity { turn.lastStreamError ?? "The custom validation turn did not complete.", ); + tracker.recordCompletedThreadUsage(turn.threadId, turn.usage); usage = sumTokenUsage(usage, turn.usage); return turn.finalResponse; }, }); customValidationComplete = true; } - const snapshot = await tracker.stop(usage).catch((error: unknown) => { - if (options.maxCostUsd !== undefined) throw error; - reportTrackingError(error); - return { usage, cost: estimateScanCost(model, usage) }; - }); + const snapshot = await tracker + .stop(discoveryUsage) + .catch(async (error: unknown) => { + if (options.maxCostUsd !== undefined) { + throwIfAborted(signal, scanDir); + try { + return await tracker.stop(discoveryUsage); + } catch { + runPostScan = null; + throwIfAborted(signal, scanDir); + throw error; + } + } + reportTrackingError(error); + return { usage, cost: estimateScanCost(model, usage) }; + }); throwIfAborted(signal, scanDir); if (options.maxCostUsd !== undefined && snapshot.cost === null) { notifyObserver( @@ -1470,11 +1504,48 @@ export class CodexSecurity { // Recorded first: everything below can throw a different error for this same failed // scan, and cleanup must treat all of those as a failure it is not allowed to mask. scanFailure = true; - const snapshot = await costTracker?.stop().catch(() => null); - const failure = + const trackedSnapshot = await costTracker?.stop().catch(() => null); + const signaledOverage = signal.reason instanceof ScanCostLimitExceededError ? signal.reason - : error; + : null; + const trackedCost = trackedSnapshot?.cost; + const snapshot = + signaledOverage !== null && + (trackedCost === undefined || + trackedCost === null || + signaledOverage.cost.estimatedUsd > trackedCost.estimatedUsd) + ? { + cost: signaledOverage.cost, + usage: signaledCostUsage ?? { + input_tokens: signaledOverage.cost.inputTokens, + cached_input_tokens: signaledOverage.cost.cachedInputTokens, + cache_write_input_tokens: + signaledOverage.cost.cacheWriteInputTokens, + output_tokens: signaledOverage.cost.outputTokens, + reasoning_output_tokens: 0, + }, + } + : trackedSnapshot; + const knownCost = snapshot?.cost; + let failure: unknown = signaledOverage ?? error; + if ( + options.maxCostUsd !== undefined && + knownCost !== undefined && + knownCost !== null && + knownCost.estimatedUsd > options.maxCostUsd && + (signaledOverage !== null || + (!this.#abortController.signal.aborted && + options.signal?.aborted !== true)) && + (signaledOverage === null || + knownCost.estimatedUsd > signaledOverage.cost.estimatedUsd) + ) { + failure = new ScanCostLimitExceededError( + options.maxCostUsd, + knownCost, + scanDir, + ); + } if ( failure instanceof ScanCostLimitExceededError && budgetRecovery !== null && diff --git a/sdk/typescript/src/cost.ts b/sdk/typescript/src/cost.ts index 05b72254f..abf1bf151 100644 --- a/sdk/typescript/src/cost.ts +++ b/sdk/typescript/src/cost.ts @@ -1,4 +1,4 @@ -import { open, readdir } from "node:fs/promises"; +import { open, readdir, realpath } from "node:fs/promises"; import { join } from "node:path"; import { estimateScanCost, @@ -40,14 +40,18 @@ interface SessionUsage { offset: number; pendingLine: Buffer[]; pendingLineBytes: number; - unreadable: boolean; + unreadable: { error: unknown } | null; threadId: string | null; parentThreadId: string | null; workingDirectory: string | null; startedAt: number | null; inheritedUsage: ScanTokenUsage | null; + previousRawUsage: ScanTokenUsage | null; + accumulatedOwnUsage: ScanTokenUsage | null; replaying: boolean; - usage: ScanTokenUsage | null; + accounting: { usage: ScanTokenUsage; cost: ScanCost | null } | null; + accountingError: Error | null; + taskCompleted: boolean; calls: Map; activities: ScanActivity[]; progress: ScanProgress[]; @@ -66,11 +70,14 @@ interface ScanCostTrackerOptions { scanDirectory?: string; maxCostUsd?: number; expectedFilesTotal?: number; - onCost?: (cost: Readonly) => void; + onCost?: (cost: Readonly, usage: unknown) => void; onActivity?: (activity: ScanActivity) => void; onProgress?: (progress: ScanProgress) => void; onSessionEvent?: (event: ScanSessionEvent) => void; onError?: (error: unknown) => void; + resolveOwnedSessionPaths?: ( + rootThreadId: string, + ) => Promise>; } interface ScanCostSnapshot { @@ -78,6 +85,16 @@ interface ScanCostSnapshot { cost: ScanCost | null; } +interface ObservedSessionUsage { + root: ScanTokenUsage | null; + workers: ScanTokenUsage | null; + completedRoot: ScanTokenUsage | null; + rootCompleted: boolean; + unverified: boolean; + unfinishedWorkers: boolean; + unidentifiedSessions: Set; + accountedSessions: Set; +} const COST_POLL_INTERVAL_MS = 100; const SESSION_READ_SIZE = 64 * 1_024; @@ -86,14 +103,18 @@ function createSessionUsage(): SessionUsage { offset: 0, pendingLine: [], pendingLineBytes: 0, - unreadable: false, + unreadable: null, threadId: null, parentThreadId: null, workingDirectory: null, startedAt: null, inheritedUsage: null, + previousRawUsage: null, + accumulatedOwnUsage: null, replaying: false, - usage: null, + accounting: null, + accountingError: null, + taskCompleted: false, calls: new Map(), activities: [], progress: [], @@ -115,7 +136,20 @@ export class ScanCostTracker { #timer: NodeJS.Timeout | null = null; #pending: Promise = Promise.resolve(); #snapshot: ScanCostSnapshot = { usage: null, cost: null }; + #finalSnapshot: ScanCostSnapshot | null = null; + #completedThreadUsage = new Map(); + #observedUsage: ObservedSessionUsage = { + root: null, + workers: null, + completedRoot: null, + rootCompleted: false, + unverified: false, + unfinishedWorkers: false, + unidentifiedSessions: new Set(), + accountedSessions: new Set(), + }; #lastCost: number | null = null; + #rootOnlyReadError = false; #highestFilesCompleted = 0; #expectedFilesTotal: number | undefined; @@ -128,6 +162,13 @@ export class ScanCostTracker { this.#expectedFilesTotal = filesTotal; } + public recordCompletedThreadUsage( + threadId: string | null, + usage: unknown, + ): void { + this.#completedThreadUsage.set(threadId, tokenUsage(usage)); + } + public start(threadId: string): void { if (this.#threadId !== null) return; this.#threadId = threadId; @@ -175,20 +216,114 @@ export class ScanCostTracker { } public async stop(fallbackUsage?: unknown): Promise { + const finalizing = arguments.length > 0; + if (this.#finalSnapshot !== null) return this.#finalSnapshot; if (this.#timer !== null) { clearInterval(this.#timer); this.#timer = null; } - await this.refresh(); - if (this.#snapshot.usage !== null) return this.#snapshot; - const cost = estimateScanCost(this.#options.model, fallbackUsage); - this.#snapshot = { usage: fallbackUsage ?? null, cost }; - this.#reportCost(cost); - return this.#snapshot; + const suppliedRoot = tokenUsage(fallbackUsage); + let refreshFailure: { error: unknown } | null = null; + try { + await this.refresh(); + } catch (error) { + refreshFailure = { error }; + } + const observed = this.#observedUsage; + const completedRoot = higherCostUsage( + this.#options.model, + observed.completedRoot, + suppliedRoot, + ); + observed.completedRoot = completedRoot; + const rootUsage = higherCostUsage( + this.#options.model, + observed.root, + completedRoot, + ); + let completedUsage: unknown = + rootUsage === suppliedRoot ? fallbackUsage : rootUsage; + const workerUsage = observed.workers; + if (workerUsage !== null) { + completedUsage = addTokenUsage(rootUsage, workerUsage); + } + const cost = estimateScanCost(this.#options.model, completedUsage); + const snapshot = + this.#snapshot.usage !== null && + ((rootUsage === null && workerUsage === null) || + (this.#snapshot.cost !== null && + (cost === null || + this.#snapshot.cost.estimatedUsd > cost.estimatedUsd))) + ? this.#snapshot + : { usage: completedUsage ?? null, cost }; + this.#snapshot = snapshot; + if ( + this.#options.maxCostUsd !== undefined && + snapshot.cost !== null && + snapshot.cost.estimatedUsd > this.#options.maxCostUsd + ) { + this.#reportCost(snapshot.cost); + if (!finalizing) return snapshot; + } + if (refreshFailure !== null) { + if ( + fallbackUsage === undefined || + (this.#options.maxCostUsd !== undefined && + (completedRoot === null || !this.#rootOnlyReadError)) + ) { + throw refreshFailure.error; + } + if (this.#options.maxCostUsd === undefined) { + this.#options.onError?.(refreshFailure.error); + } + } + let unidentifiedOwnedSession = false; + if (finalizing && this.#options.maxCostUsd !== undefined) { + const resolveOwnedSessionPaths = this.#options.resolveOwnedSessionPaths; + if (resolveOwnedSessionPaths === undefined || this.#threadId === null) { + unidentifiedOwnedSession = observed.unidentifiedSessions.size > 0; + } else { + const ownedPaths = await resolveOwnedSessionPaths(this.#threadId); + const accountedPaths = new Set( + await Promise.all( + [...observed.accountedSessions].map(async (path) => realpath(path)), + ), + ); + for (const path of ownedPaths) { + if (!accountedPaths.has(path)) { + unidentifiedOwnedSession = true; + break; + } + } + } + } + if ( + this.#options.maxCostUsd !== undefined && + (rootUsage === null || + cost === null || + observed.unverified || + unidentifiedOwnedSession || + (finalizing && + ((completedRoot === null && !observed.rootCompleted) || + observed.unfinishedWorkers))) + ) { + throw ( + refreshFailure?.error ?? + new Error( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ) + ); + } + if (finalizing) this.#finalSnapshot = snapshot; + this.#reportCost(snapshot.cost); + return snapshot; } async #readSessions(): Promise { - if (this.#threadId === null) return; + const rootThreadId = this.#threadId; + if (rootThreadId === null) return; + this.#rootOnlyReadError = false; + const presentSessions = new Set(); const unreadable: Array<{ session: SessionUsage; error: unknown }> = []; for await (const path of sessionFiles( join(this.#options.codexHome, "sessions"), @@ -199,38 +334,94 @@ export class ScanCostTracker { this.#sessions.set(path, session); } try { - await readSessionUsage(path, session, this.#options.repository); + presentSessions.add(path); + if ( + !(await readSessionUsage( + path, + session, + this.#options.model, + this.#options.repository, + this.#options.maxCostUsd !== undefined, + )) + ) { + presentSessions.delete(path); + } } catch (error) { if (session.threadId === null) throw error; unreadable.push({ session, error }); } } - const included = new Set([this.#threadId]); + const parents = new Map(); + const conflictingParents = new Set(); + for (const session of this.#sessions.values()) { + if (session.threadId === null) continue; + const previous = parents.get(session.threadId); + if (previous !== undefined && previous !== session.parentThreadId) { + conflictingParents.add(session.threadId); + } else { + parents.set(session.threadId, session.parentThreadId); + } + } + const knownOwner = (threadId: string): string | null => { + const seen = new Set(); + while (threadId !== rootThreadId) { + if (seen.has(threadId) || conflictingParents.has(threadId)) return null; + seen.add(threadId); + const parent = parents.get(threadId); + if (parent === undefined) return null; + if (parent === null) return threadId; + threadId = parent; + } + return rootThreadId; + }; + + const included = new Set([ + rootThreadId, + ...[...this.#completedThreadUsage.keys()].filter( + (threadId): threadId is string => threadId !== null, + ), + ]); + const ambiguousWorkers = new Set(); if (this.#options.scanDirectory !== undefined) { const scanStartedAt = [...this.#sessions.values()].find( - (session) => session.threadId === this.#threadId, + (session) => session.threadId === rootThreadId, )?.startedAt ?? null; for (const session of this.#sessions.values()) { if ( session.threadId === null || - session.parentThreadId !== null || - session.workingDirectory === null || - scanStartedAt === null || - session.startedAt === null || - session.startedAt < scanStartedAt + session.threadId === rootThreadId || + session.workingDirectory === null ) { continue; } if ( - isScanArtifactDirectory( + !isScanArtifactDirectory( this.#options.scanDirectory, session.workingDirectory, ) ) { - included.add(session.threadId); + continue; } + if ( + scanStartedAt !== null && + session.startedAt !== null && + session.startedAt < scanStartedAt + ) { + continue; + } + const owner = knownOwner(session.threadId); + if (owner === null) { + ambiguousWorkers.add(session.threadId); + continue; + } + if (owner !== session.threadId) continue; + if (scanStartedAt === null || session.startedAt === null) { + ambiguousWorkers.add(session.threadId); + continue; + } + included.add(session.threadId); } } let changed = true; @@ -243,61 +434,195 @@ export class ScanCostTracker { included.has(session.parentThreadId) && !included.has(session.threadId) ) { + if (conflictingParents.has(session.threadId)) { + ambiguousWorkers.add(session.threadId); + continue; + } included.add(session.threadId); changed = true; } } } + const hasUnverifiedWorkerAttribution = [...ambiguousWorkers].some( + (threadId) => !included.has(threadId), + ); + const readFailures: Array<{ + session: SessionUsage; + error: unknown; + rootOnlyRecoverable: boolean; + }> = []; + if (this.#options.maxCostUsd !== undefined) { + for (const [path, session] of this.#sessions) { + if ( + session.threadId !== null && + included.has(session.threadId) && + !presentSessions.has(path) + ) { + readFailures.push({ + session, + error: new Error( + "A tracked scan session disappeared before its cost could be verified.", + ), + rootOnlyRecoverable: false, + }); + } + } + } for (const { session, error } of unreadable) { - if (included.has(session.threadId!)) throw error; + if (included.has(session.threadId!)) { + readFailures.push({ session, error, rootOnlyRecoverable: true }); + } else if (isSessionAccessDenied(error)) { + quarantineSession(session, error); + } } - let usage: ScanTokenUsage | null = null; - for (const [path, tracked] of this.#sessions) { - const threadId = tracked.threadId; - if (threadId === null || !included.has(threadId)) continue; - let session = tracked; + const observed: ObservedSessionUsage = { + root: null, + workers: null, + completedRoot: this.#observedUsage.completedRoot, + rootCompleted: false, + unverified: hasUnverifiedWorkerAttribution, + unfinishedWorkers: false, + unidentifiedSessions: new Set(), + accountedSessions: new Set(), + }; + const accountedThreads = new Set(); + for (const [path, session] of this.#sessions) { if ( - this.#options.onSessionEvent !== undefined && - session.events === undefined + this.#options.maxCostUsd !== undefined && + session.threadId === null && + presentSessions.has(path) ) { - // Replay only newly associated sessions, including their early events. - session = createSessionUsage(); - session.events = []; - await readSessionUsage(path, session, this.#options.repository); - this.#sessions.set(path, session); - } - let worker: number | undefined; - if (threadId !== this.#threadId) { - worker = this.#workers.get(threadId) ?? this.#workers.size + 1; - this.#workers.set(threadId, worker); + observed.unidentifiedSessions.add(path); } - for (const event of session.events?.splice(0) ?? []) { - this.#options.onSessionEvent?.({ - threadId, - parentThreadId: session.parentThreadId, - worker, - event, - }); - } - if (worker !== undefined) { - for (const activity of session.activities.splice(0)) { - this.#options.onActivity?.({ - ...activity, - id: `${threadId}:${activity.id}`, - worker, + if (session.threadId !== null && included.has(session.threadId)) { + accountedThreads.add(session.threadId); + if (presentSessions.has(path)) observed.accountedSessions.add(path); + await this.#reportSessionEvents(path, session); + if ( + this.#options.maxCostUsd !== undefined && + session.accountingError !== null + ) { + readFailures.push({ + session, + error: session.accountingError, + rootOnlyRecoverable: true, }); } - this.#reportWorkerProgress(session); + if (session.pendingLineBytes > 0) observed.unverified = true; + const usage = higherCostUsage( + this.#options.model, + session.accounting?.usage ?? null, + this.#completedThreadUsage.get(session.threadId) ?? null, + ); + if (session.threadId === this.#threadId) { + observed.rootCompleted = + session.taskCompleted || + this.#completedThreadUsage.has(session.threadId); + if (usage !== null) { + observed.root = addTokenUsage(observed.root, usage); + } + } else { + if (usage === null) { + observed.unverified = true; + } else { + observed.workers = addTokenUsage(observed.workers, usage); + } + if ( + !session.taskCompleted && + !this.#completedThreadUsage.has(session.threadId) + ) + observed.unfinishedWorkers = true; + } } - if (session.usage !== null) { - usage = addTokenUsage(usage, session.usage); + } + for (const [threadId, usage] of this.#completedThreadUsage) { + if (accountedThreads.has(threadId)) continue; + if (usage === null) { + observed.unverified = true; + } else if (threadId === rootThreadId) { + observed.root = addTokenUsage(observed.root, usage); + observed.rootCompleted = true; + } else { + observed.workers = addTokenUsage(observed.workers, usage); } } - if (usage === null) return; - const cost = estimateScanCost(this.#options.model, usage); - this.#snapshot = { usage, cost }; - this.#reportCost(cost); + this.#observedUsage = observed; + const usage = + observed.workers === null + ? observed.root + : addTokenUsage(observed.root, observed.workers); + if (usage !== null) { + const cost = estimateScanCost(this.#options.model, usage); + if ( + this.#snapshot.cost === null || + (cost !== null && cost.estimatedUsd >= this.#snapshot.cost.estimatedUsd) + ) { + this.#snapshot = { usage, cost }; + } + this.#reportCost(this.#snapshot.cost); + } + const readFailure = readFailures[0]; + if (readFailure !== undefined) { + this.#rootOnlyReadError = + readFailures.every( + (failure) => + failure.rootOnlyRecoverable && + failure.session.threadId === rootThreadId, + ) && !hasUnverifiedWorkerAttribution; + throw readFailure.error; + } + } + + async #reportSessionEvents( + path: string, + session: SessionUsage, + ): Promise { + const threadId = session.threadId; + if (threadId === null) return; + if ( + this.#options.onSessionEvent !== undefined && + session.events === undefined + ) { + const replay = createSessionUsage(); + replay.events = []; + try { + // Replay only bytes already accounted for, without replacing cost state. + await readSessionUsage( + path, + replay, + this.#options.model, + this.#options.repository, + false, + session.offset, + ); + } catch { + // Detail replay is optional. The accounting reader retains its own errors. + } + session.events = replay.threadId === threadId ? replay.events : []; + } + let worker: number | undefined; + if (threadId !== this.#threadId) { + worker = this.#workers.get(threadId) ?? this.#workers.size + 1; + this.#workers.set(threadId, worker); + } + for (const event of session.events?.splice(0) ?? []) { + this.#options.onSessionEvent?.({ + threadId, + parentThreadId: session.parentThreadId, + worker, + event, + }); + } + if (worker === undefined) return; + for (const activity of session.activities.splice(0)) { + this.#options.onActivity?.({ + ...activity, + id: `${threadId}:${activity.id}`, + worker, + }); + } + this.#reportWorkerProgress(session); } #reportWorkerProgress(session: SessionUsage): void { @@ -343,7 +668,7 @@ export class ScanCostTracker { #reportCost(cost: ScanCost | null): void { if (cost === null || cost.estimatedUsd === this.#lastCost) return; this.#lastCost = cost.estimatedUsd; - this.#options.onCost?.(cost); + this.#options.onCost?.(cost, this.#snapshot.usage); } } @@ -368,33 +693,45 @@ export async function* sessionFiles(directory: string): AsyncGenerator { async function readSessionUsage( path: string, session: SessionUsage, + model: string, repository?: string, -): Promise { - if (session.unreadable) return; + requireReadableSessions = false, + endOffset?: number, +): Promise { + if (session.unreadable !== null) { + if (requireReadableSessions) throw session.unreadable.error; + return true; + } let file; try { file = await open(path, "r"); } catch (error) { - if (isMissingFile(error)) return; + if (isMissingFile(error)) return false; + if (session.threadId === null && isSessionAccessDenied(error)) { + quarantineSession(session, error); + } throw error; } try { const buffer = Buffer.alloc(SESSION_READ_SIZE); while (true) { - const { bytesRead } = await file.read( - buffer, - 0, - buffer.length, - session.offset, - ); - if (bytesRead === 0) return; + const length = + endOffset === undefined + ? buffer.length + : Math.min(buffer.length, endOffset - session.offset); + if (length <= 0) return true; + const { bytesRead } = await file.read(buffer, 0, length, session.offset); + if (bytesRead === 0) return true; session.offset += bytesRead; try { - readSessionChunk(buffer.subarray(0, bytesRead), session, repository); + readSessionChunk( + buffer.subarray(0, bytesRead), + session, + model, + repository, + ); } catch (error) { - session.unreadable = true; - session.pendingLine = []; - session.pendingLineBytes = 0; + quarantineSession(session, error); throw error; } } @@ -403,9 +740,16 @@ async function readSessionUsage( } } +function quarantineSession(session: SessionUsage, error: unknown): void { + session.unreadable = { error }; + session.pendingLine = []; + session.pendingLineBytes = 0; +} + function readSessionChunk( contents: Buffer, session: SessionUsage, + model: string, repository?: string, ): void { let lineStart = 0; @@ -424,12 +768,13 @@ function readSessionChunk( } if (session.pendingLineBytes === 0) { - readSessionEvent(fragment.toString("utf8"), session, repository); + readSessionEvent(fragment.toString("utf8"), session, model, repository); } else { if (fragment.length > 0) session.pendingLine.push(Buffer.from(fragment)); readSessionEvent( Buffer.concat(session.pendingLine, lineBytes).toString("utf8"), session, + model, repository, ); session.pendingLine = []; @@ -439,9 +784,22 @@ function readSessionChunk( } } +function uuid7Order(value: unknown): bigint | null { + if ( + typeof value !== "string" || + !/^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/iu.test( + value, + ) + ) { + return null; + } + return BigInt(`0x${value.replaceAll("-", "")}`); +} + function readSessionEvent( line: string, session: SessionUsage, + model: string, repository?: string, ): void { if (line.length === 0) return; @@ -449,6 +807,11 @@ function readSessionEvent( try { event = JSON.parse(line) as unknown; } catch { + if (!session.replaying) { + session.accountingError ??= new Error( + "The scan cost limit could not be verified because a tracked session record could not be read.", + ); + } return; } if (!isRecord(event) || !isRecord(event["payload"])) return; @@ -456,6 +819,7 @@ function readSessionEvent( if (event["type"] === "session_meta") { if (session.threadId !== null) { session.replaying = payload["id"] !== session.threadId; + session.taskCompleted = false; if (!session.replaying) session.events?.push(event); return; } @@ -467,6 +831,8 @@ function readSessionEvent( } session.startedAt = sessionStartedAt(payload["timestamp"]); session.parentThreadId = sessionParentThreadId(payload); + const forkedFrom = payload["forked_from_id"]; + session.replaying = typeof forkedFrom === "string" && forkedFrom.length > 0; session.events?.push(event); return; } @@ -474,20 +840,43 @@ function readSessionEvent( if (event["type"] !== "event_msg") return; if (payload["type"] === "token_count" && isRecord(payload["info"])) { const usage = tokenUsage(payload["info"]["total_token_usage"]); - if (usage !== null) session.inheritedUsage = usage; + if (usage !== null) { + session.inheritedUsage = usage; + session.previousRawUsage = usage; + } } - if ( - payload["type"] === "task_started" && - typeof payload["started_at"] === "number" && - session.startedAt !== null && - payload["started_at"] >= Math.floor(session.startedAt / 1_000) - ) { - session.replaying = false; - session.events?.push(event); + if (payload["type"] === "task_started") { + const threadOrder = uuid7Order(session.threadId); + const turnOrder = uuid7Order(payload["turn_id"]); + const owned = + threadOrder === null + ? typeof payload["started_at"] === "number" && + session.startedAt !== null && + payload["started_at"] >= Math.floor(session.startedAt / 1_000) + : turnOrder !== null && turnOrder >= threadOrder; + if (owned) { + session.replaying = false; + session.taskCompleted = false; + session.events?.push(event); + } } return; } session.events?.push(event); + if (event["type"] === "event_msg") { + if (payload["type"] === "task_started") { + session.taskCompleted = false; + return; + } + if ( + payload["type"] === "task_complete" || + payload["type"] === "turn_complete" || + payload["type"] === "turn_aborted" + ) { + session.taskCompleted = true; + return; + } + } if (event["type"] === "response_item") { session.progress.push(...sessionProgressUpdates(payload)); if (repository === undefined) return; @@ -620,12 +1009,56 @@ function readSessionEvent( return; } const usage = tokenUsage(payload["info"]["total_token_usage"]); - if (usage === null) return; + const accumulated = + usage === null + ? null + : accumulateTokenUsage( + session.previousRawUsage, + session.accumulatedOwnUsage, + usage, + ); const ownUsage = - session.inheritedUsage === null - ? usage - : subtractTokenUsage(usage, session.inheritedUsage); - if (ownUsage !== null) session.usage = ownUsage; + usage === null + ? null + : session.inheritedUsage === null + ? usage + : subtractTokenUsage(usage, session.inheritedUsage); + const cost = ownUsage === null ? null : estimateScanCost(model, ownUsage); + const accumulatedUsage = accumulated ?? session.accumulatedOwnUsage; + const accumulatedCost = estimateScanCost(model, accumulatedUsage); + if ( + accumulated === null || + accumulatedCost === null || + (ownUsage !== null && cost === null) + ) { + session.accountingError ??= new Error( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + } + if (usage === null) return; + session.previousRawUsage = usage; + if (accumulated !== null) session.accumulatedOwnUsage = accumulated; + for (const candidate of [ + ownUsage === null ? null : { usage: ownUsage, cost }, + accumulatedUsage === null + ? null + : { usage: accumulatedUsage, cost: accumulatedCost }, + ]) { + if (candidate === null) continue; + const previous = session.accounting; + if ( + previous === null || + (candidate.cost !== null + ? previous.cost === null || + candidate.cost.estimatedUsd >= previous.cost.estimatedUsd + : previous.cost === null && + higherCostUsage(model, previous.usage, candidate.usage) === + candidate.usage) + ) { + session.accounting = candidate; + } + } + session.taskCompleted = false; } function readSessionReasoning( @@ -749,6 +1182,82 @@ function sessionContentText( .join("\n"); } +function higherCostUsage( + model: string, + previous: ScanTokenUsage | null, + next: ScanTokenUsage | null, +): ScanTokenUsage | null { + if (next === null) return previous; + const previousCost = estimateScanCost(model, previous); + const nextCost = estimateScanCost(model, next); + if (previous !== null && previousCost === null && nextCost === null) { + const previousTotal = + BigInt(previous.input_tokens) + BigInt(previous.output_tokens); + const nextTotal = BigInt(next.input_tokens) + BigInt(next.output_tokens); + return previousTotal > nextTotal ? previous : next; + } + return previousCost !== null && + nextCost !== null && + previousCost.estimatedUsd > nextCost.estimatedUsd + ? previous + : next; +} + +function accumulateTokenUsage( + previousRaw: ScanTokenUsage | null, + accumulated: ScanTokenUsage | null, + next: ScanTokenUsage, +): ScanTokenUsage | null { + const previousTotal = + BigInt(previousRaw?.input_tokens ?? 0) + + BigInt(previousRaw?.output_tokens ?? 0); + const nextTotal = BigInt(next.input_tokens) + BigInt(next.output_tokens); + const reset = nextTotal < previousTotal; + if ( + next.input_tokens === (previousRaw?.input_tokens ?? 0) && + next.output_tokens === (previousRaw?.output_tokens ?? 0) + ) { + return accumulated ?? tokenUsage({ input_tokens: 0, output_tokens: 0 }); + } + type TokenField = Exclude; + const fieldDelta = (field: TokenField): bigint => { + const previous = BigInt(previousRaw?.[field] ?? 0); + const value = BigInt(next[field]); + return reset || value < previous ? value : value - previous; + }; + const inputDelta = fieldDelta("input_tokens"); + const outputDelta = fieldDelta("output_tokens"); + const cacheWriteRaw = fieldDelta("cache_write_input_tokens"); + const cacheWriteDelta = + cacheWriteRaw > inputDelta ? inputDelta : cacheWriteRaw; + const remainingInputDelta = inputDelta - cacheWriteDelta; + const cachedRaw = fieldDelta("cached_input_tokens"); + const cachedDelta = + cachedRaw > remainingInputDelta ? remainingInputDelta : cachedRaw; + const reasoningRaw = fieldDelta("reasoning_output_tokens"); + const reasoningDelta = + reasoningRaw > outputDelta ? outputDelta : reasoningRaw; + const addDelta = (field: TokenField, delta: bigint): number | null => { + const total = BigInt(accumulated?.[field] ?? 0) + delta; + return total <= BigInt(Number.MAX_SAFE_INTEGER) ? Number(total) : null; + }; + const usage = { + input_tokens: addDelta("input_tokens", inputDelta), + cached_input_tokens: addDelta("cached_input_tokens", cachedDelta), + cache_write_input_tokens: addDelta( + "cache_write_input_tokens", + cacheWriteDelta, + ), + output_tokens: addDelta("output_tokens", outputDelta), + reasoning_output_tokens: addDelta( + "reasoning_output_tokens", + reasoningDelta, + ), + }; + return Object.values(usage).some((value) => value === null) + ? null + : tokenUsage(usage); +} function addTokenUsage( previous: ScanTokenUsage | null, next: ScanTokenUsage, @@ -797,3 +1306,9 @@ function isRecord(value: unknown): value is Record { function isMissingFile(error: unknown): boolean { return isRecord(error) && error["code"] === "ENOENT"; } + +function isSessionAccessDenied(error: unknown): boolean { + return ( + isRecord(error) && (error["code"] === "EACCES" || error["code"] === "EPERM") + ); +} diff --git a/sdk/typescript/src/runtime.ts b/sdk/typescript/src/runtime.ts index 17cc88dc7..d784258f4 100644 --- a/sdk/typescript/src/runtime.ts +++ b/sdk/typescript/src/runtime.ts @@ -1483,6 +1483,84 @@ export async function preparePersistentOutputRoot( return root; } +function workbenchEnvironment(options: WorkbenchCommandOptions) { + return Object.fromEntries( + Object.entries(options.environment).filter( + ([name]) => + name.toUpperCase() !== "OPENAI_API_KEY" && + name.toUpperCase() !== "CODEX_API_KEY" && + name.toUpperCase() !== "OPENROUTER_API_KEY" && + name.toUpperCase() !== "FIREWORKS_API_KEY", + ), + ); +} + +/** @internal */ +export async function resolveScanSessionPaths( + options: WorkbenchCommandOptions, + scanId: string, + rootThreadId: string, +): Promise> { + const source = [ + "import json, os, sqlite3, sys", + "from pathlib import Path", + "sys.path.insert(0, sys.argv[1])", + "import workbench_scan_usage as usage", + "path = (Path(os.environ['CODEX_SECURITY_STATE_DIR']) / 'workbench.sqlite3').expanduser().resolve()", + "connection = sqlite3.connect(path.as_uri() + '?mode=ro', uri=True, timeout=1)", + "try:", + " connection.row_factory = sqlite3.Row", + " connection.execute('PRAGMA query_only = ON')", + " scan = connection.execute('SELECT * FROM scans WHERE id = ?', (sys.argv[2],)).fetchone()", + " if scan is None: raise RuntimeError('Scan ownership is unavailable.')", + " roots = usage._scan_root_thread_ids(connection, scan, sys.argv[3])", + "finally:", + " connection.close()", + "database = usage._codex_state_database()", + "if database is None: raise RuntimeError('Codex session ownership is unavailable.')", + "warnings = set()", + "sessions, missing = usage._discover_rollout_sessions(database, roots, warnings)", + "if missing or warnings or not any(session.thread_id == sys.argv[3] for session in sessions):", + " raise RuntimeError('Scan session ownership is incomplete.')", + "print(json.dumps([str(session.path) for session in sessions], allow_nan=False))", + ].join("\n"); + try { + const { stdout } = await execFile( + options.python, + [ + "-I", + "-B", + "-c", + source, + join(options.pluginRoot, "scripts"), + scanId, + rootThreadId, + ], + { + env: workbenchEnvironment(options), + encoding: "utf8", + maxBuffer: Infinity, + windowsHide: true, + signal: options.signal, + }, + ); + const paths: unknown = JSON.parse(stdout); + if ( + !Array.isArray(paths) || + !paths.every((path) => typeof path === "string" && isAbsolute(path)) + ) { + throw new Error("The scan session ownership response is invalid."); + } + return new Set(paths); + } catch (error) { + if (options.signal?.aborted) throw error; + throw new CodexSecurityError( + "The scan session ownership could not be verified.", + { cause: error }, + ); + } +} + export async function runWorkbench( options: WorkbenchCommandOptions, args: readonly string[], @@ -1490,15 +1568,6 @@ export async function runWorkbench( ): Promise { let stdout: string; try { - const environment = Object.fromEntries( - Object.entries(options.environment).filter( - ([name]) => - name.toUpperCase() !== "OPENAI_API_KEY" && - name.toUpperCase() !== "CODEX_API_KEY" && - name.toUpperCase() !== "OPENROUTER_API_KEY" && - name.toUpperCase() !== "FIREWORKS_API_KEY", - ), - ); const result = await runCodexCommand( { command: options.python }, [ @@ -1509,7 +1578,7 @@ export async function runWorkbench( join(options.pluginRoot, "scripts", "workbench_db.py"), ...args, ], - pythonUtf8Environment(environment), + pythonUtf8Environment(workbenchEnvironment(options)), input, options.signal, ); diff --git a/sdk/typescript/src/version.ts b/sdk/typescript/src/version.ts index 95861c52e..450e28021 100644 --- a/sdk/typescript/src/version.ts +++ b/sdk/typescript/src/version.ts @@ -9,7 +9,7 @@ const PACKAGE_VERSIONS = packageVersions( export const VERSION = PACKAGE_VERSIONS.package; export const CODEX_SDK_VERSION = PACKAGE_VERSIONS.sdk; export const CODEX_EXECUTABLE_VERSION = PACKAGE_VERSIONS.executable; -export const BUNDLED_PLUGIN_VERSION = "0.1.60" as const; +export const BUNDLED_PLUGIN_VERSION = "0.1.88" as const; const PACKAGE_NAME = "@openai/codex-security"; diff --git a/sdk/typescript/tests-ts/api.test.ts b/sdk/typescript/tests-ts/api.test.ts index 74836e58d..5252fd5ce 100644 --- a/sdk/typescript/tests-ts/api.test.ts +++ b/sdk/typescript/tests-ts/api.test.ts @@ -23,7 +23,7 @@ import { type ThreadEvent, type ThreadOptions, } from "@openai/codex-sdk"; -import { afterEach, describe, expect, mock, test } from "bun:test"; +import { afterEach, describe, expect, mock, spyOn, test } from "bun:test"; import { parse as parseToml } from "smol-toml"; import { AuthenticationRequiredError, @@ -48,7 +48,11 @@ import { OPENROUTER_CODEX_PROVIDER, type JsonObject, } from "../src/config.js"; -import { estimateScanCost, type ScanCost } from "../src/cost.js"; +import { + estimateScanCost, + type ScanCost, + ScanCostTracker, +} from "../src/cost.js"; import { resolveCodexCommand, runWorkbench } from "../src/runtime.js"; import { normalizeTarget } from "../src/targets.js"; import { SYNTHETIC_CREDENTIALS } from "./cli-fixtures.js"; @@ -77,6 +81,12 @@ const EXAMPLE = join(PLUGIN_ROOT, "examples", "completed-scan"); const { cleanup, copyCompletedScan, temporaryDirectory } = createApiTestFixtures(); afterEach(cleanup); +const RESET_BUDGET_USAGE = { + input_tokens: 1_200, + cached_input_tokens: 200, + cache_write_input_tokens: 0, + output_tokens: 25, +}; const EXTERNAL_PROVIDER_CASES = [ [ @@ -157,11 +167,12 @@ async function writeUsageSession( threadId: string, usage: Record, parentThreadId?: string, -): Promise { +): Promise { const directory = join(codexHome, "sessions", "2026", "07", "26"); await mkdir(directory, { recursive: true }); + const path = join(directory, `rollout-${threadId}.jsonl`); await writeFile( - join(directory, `rollout-${threadId}.jsonl`), + path, [ JSON.stringify({ type: "session_meta", @@ -182,6 +193,7 @@ async function writeUsageSession( "", ].join("\n"), ); + return path; } describe("CodexSecurity finding validation", () => { @@ -2970,7 +2982,6 @@ describe("CodexSecurity orchestration", () => { const scan = client.run(repository, { ...(enforceCostLimit ? { maxCostUsd: 1 } : {}), - onActivity: () => {}, onWarning: (warning) => warnings.push(warning), }); if (enforceCostLimit) { @@ -2987,6 +2998,192 @@ describe("CodexSecurity orchestration", () => { }, ); + test.each([ + "complete", + "complete-followup", + "unverified", + "canceled-before", + "canceled-during", + "canceled-rejection", + "over-budget", + ] as const)( + "rechecks strict final accounting when the fresh result is %s", + async (state) => { + const root = await temporaryDirectory(); + const repository = join(root, "repository"); + const codexHome = join(root, "codex-home"); + const scanDir = join(root, "scan"); + await Promise.all([ + mkdir(repository), + mkdir(codexHome), + mkdir(scanDir, { mode: 0o700 }), + ]); + const model = "gpt-5.6-sol"; + const rootUsage = { + input_tokens: 100, + cached_input_tokens: 0, + cache_write_input_tokens: 0, + output_tokens: 10, + reasoning_output_tokens: 0, + }; + const finalUsage = { + ...rootUsage, + input_tokens: 200, + output_tokens: 20, + }; + const rootCost = estimateScanCost(model, rootUsage)!; + const finalCost = estimateScanCost(model, finalUsage)!; + const completes = state === "complete" || state === "complete-followup"; + const overBudget = state === "over-budget"; + const maxCostUsd = overBudget ? 0.001 : 1; + const firstError = new Error("Initial strict accounting check failed."); + const retryError = new Error("Fresh strict accounting check failed."); + const cancellation = new AbortController(); + const retryStarted = Promise.withResolvers(); + const releaseRetry = Promise.withResolvers(); + const commands: Array = []; + const finalizations: unknown[] = []; + let cleanupStops = 0; + let turns = 0; + const client = new TestClient( + { codexOverrides: { model } }, + { + environment: {}, + prepareRuntime: async () => preparedRuntime(codexHome), + resolvePluginPython: async () => "/managed/python", + prepareOutputDir: async () => scanDir, + repositoryRevision: async () => "deadbeef", + runWorkbench: async ( + _options: unknown, + args: readonly string[], + input?: string, + ) => { + commands.push(args); + return mockWorkbench(args, input); + }, + createCodex: () => ({ + startThread: () => ({ + id: null, + async runStreamed() { + turns += 1; + await copyCompletedScan(root); + async function* events(): AsyncGenerator { + for await (const event of completedEvents()) { + yield event.type === "turn.completed" + ? { ...event, usage: rootUsage } + : event; + } + } + return { events: events() }; + }, + }), + }), + }, + ); + const start = spyOn( + ScanCostTracker.prototype, + "start", + ).mockImplementation(() => {}); + const originalStop = ScanCostTracker.prototype.stop; + const stop = spyOn(ScanCostTracker.prototype, "stop").mockImplementation( + async function ( + this: ScanCostTracker, + ...args: Parameters + ) { + if (args.length === 0) { + cleanupStops += 1; + return overBudget + ? { usage: rootUsage, cost: rootCost } + : { usage: finalUsage, cost: finalCost }; + } + finalizations.push(args[0]); + if (finalizations.length === 1) { + if (state === "canceled-before") cancellation.abort(); + throw firstError; + } + retryStarted.resolve(); + await releaseRetry.promise; + if (state === "unverified" || state === "canceled-rejection") { + throw retryError; + } + const snapshot = await originalStop.call(this, finalUsage); + if (overBudget) throw retryError; + return snapshot; + }, + ); + const outcome = client + .run(repository, { + maxCostUsd, + signal: cancellation.signal, + ...(state === "unverified" || state === "complete-followup" + ? { postScanPrompt: "Record the completed scan accounting." } + : {}), + }) + .then( + (result) => ({ ok: true as const, result }), + (error: unknown) => ({ ok: false as const, error }), + ); + try { + if (state !== "canceled-before") { + await Promise.race([ + retryStarted.promise, + outcome.then(() => { + throw new Error("Scan ended before strict re-verification."); + }), + ]); + expect( + commands.some((args) => args[0] === "prepare-scan-completion"), + ).toBe(false); + if (state === "canceled-during" || state === "canceled-rejection") { + cancellation.abort(); + } + releaseRetry.resolve(); + } + const settled = await outcome; + if (completes) { + expect(settled.ok).toBe(true); + if (!settled.ok) throw settled.error; + expect(settled.result.turnResult.usage).toMatchObject(finalUsage); + expect(settled.result.cost).toEqual(finalCost); + const completion = commands.find( + (args) => args[0] === "complete-scan", + ); + expect(completion).toContain(JSON.stringify(finalCost)); + expect(commands.some((args) => args[0] === "fail-scan")).toBe(false); + } else { + expect(settled.ok).toBe(false); + if (settled.ok) throw new Error("Unverified scan completed."); + if (state.startsWith("canceled-")) { + expect(settled.error).toBeInstanceOf(ScanInterruptedError); + } else if (overBudget) { + expect(settled.error).toMatchObject({ + name: ScanCostLimitExceededError.name, + maxCostUsd, + cost: finalCost, + }); + } else { + expect(settled.error).toBe(firstError); + } + const failure = commands.find((args) => args[0] === "fail-scan"); + expect(failure).toContain(JSON.stringify(finalCost)); + expect(commands.some((args) => args[0] === "complete-scan")).toBe( + false, + ); + } + expect(finalizations).toHaveLength(state === "canceled-before" ? 1 : 2); + for (const usage of finalizations) expect(usage).toBe(rootUsage); + expect(cleanupStops).toBe(completes ? 0 : 1); + expect(turns).toBe(state === "complete-followup" ? 2 : 1); + } finally { + releaseRetry.resolve(); + await outcome; + stop.mockRestore(); + start.mockRestore(); + await client.close(); + } + }, + ); + test("uses the actual scanner inventory instead of a stale workbench estimate", async () => { const root = await temporaryDirectory(); const repository = join(root, "repository"); @@ -3917,146 +4114,263 @@ describe("CodexSecurity orchestration", () => { }, ); - test("stops and records a scan as soon as its live cost exceeds the limit", async () => { - const root = await temporaryDirectory(); - const repository = join(root, "repository"); - const codexHome = join(root, "codex-home"); - const scanDir = join(root, "scan"); - await mkdir(repository); - await mkdir(codexHome); - await mkdir(scanDir, { mode: 0o700 }); - const commands: Array = []; - const costs: number[] = []; - let turns = 0; - const cost = { - model: "gpt-5.6-sol", - inputTokens: 1_250, - cachedInputTokens: 200, - cacheWriteInputTokens: 0, - outputTokens: 30, - estimatedUsd: 0.00625, - }; - const client = new TestClient( - {}, - { - environment: {}, - prepareRuntime: async () => preparedRuntime(codexHome), - resolvePluginPython: async () => "/managed/python", - prepareOutputDir: async () => scanDir, - repositoryRevision: async () => "deadbeef", - runWorkbench: async ( - _options: unknown, - args: readonly string[], - input?: string, - ): Promise => { - commands.push(args); - if (args[0] === "register-cli-scan") { - return mockScanRegistration(args, input); - } - if (args[0] === "get-scan-feedback") { - return { - scanId: "scan_example_001", - targetId: "target_sha256_example", - falsePositives: [], + test.each([ + ["live polling", "live"], + ["a smaller cleanup counter snapshot", "reset"], + ["turn completion", "completed"], + ["turn completion after a tracking failure", "tracking-failure"], + ["failure cleanup after a tracking failure", "cleanup"], + ["turn completion after user cancellation", "canceled"], + ] as const)( + "stops and records a scan with over-budget usage from %s", + async (_description, costSource) => { + const root = await temporaryDirectory(); + const repository = join(root, "repository"); + const codexHome = join(root, "codex-home"); + const scanDir = join(root, "scan"); + await mkdir(repository); + await mkdir(codexHome); + await mkdir(scanDir, { mode: 0o700 }); + const commands: Array = []; + const costs: number[] = []; + const cancellation = new AbortController(); + const userCanceled = costSource === "canceled"; + const liveUsage = costSource === "live" || costSource === "reset"; + let turns = 0; + let rejectPoll: ((error: unknown) => void) | undefined; + const rootUsage: Record = + costSource === "cleanup" + ? { + input_tokens: 1_250, + cached_input_tokens: 200, + output_tokens: 30, + } + : liveUsage + ? { + input_tokens: 500, + cached_input_tokens: 100, + output_tokens: 10, + } + : { input_tokens: 100, output_tokens: 1 }; + const workerUsage: Record = + costSource === "cleanup" + ? {} + : { + input_tokens: liveUsage ? 750 : 250, + cached_input_tokens: 100, + output_tokens: liveUsage ? 20 : 10, }; - } - return {}; - }, - createCodex: () => ({ - startThread: () => ({ - id: null, - async runStreamed( - _input: string, - options: { signal: AbortSignal }, - ) { - turns += 1; - async function* events(): AsyncGenerator { - yield { type: "thread.started", thread_id: "scan-thread" }; - await Promise.all([ - writeUsageSession(codexHome, "scan-thread", { - input_tokens: 500, - cached_input_tokens: 100, - output_tokens: 10, - }), - writeUsageSession( - codexHome, - "worker-thread", - { - input_tokens: 750, - cached_input_tokens: 100, - output_tokens: 20, - }, - "scan-thread", - ), - ]); - await new Promise((resolve) => { - if (options.signal.aborted) { - resolve(); - } else { - options.signal.addEventListener("abort", () => resolve(), { - once: true, - }); + const cost = { + model: "gpt-5.6-sol", + inputTokens: 1_250, + cachedInputTokens: 200, + cacheWriteInputTokens: 0, + outputTokens: 30, + estimatedUsd: 0.00625, + }; + const client = new TestClient( + {}, + { + environment: {}, + prepareRuntime: async () => preparedRuntime(codexHome), + resolvePluginPython: async () => "/managed/python", + prepareOutputDir: async () => scanDir, + repositoryRevision: async () => "deadbeef", + runWorkbench: async ( + _options: unknown, + args: readonly string[], + input?: string, + ) => { + commands.push(args); + if (args[0] === "register-cli-scan") { + return mockScanRegistration(args, input); + } + if (args[0] === "get-scan-feedback") { + return { + scanId: "scan_example_001", + targetId: "target_sha256_example", + falsePositives: [], + }; + } + return {}; + }, + createCodex: () => ({ + startThread: () => ({ + id: null, + async runStreamed( + _input: string, + options: { signal: AbortSignal }, + ) { + turns += 1; + async function* events(): AsyncGenerator { + yield { type: "thread.started", thread_id: "scan-thread" }; + await Promise.all([ + writeUsageSession(codexHome, "scan-thread", rootUsage), + writeUsageSession( + codexHome, + "worker-thread", + workerUsage, + "scan-thread", + ), + ]); + if (costSource === "cleanup") { + expect(rejectPoll).toBeDefined(); + rejectPoll?.(new Error("session read failed")); + rejectPoll = undefined; } - }); - throw new DOMException("aborted", "AbortError"); - } - return { events: events() }; - }, + if (!liveUsage && costSource !== "cleanup") { + await copyCompletedScan(root); + yield { + type: "turn.completed", + usage: { + input_tokens: 1_000, + cached_input_tokens: 100, + cache_write_input_tokens: 0, + output_tokens: 20, + reasoning_output_tokens: 0, + }, + }; + return; + } + await new Promise((resolve) => { + if (options.signal.aborted) { + resolve(); + } else { + options.signal.addEventListener( + "abort", + () => resolve(), + { + once: true, + }, + ); + } + }); + throw new DOMException("aborted", "AbortError"); + } + return { events: events() }; + }, + }), }), - }), - }, - ); + }, + ); - // The fake Codex stream has no process handle to keep its unref'ed poll alive. - const keepEventLoopAlive = setTimeout(() => {}, 10_000); - try { - await expect( - client.run(repository, { + // The fake Codex stream has no process handle to keep its unref'ed poll alive. + const keepEventLoopAlive = setTimeout(() => {}, 10_000); + const originalRefresh = ScanCostTracker.prototype.refresh; + const originalStop = ScanCostTracker.prototype.stop; + let firstRefresh = true; + const stop = + costSource === "reset" + ? spyOn(ScanCostTracker.prototype, "stop").mockImplementation( + async function ( + this: ScanCostTracker, + ...args: Parameters + ) { + const snapshot = await originalStop.apply(this, args); + return args.length === 0 + ? { + usage: RESET_BUDGET_USAGE, + cost: estimateScanCost("gpt-5.6-sol", RESET_BUDGET_USAGE), + } + : snapshot; + }, + ) + : null; + const refresh = + costSource === "tracking-failure" || + costSource === "cleanup" || + userCanceled + ? spyOn(ScanCostTracker.prototype, "refresh").mockImplementation( + function (this: ScanCostTracker) { + if (firstRefresh) { + firstRefresh = false; + return new Promise< + Awaited> + >((_resolve, reject) => { + rejectPoll = reject; + }); + } + if (rejectPoll !== undefined) { + if (userCanceled) cancellation.abort(); + rejectPoll(new Error("session read failed")); + rejectPoll = undefined; + } + return originalRefresh.call(this); + }, + ) + : null; + try { + const scan = client.run(repository, { maxCostUsd: 0.005, postScanPrompt: "Record the scan cost.", onCost: (cost) => costs.push(cost.estimatedUsd), - signal: AbortSignal.timeout(5_000), - }), - ).rejects.toMatchObject({ - name: ScanCostLimitExceededError.name, - maxCostUsd: 0.005, - scanDir, - cost, - }); - } finally { - clearTimeout(keepEventLoopAlive); - } - expect(turns).toBe(1); - expect(costs.at(-1)).toBe(0.00625); - expect(commands[1]).toEqual([ - "get-scan-feedback", - "--scan-id", - "scan_example_001", - ]); - expect(commands[2]).toEqual([ - "set-scan-thread", - "--scan-id", - "scan_example_001", - "--thread-id", - "scan-thread", - ]); - expect(commands[3]).toEqual([ - "fail-scan", - "--scan-id", - "scan_example_001", - "--message", - `Scan stopped: estimated cost $0.00625 exceeded the $0.005 limit; partial output remains at ${scanDir}.`, - "--cost-json", - JSON.stringify(cost), - ]); - expect(commands.some((args) => args[0] === "complete-scan")).toBe(false); - await expect(stat(scanDir)).resolves.toBeDefined(); - await client.close(); - }); + signal: AbortSignal.any([ + cancellation.signal, + AbortSignal.timeout(5_000), + ]), + }); + if (userCanceled) { + const failure = await scan.catch((error: unknown) => error); + expect(failure).toBeInstanceOf(ScanInterruptedError); + expect(failure).not.toBeInstanceOf(ScanCostLimitExceededError); + } else { + await expect(scan).rejects.toMatchObject({ + name: ScanCostLimitExceededError.name, + maxCostUsd: 0.005, + scanDir, + cost, + }); + } + } finally { + stop?.mockRestore(); + refresh?.mockRestore(); + clearTimeout(keepEventLoopAlive); + } + expect(turns).toBe(1); + expect(costs.at(-1)).toBe(0.00625); + expect(commands[1]).toEqual([ + "get-scan-feedback", + "--scan-id", + "scan_example_001", + ]); + expect(commands[2]).toEqual([ + "set-scan-thread", + "--scan-id", + "scan_example_001", + "--thread-id", + "scan-thread", + ]); + if (userCanceled) { + expect(commands[3]?.[0]).toBe("fail-scan"); + expect(commands[3]).toContain(JSON.stringify(cost)); + } else { + expect(commands[3]).toEqual([ + "fail-scan", + "--scan-id", + "scan_example_001", + "--message", + `Scan stopped: estimated cost $0.00625 exceeded the $0.005 limit; partial output remains at ${scanDir}.`, + "--cost-json", + JSON.stringify(cost), + ]); + } + expect(commands.some((args) => args[0] === "complete-scan")).toBe(false); + await expect(stat(scanDir)).resolves.toBeDefined(); + await client.close(); + }, + ); - test.each(["partial", "invalid", "unavailable"] as const)( - "recovers exhausted deep-scan budget when completion is %s", - async (completion) => { + test.each([ + ["partial", "live"], + ["invalid", "live"], + ["unavailable", "live"], + ["partial", "completed"], + ["partial", "reset"], + ["partial", "cleanup-error"], + ["partial", "flushed"], + ] as const)( + "recovers exhausted deep-scan budget when completion is %s using %s usage", + async (completion, costSource) => { const root = await temporaryDirectory(); const repository = join(root, "repository"); const codexHome = join(root, "codex-home"); @@ -4068,7 +4382,9 @@ describe("CodexSecurity orchestration", () => { ]); const commands: Array = []; const warnings: string[] = []; + const expectedCostUsd = costSource === "flushed" ? 0.0078 : 0.00625; let turns = 0; + let workerSession: string | null = null; const client = new TestClient( {}, { @@ -4126,10 +4442,44 @@ describe("CodexSecurity orchestration", () => { turns += 1; async function* events(): AsyncGenerator { yield { type: "thread.started", thread_id: "scan-thread" }; + if (costSource === "completed" || costSource === "flushed") { + const sessions = await Promise.all([ + writeUsageSession(codexHome, "scan-thread", { + input_tokens: 100, + output_tokens: 1, + }), + writeUsageSession( + codexHome, + "worker-thread", + { + input_tokens: 250, + cached_input_tokens: 100, + output_tokens: 10, + }, + "scan-thread", + ), + ]); + workerSession = sessions[1]!; + yield { + type: "turn.completed", + usage: { + input_tokens: 1_000, + cached_input_tokens: 100, + cache_write_input_tokens: 0, + output_tokens: 20, + reasoning_output_tokens: 0, + }, + }; + return; + } await writeUsageSession(codexHome, "scan-thread", { input_tokens: 1_250, cached_input_tokens: 200, output_tokens: 30, + reasoning_output_tokens: + costSource === "reset" || costSource === "cleanup-error" + ? 12 + : 0, }); await new Promise((resolve) => { if (options.signal.aborted) resolve(); @@ -4152,6 +4502,66 @@ describe("CodexSecurity orchestration", () => { }, ); const keepAlive = setTimeout(() => {}, 10_000); + const originalStop = ScanCostTracker.prototype.stop; + const stop = + costSource === "reset" || + costSource === "cleanup-error" || + costSource === "flushed" + ? spyOn(ScanCostTracker.prototype, "stop").mockImplementation( + async function ( + this: ScanCostTracker, + ...args: Parameters + ) { + try { + const snapshot = await originalStop.apply(this, args); + if (costSource === "cleanup-error" && args.length === 0) { + throw new Error("Synthetic cleanup accounting failure"); + } + return costSource === "reset" && args.length === 0 + ? { + usage: RESET_BUDGET_USAGE, + cost: estimateScanCost( + "gpt-5.6-sol", + RESET_BUDGET_USAGE, + ), + } + : snapshot; + } finally { + if ( + costSource === "flushed" && + args.length > 0 && + workerSession !== null + ) { + const path = workerSession; + workerSession = null; + await appendFile( + path, + [ + JSON.stringify({ + type: "event_msg", + payload: { + type: "token_count", + info: { + total_token_usage: { + input_tokens: 500, + cached_input_tokens: 100, + output_tokens: 20, + }, + }, + }, + }), + JSON.stringify({ + type: "event_msg", + payload: { type: "task_complete" }, + }), + "", + ].join("\n"), + ); + } + } + }, + ) + : null; try { const result = client.run(repository, { mode: "deep", @@ -4176,9 +4586,13 @@ describe("CodexSecurity orchestration", () => { expect(recovered.coverage.completeness).toBe(completion); expect(recovered.findings.findings).toHaveLength(1); expect(recovered.threadId).toBe("scan-thread"); - expect(recovered.cost?.estimatedUsd).toBe(0.00625); + expect(recovered.cost?.estimatedUsd).toBe(expectedCostUsd); + expect(recovered.turnResult.usage).toMatchObject({ + reasoning_output_tokens: + costSource === "reset" || costSource === "cleanup-error" ? 12 : 0, + }); expect(warnings).toEqual([ - `Scan stopped: estimated cost $0.00625 exceeded the $0.005 limit; partial output remains at ${scanDir}.`, + `Scan stopped: estimated cost $${expectedCostUsd} exceeded the $0.005 limit; partial output remains at ${scanDir}.`, ]); expect(commands.some((args) => args[0] === "fail-scan")).toBe(false); } @@ -4188,14 +4602,98 @@ describe("CodexSecurity orchestration", () => { ); expect(recovery?.includes("--cost-json")).toBe(true); expect(recovery?.includes("--message")).toBe(true); + expect( + JSON.parse(recovery![recovery!.indexOf("--cost-json") + 1]!), + ).toMatchObject({ estimatedUsd: expectedCostUsd }); } finally { + stop?.mockRestore(); clearTimeout(keepAlive); await client.close(); } }, ); - test("saves a budgeted scan with a warning when token usage is unavailable", async () => { + test.each([ + ["explicitly budgeted", true], + ["optionally accounted", false], + ] as const)( + "requires every owned session only for an %s scan", + async (_description, enforceCostLimit) => { + const root = await temporaryDirectory(); + const repository = join(root, "repository"); + const codexHome = join(root, "codex-home"); + const scanDir = join(root, "scan"); + await Promise.all([ + mkdir(repository), + mkdir(codexHome), + mkdir(scanDir, { mode: 0o700 }), + ]); + const commands: string[] = []; + let rootSession: string | null = null; + let ownershipChecks = 0; + const client = new TestClient( + {}, + { + environment: {}, + prepareRuntime: async () => preparedRuntime(codexHome), + resolvePluginPython: async () => "/managed/python", + resolveScanSessionPaths: async ( + _options: unknown, + scanId: string, + threadId: string, + ) => { + ownershipChecks += 1; + expect(scanId).toBe("scan_example_001"); + expect(threadId).toBe("thread-1"); + return new Set([ + rootSession!, + join(codexHome, "missing-owned-worker.jsonl"), + ]); + }, + prepareOutputDir: async () => scanDir, + repositoryRevision: async () => "deadbeef", + runWorkbench: async ( + _options: unknown, + args: readonly string[], + input?: string, + ) => { + commands.push(args[0]!); + return mockWorkbench(args, input); + }, + createCodex: () => ({ + startThread: () => ({ + id: null, + async runStreamed() { + await copyCompletedScan(root); + rootSession = await writeUsageSession(codexHome, "thread-1", { + input_tokens: 100, + output_tokens: 10, + }); + return { events: completedEvents() }; + }, + }), + }), + }, + ); + + const scan = client.run(repository, { + ...(enforceCostLimit ? { maxCostUsd: 1 } : {}), + }); + if (enforceCostLimit) { + await expect(scan).rejects.toThrow("cost limit could not be verified"); + expect(ownershipChecks).toBe(2); + expect(commands).toContain("fail-scan"); + expect(commands).not.toContain("complete-scan"); + } else { + await expect(scan).resolves.toMatchObject({ threadId: "thread-1" }); + expect(ownershipChecks).toBe(0); + expect(commands).toContain("complete-scan"); + } + await client.close(); + }, + ); + + test("fails a budgeted scan when only unfinished root usage is available", async () => { const root = await temporaryDirectory(); const repository = join(root, "repository"); const codexHome = join(root, "codex-home"); @@ -4242,6 +4740,10 @@ describe("CodexSecurity orchestration", () => { }; executable._exec.run = async function* () { await copyCompletedScan(root); + await writeUsageSession(codexHome, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); yield JSON.stringify({ type: "thread.started", thread_id: "scan-thread", @@ -4254,26 +4756,24 @@ describe("CodexSecurity orchestration", () => { }, ); - const result = await client.run(repository, { - maxCostUsd: 1, - onWarning: (warning) => { - warnings.push(warning); - }, - }); - expect(result.threadId).toBe("scan-thread"); - expect(result.cost).toBeNull(); - expect(warnings).toEqual([ - "Scan completed, but its cost limit could not be verified because model pricing or token usage is unavailable.", - ]); + await expect( + client.run(repository, { + maxCostUsd: 1, + onWarning: (warning) => { + warnings.push(warning); + }, + }), + ).rejects.toThrow( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + expect(warnings).toEqual([]); expect(commands.map(([command]) => command)).toEqual([ "register-cli-scan", "get-scan-feedback", "set-scan-thread", - "prepare-scan-completion", - "complete-scan", - "list-global-findings", + "fail-scan", ]); - expect(commands.some((args) => args[0] === "fail-scan")).toBe(false); + expect(commands.some((args) => args[0] === "complete-scan")).toBe(false); await client.close(); }); diff --git a/sdk/typescript/tests-ts/cost.test.ts b/sdk/typescript/tests-ts/cost.test.ts index 3d0fbe317..d4fff526c 100644 --- a/sdk/typescript/tests-ts/cost.test.ts +++ b/sdk/typescript/tests-ts/cost.test.ts @@ -1,6 +1,8 @@ import { spawnSync } from "node:child_process"; +import * as fsPromises from "node:fs/promises"; import { appendFile, + chmod, mkdir, mkdtemp, realpath, @@ -9,7 +11,7 @@ import { } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join, parse, sep } from "node:path"; -import { afterEach, describe, expect, test } from "bun:test"; +import { afterEach, describe, expect, mock, spyOn, test } from "bun:test"; import { estimateScanCost, ScanCostTracker, @@ -19,8 +21,11 @@ import type { ScanActivity } from "../src/scan-activity.js"; import { readScanLogs } from "../src/scan-logs.js"; import { sessionParentThreadId } from "../src/scan-sessions.js"; import type { ScanProgress } from "../src/worker-progress.js"; +import { runTestInSubprocess } from "./support/test-subprocess.js"; const temporaryDirectories: string[] = []; +const testPosix = + process.platform === "win32" || process.geteuid?.() === 0 ? test.skip : test; const parentFields = ["source", "parent_thread_id", "forked_from_id"] as const; type SessionParentField = (typeof parentFields)[number]; @@ -61,10 +66,11 @@ async function codexHome(): Promise { async function writeSession( home: string, threadId: string, - usage: Record, + usage: Record | null, parentThreadId?: string, workingDirectory?: string, timestamp?: string, + completed = false, parentField: SessionParentField = "source", ): Promise { const directory = join(home, "sessions", "2026", "07", "26"); @@ -84,19 +90,213 @@ async function writeSession( : parentMetadata(parentThreadId, parentField)), }, }), - JSON.stringify({ - type: "event_msg", - payload: { - type: "token_count", - info: { total_token_usage: usage }, - }, - }), + ...(completed ? [taskEvent("task_started")] : []), + ...(usage === null ? [] : [JSON.stringify(accountingEvent(usage))]), + ...(completed ? [taskEvent("task_complete")] : []), "", ].join("\n"), ); return path; } +function taskEvent( + type: "task_started" | "task_complete" | "turn_complete" | "turn_aborted", +): string { + return JSON.stringify({ + type: "event_msg", + payload: { + type, + turn_id: "fixture-turn", + started_at: 1_785_067_320, + ...(type === "task_started" ? {} : { completed_at: 1_785_067_321 }), + ...(type === "turn_aborted" ? { reason: "interrupted" } : {}), + }, + }); +} + +type MockAccountingEvent = Readonly> | Error; + +function accountingEvent( + usage: Readonly> | null, +): MockAccountingEvent { + return { + type: "event_msg", + payload: { type: "token_count", info: { total_token_usage: usage } }, + }; +} + +const accountingFork = { + // The inherited turn, child, and owned turn share a UUID7 millisecond. + inheritedTurnId: "019f9e4d-b3ba-7000-8000-000000000001", + threadId: "019f9e4d-b3ba-7000-8000-000000000002", + ownedTurnId: "019f9e4d-b3ba-7000-8000-000000000003", + pendingTurnId: "019f9e4d-b450-4000-8000-000000000004", + timestamp: "2026-07-26T12:02:00.250Z", + startedAt: 1_785_067_320, +}; + +function accountingTaskStart( + turnId: string, + startedAt: number, +): MockAccountingEvent { + return { + type: "event_msg", + payload: { + type: "task_started", + turn_id: turnId, + started_at: startedAt, + }, + }; +} + +function accountingSession( + threadId: string, + events: readonly MockAccountingEvent[], + parentThreadId?: string, + metadata: Readonly> = {}, +): MockAccountingEvent[] { + return [ + { + type: "session_meta", + payload: { ...metadata, id: threadId, parent_thread_id: parentThreadId }, + }, + ...events, + { type: "event_msg", payload: { type: "task_complete" } }, + ]; +} + +async function withMockAccountingSessions( + sessions: Readonly>, + options: Omit[0], "codexHome">, + check: ( + tracker: ScanCostTracker, + append: (threadId: string, events: readonly MockAccountingEvent[]) => void, + omit: (threadId: string) => void, + ) => Promise, + beforeOpen?: ( + path: string, + attempt: number, + append: (threadId: string, events: readonly MockAccountingEvent[]) => void, + ) => void, +): Promise { + const home = join(tmpdir(), "codex-security-mock-cost"); + const directory = join(home, "sessions"); + const files = new Map(); + const omittedFiles = new Set(); + const opens = new Map(); + const append = ( + threadId: string, + next: readonly MockAccountingEvent[], + ): void => { + const path = join(directory, `rollout-${threadId}.jsonl`); + const lines = next.map((event) => + event instanceof Error ? "{\n" : `${JSON.stringify(event)}\n`, + ); + files.set( + path, + Buffer.concat([ + files.get(path) ?? Buffer.alloc(0), + Buffer.from(lines.join("")), + ]), + ); + }; + for (const [threadId, initial] of Object.entries(sessions)) { + append(threadId, initial); + } + const originalOpen = fsPromises.open; + const originalReaddir = fsPromises.readdir; + mock.module("node:fs/promises", () => ({ + ...fsPromises, + readdir: async (path: unknown) => { + if (String(path) !== directory) + throw new Error("Unexpected session directory"); + return [...files.keys()] + .filter((path) => !omittedFiles.has(path)) + .map((path) => ({ + name: path.slice(directory.length + 1), + isDirectory: () => false, + isFile: () => true, + })); + }, + open: async (path: unknown) => { + const name = String(path); + const attempt = (opens.get(name) ?? 0) + 1; + opens.set(name, attempt); + beforeOpen?.(name, attempt, append); + const contents = files.get(name); + if (contents === undefined) throw new Error("Unexpected session file"); + return { + read: async ( + buffer: Buffer, + offset: number, + length: number, + position: number, + ) => ({ + bytesRead: + position >= contents.length + ? 0 + : contents.copy(buffer, offset, position, position + length), + buffer, + }), + close: async () => {}, + }; + }, + })); + const tracker = new ScanCostTracker({ ...options, codexHome: home }); + tracker.start("scan-thread"); + try { + await check(tracker, append, (threadId) => { + omittedFiles.add(join(directory, `rollout-${threadId}.jsonl`)); + }); + } finally { + await tracker.stop().catch(() => {}); + mock.module("node:fs/promises", () => ({ + ...fsPromises, + open: originalOpen, + readdir: originalReaddir, + })); + } +} + +async function workerScan({ + rootUsage = { input_tokens: 100, output_tokens: 10 }, + workerUsage = { input_tokens: 100, output_tokens: 10 }, + workerCompleted = true, + maxCostUsd, + onCost, +}: { + rootUsage?: Record | null; + workerUsage?: Record | null; + workerCompleted?: boolean; + maxCostUsd?: number; + onCost?: (cost: { estimatedUsd: number }) => void; +} = {}): Promise<{ + home: string; + root: string; + worker: string; + tracker: ScanCostTracker; +}> { + const home = await codexHome(); + const root = await writeSession(home, "scan-thread", rootUsage); + const worker = await writeSession( + home, + "worker-thread", + workerUsage, + "scan-thread", + undefined, + undefined, + workerCompleted, + ); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd, + onCost, + }); + tracker.start("scan-thread"); + return { home, root, worker, tracker }; +} + async function appendSessionItem( path: string, payload: Readonly>, @@ -107,6 +307,13 @@ async function appendSessionItem( ); } +async function appendIncompleteTokenUsage(path: string): Promise { + const event = JSON.stringify( + accountingEvent({ input_tokens: 10_000, output_tokens: 1_000 }), + ); + await appendFile(path, event.slice(0, -1)); +} + function progressMessage( filesCompleted: number, filesTotal = 8, @@ -371,6 +578,56 @@ describe("scan cost", () => { }); describe("live scan cost tracking", () => { + test.each([null, 40, 80])( + "counts SDK-completed validation and its workers once with rollout usage %s", + async (validationInput) => { + const home = await codexHome(); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + await writeSession( + home, + "scan-thread", + rootUsage, + undefined, + undefined, + undefined, + true, + ); + if (validationInput !== null) + await writeSession(home, "validation-thread", { + input_tokens: validationInput, + output_tokens: validationInput / 10, + }); + await writeSession( + home, + "worker-thread", + { input_tokens: 20, output_tokens: 2 }, + "validation-thread", + undefined, + undefined, + true, + ); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + maxCostUsd: 1, + }); + tracker.start("scan-thread"); + tracker.recordCompletedThreadUsage("validation-thread", { + input_tokens: 40, + output_tokens: 4, + }); + + const snapshot = await tracker.stop(rootUsage); + + expect(snapshot.cost?.inputTokens).toBe( + 120 + Math.max(40, validationInput ?? 0), + ); + expect(snapshot.cost?.outputTokens).toBe( + 12 + Math.max(4, (validationInput ?? 0) / 10), + ); + }, + ); + test("coalesces overlapping polling ticks and bounds final work", async () => { const home = await codexHome(); await writeSession(home, "scan-thread", { @@ -681,7 +938,8 @@ describe("live scan cost tracking", () => { { input_tokens: 50, output_tokens: 1 }, "deep-worker", undefined, - undefined, + "2026-07-26T12:02:00Z", + true, parentField, ); await writeSession( @@ -761,6 +1019,385 @@ describe("live scan cost tracking", () => { }, ); + test.each([ + ["root metadata is missing", "missing", "independent", 1, true], + ["root timing is missing", "untimed", "independent", 1, true], + ["root timing is invalid", "invalid", "independent", 1, true], + ["worker timing is missing", "timed", "untimed", 1, true], + ["worker timing is invalid", "timed", "invalid", 1, true], + ["the worker parent is unobserved", "missing", "orphaned", 1, true], + ["the worker has a parent", "missing", "parented", 1, false], + ["the worker is unrelated", "missing", "unrelated", 1, false], + ["tracking is optional", "missing", "independent", undefined, false], + ["no worker needs attribution", "missing", "none", 1, false], + ] as const)( + "verifies independent Deep worker ownership when %s", + async (_scenario, rootState, workerState, maxCostUsd, shouldReject) => { + const home = await codexHome(); + const scanDirectory = join(home, "scans", "current"); + if (rootState !== "missing") { + await writeSession( + home, + "scan-thread", + { input_tokens: 100, output_tokens: 10 }, + undefined, + scanDirectory, + rootState === "timed" + ? "2026-07-26T12:00:00Z" + : rootState === "invalid" + ? "not a timestamp" + : undefined, + ); + } + if (workerState !== "none") { + await writeSession( + home, + "deep-worker", + { input_tokens: 1_000, output_tokens: 100 }, + workerState === "parented" + ? "scan-thread" + : workerState === "orphaned" + ? "unobserved-coordinator" + : undefined, + workerState === "unrelated" + ? join(home, "another-scan", "artifacts") + : join( + scanDirectory, + "artifacts", + "deep_discovery", + "workers", + "worker", + "output", + ), + workerState === "untimed" + ? undefined + : workerState === "invalid" + ? "not a timestamp" + : "2026-07-26T12:01:00Z", + true, + ); + } + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + scanDirectory, + maxCostUsd, + }); + tracker.start("scan-thread"); + await tracker.refresh(); + + const completed = tracker.stop({ input_tokens: 100, output_tokens: 10 }); + if (shouldReject) { + await expect(completed).rejects.toThrow( + "The scan cost limit could not be verified", + ); + } else { + await expect(completed).resolves.toMatchObject({ + cost: { + inputTokens: workerState === "parented" ? 1_100 : 100, + }, + }); + } + }, + ); + + test("resolves mocked worker ancestry before directory attribution", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "resolves mocked worker ancestry before directory attribution", + ) + ) { + return; + } + const scanDirectory = join(tmpdir(), "codex-security-mock-scan"); + const artifacts = join(scanDirectory, "artifacts"); + const workerDirectory = join( + artifacts, + "deep_discovery", + "workers", + "worker", + "output", + ); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const session = ( + id: string, + parent?: string, + cwd?: string, + timestamp?: string, + input = 1_000, + ) => + accountingSession( + id, + [accountingEvent({ input_tokens: input, output_tokens: input / 10 })], + parent, + { cwd, timestamp }, + ); + const roots = { + "scan-thread": session( + "scan-thread", + undefined, + scanDirectory, + "2026-07-26T12:00:00Z", + 100, + ), + "previous-root": session( + "previous-root", + undefined, + scanDirectory, + "2026-07-26T11:00:00Z", + ), + }; + const previousWorkers = { + "previous-coordinator": session( + "previous-coordinator", + "previous-root", + artifacts, + ), + "previous-worker": session( + "previous-worker", + "previous-coordinator", + workerDirectory, + "invalid timestamp", + ), + "resumed-previous-worker": session( + "resumed-previous-worker", + "previous-root", + workerDirectory, + "2026-07-26T12:02:00Z", + ), + "older-independent": session( + "older-independent", + undefined, + workerDirectory, + "2026-07-26T11:59:00Z", + ), + "older-orphan": session( + "older-orphan", + "missing-parent", + workerDirectory, + "2026-07-26T11:59:00Z", + ), + "older-cycle-a": session( + "older-cycle-a", + "older-cycle-b", + workerDirectory, + "2026-07-26T11:59:00Z", + ), + "older-cycle-b": session("older-cycle-b", "older-cycle-a"), + }; + await withMockAccountingSessions( + { + ...roots, + ...previousWorkers, + "current-child": session( + "current-child", + "current-independent", + workerDirectory, + "2026-07-26T11:59:00Z", + 300, + ), + "current-independent": session( + "current-independent", + undefined, + artifacts, + "2026-07-26T12:01:00Z", + 200, + ), + "unrelated-conflict-a": session("unrelated-conflict", "previous-root"), + "unrelated-conflict-b": session("unrelated-conflict", "missing-parent"), + }, + { model: "gpt-5.6-terra", maxCostUsd: 1, scanDirectory }, + async (tracker) => { + expect((await tracker.stop(rootUsage)).cost).toMatchObject({ + inputTokens: 600, + outputTokens: 60, + }); + }, + ); + await withMockAccountingSessions( + { + ...roots, + ...previousWorkers, + "scan-thread": [ + new SyntaxError("mock parser diagnostic"), + ...roots["scan-thread"], + ], + }, + { model: "gpt-5.6-terra", maxCostUsd: 1, scanDirectory }, + async (tracker) => { + expect((await tracker.stop(rootUsage)).cost?.inputTokens).toBe(100); + }, + ); + const unresolvedWorkers: Array> = [ + { + orphan: session( + "orphan", + "missing-parent", + workerDirectory, + "2026-07-26T12:01:00Z", + ), + }, + { orphan: session("orphan", "missing-parent", workerDirectory) }, + { + "cycle-a": session("cycle-a", "cycle-b", workerDirectory), + "cycle-b": session("cycle-b", "cycle-a"), + }, + { + "conflict-a": session( + "conflict", + "scan-thread", + workerDirectory, + "2026-07-26T11:59:00Z", + ), + "conflict-b": session("conflict", "previous-root"), + }, + { + "conflict-a": session("conflict", "scan-thread"), + "conflict-b": session("conflict", "previous-root"), + }, + ]; + for (const workers of unresolvedWorkers) { + await withMockAccountingSessions( + { ...roots, ...workers }, + { model: "gpt-5.6-terra", maxCostUsd: 1, scanDirectory }, + async (tracker) => { + await expect(tracker.stop(rootUsage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + }, + ); + } + }); + + test("limits mocked worker-directory attribution to contained output paths", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "limits mocked worker-directory attribution to contained output paths", + ) + ) { + return; + } + const scanDirectory = join(tmpdir(), "codex-security-mock-scan"); + const artifacts = join(scanDirectory, "artifacts"); + const workers = join(artifacts, "deep_discovery", "workers"); + const timestamp = "2026-07-26T12:01:00Z"; + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + for (const [cwd, startedAt, expectedInput] of [ + [artifacts, timestamp, 1_100], + [join(workers, "worker", "output"), timestamp, 1_100], + [join(workers, "worker", "output"), undefined, null], + [join(artifacts, "deep_discovery", "output"), undefined, 100], + [join(artifacts, "deep_discovery", "output"), timestamp, 100], + [ + join(artifacts, "deep_discovery", "workers-other", "worker", "output"), + undefined, + 100, + ], + ] as const) { + await withMockAccountingSessions( + { + "scan-thread": accountingSession( + "scan-thread", + [accountingEvent(rootUsage)], + undefined, + { + cwd: scanDirectory, + timestamp: "2026-07-26T12:00:00Z", + }, + ), + worker: accountingSession( + "worker", + [accountingEvent({ input_tokens: 1_000, output_tokens: 100 })], + undefined, + { + cwd, + timestamp: startedAt, + }, + ), + }, + { model: "gpt-5.6-terra", maxCostUsd: 1, scanDirectory }, + async (tracker) => { + const stopped = tracker.stop(rootUsage); + if (expectedInput === null) { + await expect(stopped).rejects.toThrow( + "The scan cost limit could not be verified", + ); + } else { + expect((await stopped).cost?.inputTokens).toBe(expectedInput); + } + }, + ); + } + }); + + test("keeps mocked replay errors outside owned accounting", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps mocked replay errors outside owned accounting", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + for (const location of ["replay", "owned-before", "owned-after"] as const) { + const costs: number[] = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "worker-thread": [ + { + type: "session_meta", + payload: { + id: "worker-thread", + parent_thread_id: "scan-thread", + timestamp: "2026-07-26T12:02:00Z", + }, + }, + ...(location === "owned-before" + ? [new SyntaxError("mock owned parser diagnostic")] + : []), + { type: "session_meta", payload: { id: "scan-thread" } }, + new SyntaxError("mock replay parser diagnostic"), + accountingEvent({ input_tokens: 1_000, output_tokens: 100 }), + { + type: "event_msg", + payload: { type: "task_started", started_at: 1_785_067_320 }, + }, + ...(location === "owned-after" + ? [new SyntaxError("mock owned parser diagnostic")] + : []), + accountingEvent({ input_tokens: 1_200, output_tokens: 120 }), + { type: "event_msg", payload: { type: "task_complete" } }, + ], + }, + { + model: "gpt-5.6-terra", + maxCostUsd: 1, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker) => { + const stopped = tracker.stop(rootUsage); + if (location === "replay") { + await expect(stopped).resolves.toMatchObject({ + usage: { input_tokens: 300, output_tokens: 30 }, + cost: { estimatedUsd: 0.00096 }, + }); + } else { + await expect(stopped).rejects.toThrow( + "tracked session record could not be read", + ); + } + expect(costs).toEqual([0.00096]); + }, + ); + } + }); + test.each([ [ "sessions beside the deep worker output directories", @@ -790,7 +1427,7 @@ describe("live scan cost tracking", () => { "sessions with an invalid timestamp", (scan: string) => join(scan, "artifacts"), "not-a-timestamp", - undefined, + "unrelated-parent", "source", ], [ @@ -842,6 +1479,12 @@ describe("live scan cost tracking", () => { ), "2026-07-26T12:00:00.950Z", ); + if (parentThreadId !== undefined) { + await writeSession(home, parentThreadId, { + input_tokens: 0, + output_tokens: 0, + }); + } await writeSession( home, "bystander", @@ -849,6 +1492,7 @@ describe("live scan cost tracking", () => { parentThreadId, workingDirectory(scanDirectory), timestamp, + false, parentField, ); await writeSession( @@ -1897,109 +2541,3255 @@ describe("live scan cost tracking", () => { }); }); - test("retains a partial event across incremental reads", async () => { + test("accumulates usage after a token-counter reset", async () => { const home = await codexHome(); - const path = await writeSession(home, "scan-thread", { - input_tokens: 100, - output_tokens: 10, + const root = await writeSession(home, "scan-thread", { + input_tokens: 1_000, + output_tokens: 100, }); - const events: ScanSessionEvent[] = []; + const costs: number[] = []; const tracker = new ScanCostTracker({ codexHome: home, model: "gpt-5.6-terra", - onSessionEvent: (event) => events.push(event), + maxCostUsd: 0.001, + onCost: (cost) => costs.push(cost.estimatedUsd), }); tracker.start("scan-thread"); - await tracker.refresh(); - - const event = JSON.stringify({ - type: "event_msg", - payload: { - type: "token_count", - info: { - total_token_usage: { input_tokens: 250, output_tokens: 20 }, - }, - }, - }); - const padding = " ".repeat(128 * 1_024); - await appendFile(path, `${padding}${event.slice(0, 40)}`); - expect((await tracker.refresh()).cost?.inputTokens).toBe(100); - expect(events).toHaveLength(2); + expect((await tracker.refresh()).cost?.estimatedUsd).toBe(0.0032); + await appendFile( + root, + `${JSON.stringify(accountingEvent({ input_tokens: 500, output_tokens: 50 }))}\n`, + ); - await appendFile(path, `${event.slice(40)}\n`); - expect((await tracker.stop()).cost?.inputTokens).toBe(250); - expect(events).toHaveLength(3); - expect(events.at(-1)?.event).toEqual(JSON.parse(event)); + expect((await tracker.stop()).cost?.estimatedUsd).toBe(0.0048); + expect(costs).toEqual([0.0032, 0.0048]); }); - test("reads session events larger than 16 MiB", async () => { + test("enforces budgets across equal-total field resets without counting duplicates", async () => { const home = await codexHome(); - const path = await writeSession(home, "scan-thread", { - input_tokens: 100, - output_tokens: 10, + const root = await writeSession(home, "scan-thread", { + input_tokens: 1_000, + cached_input_tokens: 400, + cache_write_input_tokens: 100, + output_tokens: 1_000, + reasoning_output_tokens: 800, }); + const costs: number[] = []; const tracker = new ScanCostTracker({ codexHome: home, model: "gpt-5.6-terra", + maxCostUsd: 0.017, + onCost: (cost) => costs.push(cost.estimatedUsd), }); tracker.start("scan-thread"); + expect((await tracker.refresh()).cost?.estimatedUsd).toBe(0.01333); + const reset = JSON.stringify( + accountingEvent({ + input_tokens: 1_500, + cached_input_tokens: 600, + cache_write_input_tokens: 200, + output_tokens: 500, + reasoning_output_tokens: 50, + }), + ); + await appendFile(root, `${reset}\n${reset}\n`); - const event = JSON.stringify({ - type: "event_msg", - payload: { - type: "token_count", - info: { - total_token_usage: { input_tokens: 250, output_tokens: 20 }, - }, - details: "x".repeat(16 * 1_024 * 1_024 + 1), - }, - }); - await appendFile(path, event.slice(0, -10)); - expect((await tracker.refresh()).cost?.inputTokens).toBe(100); - - await appendFile(path, `${event.slice(-10)}\n`); - expect((await tracker.stop()).cost?.inputTokens).toBe(250); - }); - - test("reports a changed running cost only once", async () => { - const home = await codexHome(); - await writeSession(home, "scan-thread", { - input_tokens: 1_250, - cached_input_tokens: 200, - output_tokens: 30, - }); - const updates: number[] = []; - const tracker = new ScanCostTracker({ - codexHome: home, - model: "gpt-5.6-sol", - maxCostUsd: 0.005, - onCost: (cost) => updates.push(cost.estimatedUsd), + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 1_500, + cachedInputTokens: 600, + cacheWriteInputTokens: 200, + outputTokens: 1_500, + estimatedUsd: 0.02002, }); - tracker.start("scan-thread"); - - await tracker.stop(); - - expect(updates).toEqual([0.00625]); + expect(costs).toEqual([0.01333, 0.02002]); + }); + + test.each([ + [ + "cached input", + { input_tokens: 100, cached_input_tokens: 100, output_tokens: 100 }, + { input_tokens: 101, cached_input_tokens: 99, output_tokens: 99 }, + { input_tokens: 101, cached_input_tokens: 101, output_tokens: 199 }, + 0.0015, + ], + [ + "cache writes", + { input_tokens: 100, cache_write_input_tokens: 100, output_tokens: 100 }, + { input_tokens: 101, cache_write_input_tokens: 99, output_tokens: 99 }, + { input_tokens: 101, cache_write_input_tokens: 101, output_tokens: 199 }, + 0.0016, + ], + [ + "premium cache-write allocation", + { + input_tokens: 100, + cached_input_tokens: 90, + cache_write_input_tokens: 10, + output_tokens: 100, + reasoning_output_tokens: 100, + }, + { + input_tokens: 101, + cached_input_tokens: 91, + cache_write_input_tokens: 9, + output_tokens: 99, + reasoning_output_tokens: 99, + }, + { + input_tokens: 101, + cached_input_tokens: 90, + cache_write_input_tokens: 11, + output_tokens: 199, + reasoning_output_tokens: 199, + }, + 0.0015, + ], + [ + "reasoning output", + { input_tokens: 100, output_tokens: 100, reasoning_output_tokens: 100 }, + { input_tokens: 99, output_tokens: 101, reasoning_output_tokens: 99 }, + { input_tokens: 199, output_tokens: 101, reasoning_output_tokens: 101 }, + 0.0015, + ], + ] as const)( + "preserves %s invariants without accepting a stale scan budget", + async (_description, initial, next, expected, budget) => { + const home = await codexHome(); + const root = await writeSession(home, "scan-thread", initial); + const costs: number[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: budget, + onCost: (cost) => costs.push(cost.estimatedUsd), + }); + tracker.start("scan-thread"); + expect((await tracker.refresh()).cost?.estimatedUsd).toBeLessThan(budget); + await appendFile(root, `${JSON.stringify(accountingEvent(next))}\n`); + + const completed = await tracker.stop(initial); + expect(completed.usage).toMatchObject(expected); + expect(completed.cost?.estimatedUsd).toBeGreaterThan(budget); + expect(costs.at(-1)).toBe(completed.cost?.estimatedUsd); + }, + ); + + test("keeps live and persisted usage aligned across counter resets", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps live and persisted usage aligned across counter resets", + ) + ) { + return; + } + type Usage = Record & { + input_tokens: number; + output_tokens: number; + }; + const { PLUGIN_ROOT } = await import("./plugin-root.js"); + const { resolvePluginPython } = await import("../src/runtime.js"); + const python = await resolvePluginPython({ environment: process.env }); + const probe = [ + "import io, json, sys", + "from datetime import datetime, timezone", + "from pathlib import Path", + "from unittest.mock import patch", + "sys.path.insert(0, sys.argv[1])", + "import workbench_scan_usage as usage", + "from workbench_validation import parse_scan_cost", + "case = json.loads(sys.argv[2])", + "inherited = case.get('inherited')", + "fork = case['fork']", + "thread = 'scan-thread' if inherited is None else fork['threadId']", + "parent = None if inherited is None else 'scan-thread'", + "stamp = '2026-07-26T12:02:00Z'", + "def token_event(sample):", + " snapshot = {**sample, 'total_tokens': sample['input_tokens'] + sample['output_tokens']}", + " return {'type': 'event_msg', 'timestamp': stamp, 'payload': {'type': 'token_count', 'info': {'total_token_usage': snapshot}}}", + "def task_event(turn_id):", + " return {'type': 'event_msg', 'timestamp': stamp, 'payload': {'type': 'task_started', 'turn_id': turn_id, 'started_at': fork['startedAt']}}", + "metadata = {'id': thread, 'parent_thread_id': parent}", + "if case.get('forkOnly'):", + " del metadata['parent_thread_id']", + "if inherited is not None and not case.get('copiedHistory'):", + " metadata['forked_from_id'] = parent", + "events = [{'type': 'session_meta', 'payload': metadata}]", + "if inherited is not None:", + " if case.get('copiedHistory'):", + " events.append({'type': 'session_meta', 'payload': {'id': parent}})", + " events.extend([task_event(fork['inheritedTurnId']), token_event(inherited), task_event(fork['pendingTurnId']), token_event(inherited), task_event(fork['ownedTurnId'])])", + "for index, sample in enumerate(case['samples']):", + " if inherited is not None and index == 1:", + " events.append(task_event(fork['pendingTurnId']))", + " events.append(token_event(sample))", + "session = usage.RolloutSession(thread, parent, Path('fixture-rollout'))", + "rollout = ''.join(json.dumps(event) + '\\n' for event in events).encode('utf-8')", + "with patch.object(Path, 'open', return_value=io.BytesIO(rollout)):", + " measured, warnings = usage._read_rollout_usage(session, started_at=datetime(2026, 7, 26, 12, tzinfo=timezone.utc), completed_at=None)", + "parse_scan_cost(usage.measured_scan_cost_json({'coverage': 'complete', 'source': 'codex_rollout', 'threadCount': 1, **measured}))", + "print(json.dumps({'usage': measured, 'warnings': sorted(warnings)}))", + ].join("\n"); + const inheritedReset = { + inherited: { input_tokens: 1_000, output_tokens: 1_000 }, + samples: [ + { input_tokens: 1_500, output_tokens: 100 }, + { input_tokens: 1_600, output_tokens: 150 }, + { input_tokens: 1_600, output_tokens: 150 }, + ], + expected: { input_tokens: 1_600, output_tokens: 150 }, + }; + const cases: Array<{ + samples: Usage[]; + expected: Usage; + inherited?: Usage; + copiedHistory?: boolean; + forkOnly?: boolean; + limit?: number; + }> = [ + { + samples: [ + { input_tokens: 100, cached_input_tokens: 100, output_tokens: 100 }, + { input_tokens: 101, cached_input_tokens: 99, output_tokens: 99 }, + ], + expected: { + input_tokens: 101, + cached_input_tokens: 101, + output_tokens: 199, + }, + limit: 0.0015, + }, + { + samples: [ + { + input_tokens: 100, + cached_input_tokens: 90, + cache_write_input_tokens: 10, + output_tokens: 100, + reasoning_output_tokens: 100, + }, + { + input_tokens: 101, + cached_input_tokens: 91, + cache_write_input_tokens: 9, + output_tokens: 99, + reasoning_output_tokens: 99, + }, + ], + expected: { + input_tokens: 101, + cached_input_tokens: 90, + cache_write_input_tokens: 11, + output_tokens: 199, + reasoning_output_tokens: 199, + }, + limit: 0.0015, + }, + { + samples: [ + { + input_tokens: 100, + output_tokens: 100, + reasoning_output_tokens: 100, + }, + { input_tokens: 99, output_tokens: 101, reasoning_output_tokens: 99 }, + ], + expected: { + input_tokens: 199, + output_tokens: 101, + reasoning_output_tokens: 101, + }, + limit: 0.0015, + }, + { + samples: [ + { + input_tokens: 1_000, + cached_input_tokens: 400, + cache_write_input_tokens: 100, + output_tokens: 1_000, + reasoning_output_tokens: 800, + }, + { + input_tokens: 1_500, + cached_input_tokens: 600, + cache_write_input_tokens: 200, + output_tokens: 500, + reasoning_output_tokens: 50, + }, + { + input_tokens: 1_500, + cached_input_tokens: 600, + cache_write_input_tokens: 200, + output_tokens: 500, + reasoning_output_tokens: 50, + }, + ], + expected: { + input_tokens: 1_500, + cached_input_tokens: 600, + cache_write_input_tokens: 200, + output_tokens: 1_500, + reasoning_output_tokens: 850, + }, + limit: 0.017, + }, + { + samples: [ + { input_tokens: 1_000, output_tokens: 1_000 }, + { input_tokens: 1_500, output_tokens: 100 }, + { input_tokens: 1_500, output_tokens: 100 }, + ], + expected: { input_tokens: 2_500, output_tokens: 1_100 }, + limit: 0.017, + }, + { + samples: [ + { + input_tokens: 1_000, + cached_input_tokens: 400, + cache_write_input_tokens: 100, + output_tokens: 1_000, + reasoning_output_tokens: 800, + }, + { + input_tokens: 1_500, + cached_input_tokens: 600, + cache_write_input_tokens: 200, + output_tokens: 100, + reasoning_output_tokens: 50, + }, + ], + expected: { + input_tokens: 2_500, + cached_input_tokens: 1_000, + cache_write_input_tokens: 300, + output_tokens: 1_100, + reasoning_output_tokens: 850, + }, + }, + { + samples: [ + { + input_tokens: 100, + cached_input_tokens: 20, + cache_write_input_tokens: 10, + output_tokens: 50, + reasoning_output_tokens: 10, + }, + { + input_tokens: 150, + cached_input_tokens: 10, + cache_write_input_tokens: 20, + output_tokens: 30, + reasoning_output_tokens: 5, + }, + ], + expected: { + input_tokens: 150, + cached_input_tokens: 30, + cache_write_input_tokens: 20, + output_tokens: 80, + reasoning_output_tokens: 15, + }, + }, + { + samples: [ + { input_tokens: 1_000, output_tokens: 1_000 }, + { input_tokens: 0, output_tokens: 0 }, + { input_tokens: 1_500, output_tokens: 100 }, + ], + expected: { input_tokens: 2_500, output_tokens: 1_100 }, + }, + inheritedReset, + { ...inheritedReset, forkOnly: true }, + { ...inheritedReset, copiedHistory: true }, + { + samples: [ + { input_tokens: 100, output_tokens: 100 }, + { input_tokens: 110, output_tokens: 90 }, + { input_tokens: 120, output_tokens: 100 }, + ], + expected: { input_tokens: 120, output_tokens: 200 }, + }, + ]; + for (const { + samples, + inherited, + copiedHistory, + forkOnly, + expected: counters, + limit, + } of cases) { + const expected = { + cached_input_tokens: 0, + cache_write_input_tokens: 0, + reasoning_output_tokens: 0, + ...counters, + total_tokens: counters.input_tokens + counters.output_tokens, + }; + const persisted = spawnSync( + python, + [ + "-I", + "-B", + "-c", + probe, + join(PLUGIN_ROOT, "scripts"), + JSON.stringify({ + samples, + inherited, + copiedHistory, + forkOnly, + fork: accountingFork, + }), + ], + { encoding: "utf8" }, + ); + expect(persisted.status, persisted.stderr).toBe(0); + expect(JSON.parse(persisted.stdout)).toEqual({ + usage: { + inputTokens: expected.input_tokens, + cachedInputTokens: expected.cached_input_tokens, + cacheWriteInputTokens: expected.cache_write_input_tokens, + outputTokens: expected.output_tokens, + reasoningOutputTokens: expected.reasoning_output_tokens, + totalTokens: expected.total_tokens, + }, + warnings: [], + }); + const taskStart = (turnId: string) => + accountingTaskStart(turnId, accountingFork.startedAt); + const events: MockAccountingEvent[] = []; + if (inherited !== undefined) { + if (copiedHistory) { + events.push({ + type: "session_meta", + payload: { id: "scan-thread" }, + }); + } + events.push( + taskStart(accountingFork.inheritedTurnId), + accountingEvent(inherited), + taskStart(accountingFork.pendingTurnId), + accountingEvent(inherited), + taskStart(accountingFork.ownedTurnId), + ); + } + for (const [index, sample] of samples.entries()) { + if (inherited !== undefined && index === 1) { + events.push(taskStart(accountingFork.pendingTurnId)); + } + events.push(accountingEvent(sample)); + } + const metadata = { timestamp: accountingFork.timestamp }; + const sessions: Record = + inherited === undefined + ? { + "scan-thread": accountingSession( + "scan-thread", + events, + undefined, + metadata, + ), + } + : { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent({ input_tokens: 0, output_tokens: 0 }), + ]), + [accountingFork.threadId]: accountingSession( + accountingFork.threadId, + events, + forkOnly ? undefined : "scan-thread", + { + ...metadata, + ...(copiedHistory ? {} : { forked_from_id: "scan-thread" }), + }, + ), + }; + const costs: number[] = []; + await withMockAccountingSessions( + sessions, + { + model: "gpt-5.6-terra", + maxCostUsd: limit ?? 1, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker) => { + const final = await tracker.stop(undefined); + expect(final.usage).toEqual(expected); + expect(final.cost).toEqual( + estimateScanCost("gpt-5.6-terra", expected), + ); + expect(costs.at(-1)).toBe(final.cost!.estimatedUsd); + if (limit !== undefined) { + expect(final.cost!.estimatedUsd).toBeGreaterThan(limit); + } + expect(await tracker.stop()).toBe(final); + }, + ); + } + }); + + test.each([ + ["fork-only delegated worker", { forked_from_id: "scan-thread" }, true], + [ + "nested spawn parent precedence", + { + source: { + subagent: { thread_spawn: { parent_thread_id: "scan-thread" } }, + }, + parent_thread_id: "unrelated-thread", + forked_from_id: "another-thread", + }, + true, + ], + [ + "direct parent fallback", + { + source: { subagent: { thread_spawn: { parent_thread_id: "" } } }, + parent_thread_id: "scan-thread", + forked_from_id: "unrelated-thread", + }, + true, + ], + [ + "fork parent fallback after invalid candidates", + { + source: { subagent: { thread_spawn: { parent_thread_id: 7 } } }, + parent_thread_id: "", + forked_from_id: "scan-thread", + }, + true, + ], + [ + "fork parent fallback after non-string direct parent", + { parent_thread_id: 7, forked_from_id: "scan-thread" }, + true, + ], + ["unrelated fork parent", { forked_from_id: "unrelated-thread" }, false], + ["empty fork parent", { forked_from_id: "" }, false], + ["non-string fork parent", { forked_from_id: 7 }, false], + ] as const)( + "tracks the proactive standard-scan budget for a %s", + async (_description, parentMetadata, included) => { + const home = await codexHome(); + const repository = join(home, "repository"); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const inherited = { input_tokens: 1_000, output_tokens: 100 }; + const owned = { input_tokens: 1_300, output_tokens: 130 }; + await writeSession(home, "scan-thread", rootUsage); + const worker = await writeSession(home, accountingFork.threadId, null); + await writeFile( + worker, + `${[ + { + type: "session_meta", + payload: { + id: accountingFork.threadId, + timestamp: accountingFork.timestamp, + cwd: repository, + ...parentMetadata, + }, + }, + accountingTaskStart( + accountingFork.inheritedTurnId, + accountingFork.startedAt, + ), + accountingEvent(inherited), + accountingTaskStart( + accountingFork.ownedTurnId, + accountingFork.startedAt, + ), + accountingEvent(owned), + ] + .map((event) => JSON.stringify(event)) + .join("\n")}\n`, + ); + const maxCostUsd = 0.001; + const exceeded = new AbortController(); + const costs: number[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + repository, + scanDirectory: join(home, "scan"), + maxCostUsd, + onCost: (cost) => { + costs.push(cost.estimatedUsd); + if (cost.estimatedUsd > maxCostUsd) exceeded.abort(); + }, + }); + tracker.start("scan-thread"); + + expect((await tracker.refresh()).cost).toMatchObject( + included + ? { inputTokens: 400, outputTokens: 40, estimatedUsd: 0.00128 } + : { inputTokens: 100, outputTokens: 10, estimatedUsd: 0.00032 }, + ); + expect(exceeded.signal.aborted).toBe(included); + expect(costs.at(-1)).toBe(included ? 0.00128 : 0.00032); + await tracker.stop(); + }, + ); + + test("keeps same-millisecond inherited UUID7 turns below an explicit worker budget", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps same-millisecond inherited UUID7 turns below an explicit worker budget", + ) + ) { + return; + } + const rootUsage = { input_tokens: 0, output_tokens: 0 }; + const inherited = { input_tokens: 1_000, output_tokens: 100 }; + const final = { input_tokens: 1_100, output_tokens: 110 }; + const observedCosts: number[] = []; + const maxCostUsd = 0.001; + + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + [accountingFork.threadId]: accountingSession( + accountingFork.threadId, + [ + accountingTaskStart( + accountingFork.inheritedTurnId, + accountingFork.startedAt, + ), + accountingEvent(inherited), + accountingTaskStart( + accountingFork.ownedTurnId, + accountingFork.startedAt, + ), + accountingEvent(final), + ], + "scan-thread", + { + timestamp: accountingFork.timestamp, + forked_from_id: "scan-thread", + }, + ), + }, + { + model: "gpt-5.6-sol", + maxCostUsd, + onCost: (cost) => observedCosts.push(cost.estimatedUsd), + }, + async (tracker) => { + await expect(tracker.stop(rootUsage)).resolves.toMatchObject({ + usage: { input_tokens: 100, output_tokens: 10 }, + cost: { estimatedUsd: 0.0008 }, + }); + expect(observedCosts).toEqual([0.0008]); + expect(observedCosts.every((cost) => cost < maxCostUsd)).toBe(true); + }, + ); + }); + + test("keeps unstarted fork usage out of budget totals", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps unstarted fork usage out of budget totals", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const inherited = { input_tokens: 1_000, output_tokens: 1_000 }; + for (const maxCostUsd of [undefined, 0.01]) { + const costs: number[] = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + [accountingFork.threadId]: accountingSession( + accountingFork.threadId, + [ + accountingTaskStart( + accountingFork.inheritedTurnId, + accountingFork.startedAt, + ), + accountingEvent(inherited), + accountingTaskStart( + accountingFork.pendingTurnId, + accountingFork.startedAt, + ), + accountingEvent(inherited), + ], + "scan-thread", + { + timestamp: accountingFork.timestamp, + forked_from_id: "scan-thread", + }, + ), + }, + { + model: "gpt-5.6-terra", + maxCostUsd, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker) => { + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + const final = tracker.stop(rootUsage); + if (maxCostUsd === undefined) { + await expect(final).resolves.toMatchObject({ + cost: { inputTokens: 100, outputTokens: 10 }, + }); + } else { + await expect(final).rejects.toThrow( + "The scan cost limit could not be verified", + ); + } + expect(costs.every((cost) => cost < 0.01)).toBe(true); + }, + ); + } + }); + + test("counts ordinary UUID7 workers without a fork boundary", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "counts ordinary UUID7 workers without a fork boundary", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + [accountingFork.threadId]: accountingSession( + accountingFork.threadId, + [accountingEvent({ input_tokens: 250, output_tokens: 25 })], + "scan-thread", + { timestamp: accountingFork.timestamp }, + ), + }, + { model: "gpt-5.6-terra", maxCostUsd: 1 }, + async (tracker) => { + await expect(tracker.stop(rootUsage)).resolves.toMatchObject({ + cost: { inputTokens: 350, outputTokens: 35 }, + }); + }, + ); + }); + + test("keeps mocked session-detail replay separate from accounting", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps mocked session-detail replay separate from accounting", + ) + ) { + return; + } + const initial = { input_tokens: 100, output_tokens: 10 }; + const later = { input_tokens: 150, output_tokens: 15 }; + const events: ScanSessionEvent[] = []; + const costs: number[] = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(initial), + ]), + }, + { + model: "gpt-5.6-terra", + maxCostUsd: 1, + onSessionEvent: (event) => events.push(event), + onCost: (cost) => costs.push(cost.inputTokens), + }, + async (tracker) => { + expect((await tracker.stop(later)).cost?.inputTokens).toBe(150); + expect(costs).toEqual([100, 150]); + expect(events).toHaveLength(5); + expect( + events + .map(({ event }) => event["payload"] as Record) + .filter((payload) => payload["type"] === "token_count") + .map((payload) => payload["info"]), + ).toEqual([ + { total_token_usage: initial }, + { total_token_usage: later }, + ]); + }, + (_path, attempt, append) => { + if (attempt === 2) { + append("scan-thread", [ + accountingEvent(later), + { type: "event_msg", payload: { type: "task_complete" } }, + ]); + } + }, + ); + + for (const accountingFailed of [false, true]) { + const costs: number[] = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + ...(accountingFailed + ? [new Error("Synthetic decoded record error")] + : []), + accountingEvent(initial), + ]), + }, + { + model: "gpt-5.6-terra", + maxCostUsd: 1, + onSessionEvent: () => {}, + onCost: (cost) => costs.push(cost.inputTokens), + }, + async (tracker) => { + const refresh = tracker.refresh(); + if (accountingFailed) { + await expect(refresh).rejects.toThrow( + "tracked session record could not be read", + ); + await expect(tracker.stop()).rejects.toThrow( + "tracked session record could not be read", + ); + } else { + await expect(refresh).resolves.toMatchObject({ + cost: { inputTokens: 100 }, + }); + await expect(tracker.stop(initial)).resolves.toMatchObject({ + cost: { inputTokens: 100 }, + }); + } + expect(costs).toEqual([100]); + }, + (_path, attempt) => { + if (attempt === 2) throw new Error("Synthetic detail replay failure"); + }, + ); + } + }); + + test("prefers completed root usage on mocked equal-price ties", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "prefers completed root usage on mocked equal-price ties", + ) + ) { + return; + } + const model = "gpt-5.6-terra"; + const observed = { + input_tokens: 100, + output_tokens: 1, + reasoning_output_tokens: 1, + }; + const completed = { + input_tokens: 101, + cached_input_tokens: 0, + cache_write_input_tokens: 20, + output_tokens: 0, + reasoning_output_tokens: 0, + }; + expect(estimateScanCost(model, observed)?.estimatedUsd).toBe( + estimateScanCost(model, completed)?.estimatedUsd, + ); + for (const withWorker of [false, true]) { + const sessions: Record = { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(observed), + ]), + }; + if (withWorker) { + sessions["worker-thread"] = accountingSession( + "worker-thread", + [ + accountingEvent({ + input_tokens: 20, + output_tokens: 2, + reasoning_output_tokens: 1, + }), + ], + "scan-thread", + ); + } + const expected = { + ...completed, + input_tokens: withWorker ? 121 : 101, + output_tokens: withWorker ? 2 : 0, + reasoning_output_tokens: withWorker ? 1 : 0, + }; + await withMockAccountingSessions( + sessions, + { model, maxCostUsd: 1 }, + async (tracker) => { + const final = await tracker.stop(completed); + expect(final.usage).toMatchObject(expected); + expect(final.cost).toEqual(estimateScanCost(model, expected)); + expect(await tracker.stop()).toBe(final); + }, + ); + } + }); + + test("recovers mocked root read failures only with verified workers", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "recovers mocked root read failures only with verified workers", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const workerUsage = { input_tokens: 250, output_tokens: 20 }; + const completedRoot = { input_tokens: 1_000, output_tokens: 100 }; + for (const workerState of [ + "complete", + "unreadable", + "accounting-error", + "missing", + "unfinished", + "missing-usage", + ] as const) { + let failReads = false; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "worker-thread": accountingSession( + "worker-thread", + workerState === "missing-usage" + ? [] + : [accountingEvent(workerUsage)], + "scan-thread", + ), + }, + { model: "gpt-5.6-terra", maxCostUsd: 1 }, + async (tracker, append, omit) => { + await tracker.refresh(); + if (workerState === "accounting-error") { + append("worker-thread", [ + new Error("Synthetic worker accounting error"), + { type: "event_msg", payload: { type: "task_complete" } }, + ]); + } else if (workerState === "missing") { + omit("worker-thread"); + } else if (workerState === "unfinished") { + append("worker-thread", [ + { type: "event_msg", payload: { type: "task_started" } }, + ]); + } + failReads = true; + const final = tracker.stop(completedRoot); + if (workerState === "complete") { + await expect(final).resolves.toMatchObject({ + usage: { input_tokens: 1_250, output_tokens: 120 }, + cost: { inputTokens: 1_250, outputTokens: 120 }, + }); + } else { + await expect(final).rejects.toThrow(); + } + }, + (path) => { + if ( + failReads && + (path.endsWith("rollout-scan-thread.jsonl") || + (workerState === "unreadable" && + path.endsWith("rollout-worker-thread.jsonl"))) + ) { + throw new Error("Synthetic session read failure"); + } + }, + ); + } + }); + + test("accumulates mocked owned reset epochs without double counting", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "accumulates mocked owned reset epochs without double counting", + ) + ) { + return; + } + const completed = { type: "event_msg", payload: { type: "task_complete" } }; + const initialRoot = { input_tokens: 100, output_tokens: 10 }; + const loggedRoot = { input_tokens: 280, output_tokens: 28 }; + const inherited = { + input_tokens: 1_000, + cached_input_tokens: 400, + cache_write_input_tokens: 100, + output_tokens: 100, + reasoning_output_tokens: 20, + }; + const latestWorker = { + input_tokens: 140, + cached_input_tokens: 30, + cache_write_input_tokens: 15, + output_tokens: 25, + reasoning_output_tokens: 7, + }; + const withWorkers = (root: typeof initialRoot) => ({ + input_tokens: root.input_tokens + 340, + cached_input_tokens: 80, + cache_write_input_tokens: 35, + output_tokens: root.output_tokens + 65, + reasoning_output_tokens: 17, + total_tokens: root.input_tokens + root.output_tokens + 405, + }); + for (const [authoritative, expectedRoot] of [ + [{ input_tokens: 200, output_tokens: 20 }, loggedRoot], + [ + { input_tokens: 400, output_tokens: 40 }, + { input_tokens: 400, output_tokens: 40 }, + ], + ] as const) { + const costs: number[] = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(initialRoot), + ]), + "worker-thread": accountingSession( + "worker-thread", + [ + { type: "session_meta", payload: { id: "inherited-thread" } }, + accountingEvent(inherited), + { + type: "event_msg", + payload: { type: "task_started", started_at: 1_785_067_320 }, + }, + accountingEvent({ + input_tokens: 1_200, + cached_input_tokens: 450, + cache_write_input_tokens: 120, + output_tokens: 140, + reasoning_output_tokens: 30, + }), + ], + "scan-thread", + { timestamp: "2026-07-26T12:02:00Z" }, + ), + }, + { + model: "gpt-5.6-terra", + maxCostUsd: 1, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker, append) => { + const initial = await tracker.stop(); + expect(initial.usage).toMatchObject({ + input_tokens: 300, + output_tokens: 50, + }); + append("scan-thread", [ + accountingEvent({ input_tokens: 200, output_tokens: 20 }), + accountingEvent({ input_tokens: 50, output_tokens: 5 }), + accountingEvent({ input_tokens: 80, output_tokens: 8 }), + accountingEvent({ input_tokens: 80, output_tokens: 8 }), + completed, + ]); + append("worker-thread", [ + accountingEvent({ + input_tokens: 100, + cached_input_tokens: 20, + cache_write_input_tokens: 10, + output_tokens: 20, + reasoning_output_tokens: 5, + }), + accountingEvent(latestWorker), + completed, + ]); + const current = await tracker.refresh(); + expect(current.usage).toEqual(withWorkers(loggedRoot)); + expect(await tracker.refresh()).toEqual(current); + + append("worker-thread", [accountingEvent(latestWorker)]); + await expect(tracker.stop(authoritative)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + append("worker-thread", [completed]); + const final = await tracker.stop(authoritative); + const expected = withWorkers(expectedRoot); + expect(final.usage).toEqual(expected); + expect(final.cost).toEqual( + estimateScanCost("gpt-5.6-terra", expected), + ); + expect(await tracker.stop()).toBe(final); + expect(costs).toEqual([ + initial.cost!.estimatedUsd, + current.cost!.estimatedUsd, + ...(final.cost!.estimatedUsd === current.cost!.estimatedUsd + ? [] + : [final.cost!.estimatedUsd]), + ]); + }, + ); + } + }); + + test("keeps mocked parse errors scoped to included budgeted sessions", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps mocked parse errors scoped to included budgeted sessions", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const workerUsage = { input_tokens: 200, output_tokens: 20 }; + for (const [ + failedSession, + withWorker, + maxCostUsd, + authoritative, + rejects, + ] of [ + ["worker-thread", true, 1, true, true], + ["scan-thread", true, 1, true, false], + ["scan-thread", true, 1, false, true], + ["unrelated-thread", true, 1, true, false], + ["worker-thread", true, undefined, true, false], + ["scan-thread", false, undefined, false, false], + ["scan-thread", false, 1, true, false], + ["scan-thread", false, 1, false, true], + ] as const) { + const sessions: Record = { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "unrelated-thread": accountingSession("unrelated-thread", [ + accountingEvent(workerUsage), + ]), + }; + if (withWorker) { + sessions["worker-thread"] = accountingSession( + "worker-thread", + [accountingEvent(workerUsage)], + "scan-thread", + ); + } + sessions[failedSession]!.unshift( + new SyntaxError("mock parser diagnostic"), + ); + const costs: number[] = []; + await withMockAccountingSessions( + sessions, + { + model: "gpt-5.6-terra", + maxCostUsd, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker) => { + const refresh = tracker.refresh(); + if ( + maxCostUsd !== undefined && + failedSession !== "unrelated-thread" + ) { + await expect(refresh).rejects.toThrow( + "tracked session record could not be read", + ); + } else { + await expect(refresh).resolves.toMatchObject({ + cost: { inputTokens: withWorker ? 300 : 100 }, + }); + } + const completed = tracker.stop(authoritative ? rootUsage : undefined); + if (rejects) { + await expect(completed).rejects.toThrow( + "tracked session record could not be read", + ); + } else { + await expect(completed).resolves.toMatchObject({ + cost: { inputTokens: withWorker ? 300 : 100 }, + }); + } + expect(costs).toEqual([withWorker ? 0.00096 : 0.00032]); + }, + ); + } + }); + + test("retains mocked per-session priced high-water snapshots", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "retains mocked per-session priced high-water snapshots", + ) + ) { + return; + } + type Usage = Record & { + input_tokens: number; + output_tokens: number; + }; + const reclassified: Usage[] = [ + { input_tokens: 100, cached_input_tokens: 100, output_tokens: 0 }, + { input_tokens: 100, cached_input_tokens: 0, output_tokens: 0 }, + { input_tokens: 90, cached_input_tokens: 0, output_tokens: 10 }, + { input_tokens: 100, cached_input_tokens: 0, output_tokens: 10 }, + ]; + const equalPrice: Usage[] = [ + { input_tokens: 100, output_tokens: 1 }, + { + input_tokens: 101, + cache_write_input_tokens: 20, + output_tokens: 0, + }, + ]; + const cases: Array<{ + model: string; + samples: Usage[]; + expected: Usage; + inherited?: Usage; + }> = [ + { + model: "gpt-5.6-terra", + samples: [ + { + input_tokens: 100, + cached_input_tokens: 20, + cache_write_input_tokens: 10, + output_tokens: 50, + reasoning_output_tokens: 10, + }, + { + input_tokens: 150, + cached_input_tokens: 10, + cache_write_input_tokens: 20, + output_tokens: 30, + reasoning_output_tokens: 5, + }, + ], + expected: { + input_tokens: 150, + cached_input_tokens: 30, + cache_write_input_tokens: 20, + output_tokens: 80, + reasoning_output_tokens: 15, + }, + }, + { + model: "gpt-5.6-terra", + samples: reclassified, + expected: { + input_tokens: 200, + cached_input_tokens: 100, + output_tokens: 10, + }, + }, + { + model: "unknown-model", + samples: reclassified, + expected: { + input_tokens: 200, + cached_input_tokens: 100, + output_tokens: 10, + }, + }, + { + model: "gpt-5.6-terra", + samples: equalPrice, + expected: { + input_tokens: 101, + cache_write_input_tokens: 1, + output_tokens: 1, + }, + }, + { + model: "unknown-model", + samples: equalPrice, + expected: { + input_tokens: 101, + cache_write_input_tokens: 1, + output_tokens: 1, + }, + }, + { + model: "unknown-model", + samples: [ + { input_tokens: Number.MAX_SAFE_INTEGER, output_tokens: 2 }, + { input_tokens: Number.MAX_SAFE_INTEGER, output_tokens: 1 }, + ], + // The full reset would overflow input, so retain the last valid whole. + expected: { + input_tokens: Number.MAX_SAFE_INTEGER, + output_tokens: 2, + }, + }, + { + model: "unknown-model", + samples: [ + { input_tokens: Number.MAX_SAFE_INTEGER - 5, output_tokens: 0 }, + { input_tokens: 10, output_tokens: 0 }, + { input_tokens: 11, output_tokens: 0 }, + ], + expected: { + input_tokens: Number.MAX_SAFE_INTEGER - 4, + output_tokens: 0, + }, + }, + { + model: "unknown-model", + inherited: { + input_tokens: Number.MAX_SAFE_INTEGER, + cached_input_tokens: Number.MAX_SAFE_INTEGER, + output_tokens: 0, + }, + samples: [ + { + input_tokens: Number.MAX_SAFE_INTEGER - 1, + cached_input_tokens: Number.MAX_SAFE_INTEGER - 2, + output_tokens: 0, + }, + { + input_tokens: Number.MAX_SAFE_INTEGER - 1, + cached_input_tokens: 3, + output_tokens: 1, + }, + { + input_tokens: Number.MAX_SAFE_INTEGER - 1, + cached_input_tokens: 3, + output_tokens: 2, + }, + ], + expected: { + input_tokens: Number.MAX_SAFE_INTEGER - 1, + cached_input_tokens: Number.MAX_SAFE_INTEGER - 2, + output_tokens: 2, + }, + }, + ]; + for (const { model, samples, expected, inherited } of cases) { + const events: MockAccountingEvent[] = [ + ...(inherited === undefined + ? [] + : [ + { + type: "session_meta", + payload: { id: "inherited-thread" }, + }, + accountingEvent(inherited), + { + type: "event_msg", + payload: { type: "task_started", started_at: 1_785_067_320 }, + }, + ]), + ...samples.map(accountingEvent), + ]; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", events, undefined, { + timestamp: "2026-07-26T12:02:00Z", + }), + }, + { + model, + ...(model === "unknown-model" ? {} : { maxCostUsd: 1 }), + }, + async (tracker) => { + const final = await tracker.stop(undefined); + expect(final.usage).toEqual({ + cached_input_tokens: 0, + cache_write_input_tokens: 0, + reasoning_output_tokens: 0, + ...expected, + total_tokens: expected.input_tokens + expected.output_tokens, + }); + expect(final.cost).toEqual(estimateScanCost(model, expected)); + expect(await tracker.stop()).toBe(final); + }, + ); + } + }); + + test("keeps mocked unverified evidence separate from known cost floors", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps mocked unverified evidence separate from known cost floors", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const unpriceable = accountingEvent({ + input_tokens: Number.MAX_SAFE_INTEGER, + output_tokens: 0, + }); + for (const [evidence, expectedInput, expectedCost] of [ + [new SyntaxError("mock parser diagnostic"), 1_200, 0.00384], + [unpriceable, 1_100, 0.00352], + ] as const) { + const reports: Array = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "worker-thread": accountingSession( + "worker-thread", + [ + accountingEvent({ input_tokens: 1_000, output_tokens: 100 }), + evidence, + accountingEvent(rootUsage), + ], + "scan-thread", + ), + }, + { + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + onCost: (cost) => reports.push(cost.estimatedUsd), + onError: () => reports.push("error"), + }, + async (tracker) => { + await expect(tracker.stop(rootUsage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: expectedInput, + outputTokens: expectedInput / 10, + estimatedUsd: expectedCost, + }); + expect(reports[0]).toBe(expectedCost); + expect(reports).toContain("error"); + }, + ); + } + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + unpriceable, + ]), + }, + { model: "gpt-5.6-terra", maxCostUsd: 1 }, + async (tracker) => { + expect( + (await tracker.stop({ input_tokens: 1_000, output_tokens: 100 })).cost + ?.estimatedUsd, + ).toBe(0.0032); + }, + ); + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + accountingEvent({ input_tokens: 200, output_tokens: 20 }), + ]), + }, + { model: "unknown-model" }, + async (tracker) => { + expect(await tracker.stop(undefined)).toMatchObject({ + usage: { input_tokens: 200, output_tokens: 20 }, + cost: null, + }); + }, + ); + }); + + test("keeps mocked unusable own usage scoped to accounting", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps mocked unusable own usage scoped to accounting", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const inherited = { input_tokens: 1_000, output_tokens: 100 }; + const high = accountingEvent({ input_tokens: 2_000, output_tokens: 200 }); + const low = accountingEvent({ input_tokens: 1_100, output_tokens: 110 }); + const metadata = { timestamp: "2026-07-26T12:02:00Z" }; + const history = [ + { type: "session_meta", payload: { id: "inherited-thread" } }, + accountingEvent(inherited), + ]; + const started = { + type: "event_msg", + payload: { type: "task_started", started_at: 1_785_067_320 }, + }; + const infoOnly = { + type: "event_msg", + payload: { type: "token_count", info: null }, + }; + for (const [unavailable, invalid] of [ + [accountingEvent(null), true], + [accountingEvent({ input_tokens: 900, output_tokens: 90 }), false], + ] as const) { + for (const [scope, maxCostUsd] of [ + ["worker", 0.001], + ["first-error", 0.001], + ["worker", undefined], + ["unrelated", 1], + ["root", 1], + ["replay", 1], + ] as const) { + const rejects = + scope === "first-error" || + (scope === "worker" && maxCostUsd !== undefined && invalid); + const expectedInput = + scope === "unrelated" ? 100 : scope === "root" ? 2_100 : 2_200; + const events = + scope === "replay" + ? [ + ...history, + unavailable, + accountingEvent(inherited), + started, + high, + infoOnly, + low, + ] + : [ + ...history, + started, + high, + ...(scope === "first-error" + ? [new SyntaxError("mock existing accounting diagnostic")] + : []), + unavailable, + low, + ]; + const sessions: Readonly< + Record + > = + scope === "root" + ? { + "scan-thread": accountingSession( + "scan-thread", + events, + undefined, + metadata, + ), + } + : { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "worker-thread": accountingSession( + "worker-thread", + events, + scope === "unrelated" ? undefined : "scan-thread", + metadata, + ), + }; + const costs: number[] = []; + await withMockAccountingSessions( + sessions, + { + model: "gpt-5.6-terra", + maxCostUsd, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker) => { + const completed = tracker.stop( + scope === "root" + ? { input_tokens: 1_200, output_tokens: 120 } + : rootUsage, + ); + if (rejects) { + await expect(completed).rejects.toThrow( + scope === "first-error" + ? "tracked session record could not be read" + : "model pricing or token usage is unavailable", + ); + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 2_200, + outputTokens: 220, + estimatedUsd: 0.00704, + }); + expect(costs[0]).toBe(0.00704); + } else { + expect((await completed).cost).toMatchObject({ + inputTokens: expectedInput, + outputTokens: expectedInput / 10, + }); + } + }, + ); + } + } + }); + + test("keeps mocked invalid accumulated usage fail closed", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "keeps mocked invalid accumulated usage fail closed", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const inherited = { + input_tokens: 1_000, + cached_input_tokens: 800, + output_tokens: 0, + }; + for (const [ + middle, + latest, + expectedInput, + expectedCached, + invalidSnapshot, + ] of [ + [ + accountingEvent(null), + { input_tokens: 120, cached_input_tokens: 80, output_tokens: 0 }, + 120, + 80, + true, + ], + [ + accountingEvent({ + input_tokens: 110, + cached_input_tokens: 50, + output_tokens: 0, + }), + { input_tokens: 160, cached_input_tokens: 50, output_tokens: 0 }, + 160, + 90, + false, + ], + ] as const) { + for (const [scope, maxCostUsd] of [ + ["worker", 0.0001], + ["worker", undefined], + ["unrelated", 1], + ["root", 1], + ] as const) { + const events: MockAccountingEvent[] = [ + { type: "session_meta", payload: { id: "inherited-thread" } }, + accountingEvent(inherited), + { + type: "event_msg", + payload: { type: "task_started", started_at: 1_785_067_320 }, + }, + accountingEvent({ + input_tokens: 100, + cached_input_tokens: 80, + output_tokens: 0, + }), + middle, + accountingEvent(latest), + ]; + const metadata = { timestamp: "2026-07-26T12:02:00Z" }; + const sessions: Record = + scope === "root" + ? { + "scan-thread": accountingSession( + "scan-thread", + events, + undefined, + metadata, + ), + } + : { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "worker-thread": accountingSession( + "worker-thread", + events, + scope === "unrelated" ? undefined : "scan-thread", + metadata, + ), + }; + const authoritative = + scope === "root" + ? { input_tokens: 200, output_tokens: 20 } + : rootUsage; + const expected = + scope === "root" || scope === "unrelated" + ? authoritative + : { + input_tokens: expectedInput + rootUsage.input_tokens, + cached_input_tokens: expectedCached, + output_tokens: rootUsage.output_tokens, + }; + const expectedCost = estimateScanCost("gpt-5.6-terra", expected); + const reports: number[] = []; + await withMockAccountingSessions( + sessions, + { + model: "gpt-5.6-terra", + maxCostUsd, + onCost: (cost) => reports.push(cost.estimatedUsd), + }, + async (tracker) => { + const final = tracker.stop(authoritative); + if ( + scope === "worker" && + maxCostUsd !== undefined && + invalidSnapshot + ) { + await expect(final).rejects.toThrow( + "model pricing or token usage is unavailable", + ); + expect((await tracker.stop()).cost).toEqual(expectedCost); + expect(reports[0]).toBe(expectedCost!.estimatedUsd); + await expect(tracker.stop(authoritative)).rejects.toThrow( + "model pricing or token usage is unavailable", + ); + } else { + await expect(final).resolves.toMatchObject({ + usage: expected, + cost: expectedCost, + }); + } + }, + ); + } + } + }); + + test("reports mocked readable costs before rejecting missing sessions", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "reports mocked readable costs before rejecting missing sessions", + ) + ) { + return; + } + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const freshUsage = { input_tokens: 1_000, output_tokens: 100 }; + const finalCost = { + inputTokens: 1_200, + outputTokens: 120, + estimatedUsd: 0.00384, + }; + const missingMessage = + "A tracked scan session disappeared before its cost could be verified."; + for (const [missingSession, maxCostUsd, rejects] of [ + ["scan-thread", 0.001, true], + ["retained-worker", 0.001, true], + ["unrelated-thread", 1, false], + ["retained-worker", undefined, false], + ] as const) { + const reports: Array = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + "retained-worker": accountingSession( + "retained-worker", + [accountingEvent(rootUsage)], + "scan-thread", + ), + "growing-worker": accountingSession( + "growing-worker", + [accountingEvent(rootUsage)], + "scan-thread", + ), + "unrelated-thread": accountingSession("unrelated-thread", [ + accountingEvent(freshUsage), + ]), + }, + { + model: "gpt-5.6-terra", + maxCostUsd, + onCost: (cost) => reports.push(cost.estimatedUsd), + }, + async (tracker, append, omit) => { + expect((await tracker.stop()).cost?.estimatedUsd).toBe(0.00096); + omit(missingSession); + append("growing-worker", [ + accountingEvent(freshUsage), + { type: "event_msg", payload: { type: "task_complete" } }, + ]); + const refreshed = tracker.refresh(); + if (rejects) { + await expect( + refreshed.catch((error: unknown) => { + reports.push("rejected"); + throw error; + }), + ).rejects.toThrow(missingMessage); + expect(reports).toEqual([0.00096, 0.00384, "rejected"]); + const cleanup = await tracker.stop(); + expect(cleanup.cost).toMatchObject(finalCost); + await expect(tracker.stop(rootUsage)).rejects.toThrow( + missingMessage, + ); + expect(await tracker.stop()).toEqual(cleanup); + } else { + await expect(refreshed).resolves.toMatchObject({ cost: finalCost }); + await expect(tracker.stop(rootUsage)).resolves.toMatchObject({ + cost: finalCost, + }); + expect(reports).toEqual([0.00096, 0.00384]); + } + }, + ); + } + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(rootUsage), + ]), + }, + { model: "gpt-5.6-terra", maxCostUsd: 1 }, + async (tracker, _append, omit) => { + expect((await tracker.stop()).cost?.estimatedUsd).toBe(0.00032); + omit("scan-thread"); + await expect(tracker.stop(rootUsage)).rejects.toThrow(missingMessage); + }, + ); + }); + + test("prefers mocked completed usage when prices are unavailable", async () => { + if ( + runTestInSubprocess( + import.meta.path, + "prefers mocked completed usage when prices are unavailable", + ) + ) { + return; + } + const observedRoot = { input_tokens: 100, output_tokens: 10 }; + const workerUsage = { input_tokens: 200, output_tokens: 20 }; + const completedRoot = { input_tokens: 1_000, output_tokens: 100 }; + const equalTotalRoot = { input_tokens: 900, output_tokens: 200 }; + const boundaryObserved = { + input_tokens: Number.MAX_SAFE_INTEGER, + output_tokens: 2, + }; + const boundaryCompleted = { ...boundaryObserved, output_tokens: 1 }; + expect(boundaryObserved.input_tokens + boundaryObserved.output_tokens).toBe( + boundaryCompleted.input_tokens + boundaryCompleted.output_tokens, + ); + for (const [observed, completed, expected, withWorker, refreshFails] of [ + [observedRoot, completedRoot, completedRoot, false, true], + [observedRoot, completedRoot, completedRoot, true, true], + [observedRoot, completedRoot, completedRoot, true, false], + [completedRoot, observedRoot, completedRoot, false, true], + [completedRoot, observedRoot, completedRoot, true, true], + [completedRoot, equalTotalRoot, equalTotalRoot, false, true], + [completedRoot, equalTotalRoot, equalTotalRoot, true, true], + [boundaryObserved, boundaryCompleted, boundaryObserved, false, true], + ] as const) { + const sessions: Record = { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(observed), + ]), + }; + if (withWorker) { + sessions["worker-thread"] = accountingSession( + "worker-thread", + [accountingEvent(workerUsage)], + "scan-thread", + ); + } + const reportedErrors: unknown[] = []; + const reportedCosts: number[] = []; + const refreshError = new Error("Mock final accounting refresh failed."); + await withMockAccountingSessions( + sessions, + { + model: "unknown-model", + onCost: (cost) => reportedCosts.push(cost.estimatedUsd), + onError: (error) => reportedErrors.push(error), + }, + async (tracker) => { + expect((await tracker.stop()).usage).toMatchObject({ + input_tokens: + observed.input_tokens + + (withWorker ? workerUsage.input_tokens : 0), + output_tokens: + observed.output_tokens + + (withWorker ? workerUsage.output_tokens : 0), + }); + const refresh = refreshFails + ? spyOn(tracker, "refresh").mockRejectedValue(refreshError) + : null; + try { + const final = await tracker.stop(completed); + expect(final).toMatchObject({ + usage: { + input_tokens: + expected.input_tokens + + (withWorker ? workerUsage.input_tokens : 0), + output_tokens: + expected.output_tokens + + (withWorker ? workerUsage.output_tokens : 0), + }, + cost: null, + }); + if (!withWorker && expected === completed) { + expect(final.usage).toBe(completed); + } + expect(await tracker.stop()).toBe(final); + expect(reportedErrors).toEqual(refreshFails ? [refreshError] : []); + expect(reportedCosts).toEqual([]); + } finally { + refresh?.mockRestore(); + } + }, + ); + } + const unpriceable = { + input_tokens: Number.MAX_SAFE_INTEGER, + output_tokens: 0, + }; + for (const [replacement, maxCostUsd, rejects] of [ + [observedRoot, undefined, false], + [unpriceable, undefined, false], + [unpriceable, 1, true], + ] as const) { + const costs: number[] = []; + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", [ + accountingEvent(completedRoot), + ]), + }, + { + model: "gpt-5.6-terra", + maxCostUsd, + onCost: (cost) => costs.push(cost.estimatedUsd), + }, + async (tracker) => { + const prior = await tracker.stop(); + expect(prior.cost?.estimatedUsd).toBe(0.0032); + const completed = tracker.stop(replacement); + if (rejects) { + await expect(completed).rejects.toThrow( + "The scan cost limit could not be verified", + ); + } else { + await expect(completed).resolves.toEqual(prior); + } + expect(costs).toEqual([0.0032]); + }, + ); + } + await withMockAccountingSessions( + { + "scan-thread": accountingSession("scan-thread", []), + "worker-thread": accountingSession( + "worker-thread", + [accountingEvent(workerUsage)], + "scan-thread", + ), + }, + { model: "unknown-model" }, + async (tracker, append) => { + const prior = await tracker.stop(); + expect(prior).toMatchObject({ usage: workerUsage, cost: null }); + append( + "conflicting-worker", + accountingSession("worker-thread", [], "another-thread"), + ); + await tracker.refresh(); + expect(await tracker.stop({})).toBe(prior); + }, + ); + }); + + test("retains a partial event across incremental reads", async () => { + const home = await codexHome(); + const path = await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const events: ScanSessionEvent[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + onSessionEvent: (event) => events.push(event), + }); + tracker.start("scan-thread"); + await tracker.refresh(); + + const event = JSON.stringify({ + type: "event_msg", + payload: { + type: "token_count", + info: { + total_token_usage: { input_tokens: 250, output_tokens: 20 }, + }, + }, + }); + const padding = " ".repeat(128 * 1_024); + await appendFile(path, `${padding}${event.slice(0, 40)}`); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + expect(events).toHaveLength(2); + + await appendFile(path, `${event.slice(40)}\n`); + expect((await tracker.stop()).cost?.inputTokens).toBe(250); + expect(events).toHaveLength(3); + expect(events.at(-1)?.event).toEqual(JSON.parse(event)); + }); + + test("reads session events larger than 16 MiB", async () => { + const home = await codexHome(); + const path = await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + }); + tracker.start("scan-thread"); + + const event = JSON.stringify({ + type: "event_msg", + payload: { + type: "token_count", + info: { + total_token_usage: { input_tokens: 250, output_tokens: 20 }, + }, + details: "x".repeat(16 * 1_024 * 1_024 + 1), + }, + }); + await appendFile(path, event.slice(0, -10)); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + + await appendFile(path, `${event.slice(-10)}\n`); + expect((await tracker.stop()).cost?.inputTokens).toBe(250); + }); + + testPosix( + "quarantines unreadable unrelated sessions after one failure", + async () => { + const home = await codexHome(); + const unreadable = await writeSession(home, "unrelated-thread", { + input_tokens: 99, + output_tokens: 1, + }); + await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + }); + tracker.start("scan-thread"); + await chmod(unreadable, 0o000); + + try { + await expect(tracker.refresh()).rejects.toThrow(); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + expect((await tracker.stop()).cost?.inputTokens).toBe(100); + } finally { + await chmod(unreadable, 0o600); + } + }, + ); + + testPosix( + "keeps budgeted worker sessions fail-closed after they become unreadable", + async () => { + const { worker, tracker } = await workerScan({ + workerUsage: { input_tokens: 250, output_tokens: 20 }, + maxCostUsd: 1, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(350); + await chmod(worker, 0o000); + + try { + await expect(tracker.refresh()).rejects.toThrow(); + await expect(tracker.refresh()).rejects.toThrow(); + await expect( + tracker.stop({ input_tokens: 100, output_tokens: 10 }), + ).rejects.toThrow(); + } finally { + await chmod(worker, 0o600); + } + }, + ); + + testPosix( + "combines verified workers with completed usage after a root read failure", + async () => { + const { root, tracker } = await workerScan({ + workerUsage: { input_tokens: 250, output_tokens: 20 }, + maxCostUsd: 1, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(350); + await chmod(root, 0o000); + + try { + await expect( + tracker.stop({ input_tokens: 1_000, output_tokens: 100 }), + ).resolves.toMatchObject({ + usage: { input_tokens: 1_250, output_tokens: 120 }, + cost: { inputTokens: 1_250, outputTokens: 120 }, + }); + } finally { + await chmod(root, 0o600); + } + }, + ); + + testPosix.each(["parented", "independent"] as const)( + "rejects budget fallback for a newly discovered %s worker during a failed refresh", + async (workerKind) => { + const home = await codexHome(); + const scanDirectory = join(home, "scan"); + const active = await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + scanDirectory, + maxCostUsd: 0.001, + }); + tracker.start("scan-thread"); + await tracker.refresh(); + await writeSession( + home, + "worker-thread", + { input_tokens: 10_000, output_tokens: 1_000 }, + workerKind === "parented" ? "scan-thread" : undefined, + workerKind === "independent" + ? join(scanDirectory, "artifacts") + : undefined, + "2026-07-26T12:01:00Z", + ); + await chmod(active, 0o000); + + try { + await expect( + tracker.stop({ input_tokens: 100, output_tokens: 10 }), + ).rejects.toThrow(); + } finally { + await chmod(active, 0o600); + } + }, + ); + + testPosix( + "rejects budget fallback when an unreadable session cannot be attributed", + async () => { + const home = await codexHome(); + const unreadable = await writeSession(home, "unidentified-thread", { + input_tokens: 100, + output_tokens: 10, + }); + await chmod(unreadable, 0o000); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 1, + }); + tracker.start("scan-thread"); + + try { + await expect( + tracker.stop({ input_tokens: 1_000, output_tokens: 100 }), + ).rejects.toThrow(); + } finally { + await chmod(unreadable, 0o600); + } + }, + ); + + test.each(["scan-thread", "worker-thread"] as const)( + "rejects a budgeted scan when its tracked %s session disappears", + async (missingThread) => { + const { root, worker, tracker } = await workerScan({ maxCostUsd: 1 }); + await tracker.refresh(); + await rm(missingThread === "scan-thread" ? root : worker); + + await expect( + tracker.stop({ input_tokens: 1_000, output_tokens: 100 }), + ).rejects.toThrow( + "A tracked scan session disappeared before its cost could be verified.", + ); + }, + ); + + test("ignores unrelated disappearing sessions when enforcing a budget", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const unrelated = await writeSession(home, "unrelated-thread", { + input_tokens: 1_000, + output_tokens: 100, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 1, + }); + tracker.start("scan-thread"); + await tracker.refresh(); + await rm(unrelated); + + expect((await tracker.stop()).cost?.inputTokens).toBe(100); + }); + + test("preserves known usage when an optional worker session disappears", async () => { + const { worker, tracker } = await workerScan(); + await tracker.refresh(); + await rm(worker); + + expect( + (await tracker.stop({ input_tokens: 1_000, output_tokens: 100 })).cost, + ).toMatchObject({ inputTokens: 1_100, outputTokens: 110 }); + }); + + test("reports a changed running cost only once", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { + input_tokens: 1_250, + cached_input_tokens: 200, + output_tokens: 30, + }); + const updates: number[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-sol", + maxCostUsd: 0.005, + onCost: (cost) => updates.push(cost.estimatedUsd), + }); + tracker.start("scan-thread"); + + await tracker.stop(); + + expect(updates).toEqual([0.00625]); }); test("falls back to the completed turn when session logs are unavailable", async () => { const tracker = new ScanCostTracker({ - codexHome: await codexHome(), - model: "gpt-5.6-luna", + codexHome: await codexHome(), + model: "gpt-5.6-luna", + }); + const usage = { input_tokens: 1_000, output_tokens: 20 }; + tracker.start("scan-thread"); + + expect(await tracker.stop(usage)).toEqual({ + usage, + cost: { + model: "gpt-5.6-luna", + inputTokens: 1_000, + cachedInputTokens: 0, + cacheWriteInputTokens: 0, + outputTokens: 20, + estimatedUsd: 0.000224, + }, + }); + }); + + test.each([ + ["missing", undefined], + ["null", null], + ["malformed", {}], + ] as const)( + "rejects a budgeted scan when completed-turn usage is %s", + async (_description, usage) => { + const tracker = new ScanCostTracker({ + codexHome: await codexHome(), + model: "gpt-5.6-terra", + maxCostUsd: 1, + }); + tracker.start("scan-thread"); + + await expect(tracker.stop(usage)).rejects.toThrow( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + }, + ); + + test.each([ + ["missing", undefined], + ["null", null], + ["malformed", {}], + ] as const)( + "requires completed root-session evidence when final usage is %s", + async (_description, usage) => { + const home = await codexHome(); + const root = await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 1, + }); + tracker.start("scan-thread"); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + + await expect(tracker.stop(usage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + await appendFile(root, `${taskEvent("task_complete")}\n`); + expect((await tracker.stop(usage)).cost).toMatchObject({ + inputTokens: 100, + outputTokens: 10, + }); + }, + ); + + test("rejects a budgeted scan when the completed model cannot be priced", async () => { + const tracker = new ScanCostTracker({ + codexHome: await codexHome(), + model: "unknown-model", + maxCostUsd: 1, + }); + tracker.start("scan-thread"); + + await expect( + tracker.stop({ input_tokens: 1_000, output_tokens: 100 }), + ).rejects.toThrow( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + }); + + test("allows unavailable completed-turn usage without an explicit budget", async () => { + const tracker = new ScanCostTracker({ + codexHome: await codexHome(), + model: "gpt-5.6-terra", + }); + tracker.start("scan-thread"); + + await expect(tracker.stop(null)).resolves.toEqual({ + usage: null, + cost: null, + }); + }); + + test("preserves unfinished root usage when tracking is optional", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + }); + tracker.start("scan-thread"); + + expect((await tracker.stop(undefined)).cost).toMatchObject({ + inputTokens: 100, + outputTokens: 10, + }); + }); + + test.each([ + ["missing", undefined], + ["malformed", {}], + ] as const)( + "rejects worker-only usage when the budgeted root completion is %s", + async (_description, completedRoot) => { + const { tracker } = await workerScan({ + rootUsage: null, + maxCostUsd: 1, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + + await expect(tracker.stop(completedRoot)).rejects.toThrow( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + }, + ); + + test.each([ + ["rejects", 1], + ["allows", undefined], + ] as const)( + "%s incomplete delegated-worker usage according to the explicit budget", + async (_result, maxCostUsd) => { + const { tracker } = await workerScan({ workerUsage: null, maxCostUsd }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + + const completed = tracker.stop({ + input_tokens: 1_000, + output_tokens: 100, + }); + if (maxCostUsd === undefined) { + await expect(completed).resolves.toMatchObject({ + cost: { inputTokens: 1_000, outputTokens: 100 }, + }); + } else { + await expect(completed).rejects.toThrow( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + } + }, + ); + + test.each([ + ["still-running worker", "running", 1, true], + ["completed worker", "task_complete", 1, false], + ["compatible completed worker", "turn_complete", 1, false], + ["canceled worker", "turn_aborted", 1, false], + ["worker restarted after completion", "task_started", 1, true], + ["optional still-running worker", "running", undefined, false], + ] as const)( + "verifies final delegated-worker completion for a %s", + async (_scenario, state, maxCostUsd, shouldReject) => { + const { worker, tracker } = await workerScan({ + workerCompleted: false, + maxCostUsd, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + if (state !== "running") { + if (state === "task_started") { + await appendFile(worker, `${taskEvent("task_complete")}\n`); + } + await appendFile(worker, `${taskEvent(state)}\n`); + } + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + + const completed = tracker.stop({ input_tokens: 100, output_tokens: 10 }); + if (shouldReject) { + await expect(completed).rejects.toThrow( + "The scan cost limit could not be verified", + ); + } else { + await expect(completed).resolves.toMatchObject({ + cost: { inputTokens: 200, outputTokens: 20 }, + }); + } + }, + ); + + test("preserves observed active-worker costs during budget failure cleanup", async () => { + const { tracker } = await workerScan({ + workerCompleted: false, + maxCostUsd: 1, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 200, + outputTokens: 20, + }); + }); + + test.each([ + ["null", null], + ["undefined", undefined], + ] as const)( + "rejects an active worker when completed-turn usage is %s", + async (_description, usage) => { + const { tracker } = await workerScan({ + workerCompleted: false, + maxCostUsd: 1, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + + await expect(tracker.stop(usage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + }, + ); + + test.each([ + ["budgeted root", "root", 0.001, true], + ["budgeted delegated worker", "worker", 0.001, true], + ["budgeted unrelated session", "unrelated", 0.001, false], + ["unbudgeted delegated worker", "worker", undefined, false], + ] as const)( + "handles an incomplete final event from a %s", + async (_description, session, maxCostUsd, shouldReject) => { + const { home, root, worker, tracker } = await workerScan({ maxCostUsd }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + const path = + session === "root" + ? root + : session === "worker" + ? worker + : await writeSession(home, "unrelated-thread", { + input_tokens: 100, + output_tokens: 10, + }); + await appendIncompleteTokenUsage(path); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + + const completed = tracker.stop({ input_tokens: 100, output_tokens: 10 }); + if (shouldReject) { + await expect(completed).rejects.toThrow( + "The scan cost limit could not be verified because model pricing or token usage is unavailable.", + ); + } else { + await expect(completed).resolves.toMatchObject({ + cost: { inputTokens: 200, outputTokens: 20 }, + }); + } + }, + ); + + test.each([ + ["budgeted root without completed usage", "root", 0.001, false, true], + ["budgeted root with completed usage", "root", 0.001, true, false], + ["budgeted delegated worker", "worker", 0.001, true, true], + ["budgeted unrelated session", "unrelated", 0.001, true, false], + ["unbudgeted delegated worker", "worker", undefined, true, false], + ] as const)( + "handles a malformed complete session record from a %s", + async (_description, session, maxCostUsd, authoritative, shouldReject) => { + const { home, root, worker, tracker } = await workerScan({ maxCostUsd }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + const path = + session === "root" + ? root + : session === "worker" + ? worker + : await writeSession(home, "unrelated-thread", { + input_tokens: 100, + output_tokens: 10, + }); + await appendFile( + path, + '{"type":"event_msg","payload":{"type":"token_count"\n' + + `${taskEvent("task_complete")}\n`, + ); + + const refreshed = tracker.refresh(); + if (maxCostUsd !== undefined && session !== "unrelated") { + await expect(refreshed).rejects.toThrow( + "tracked session record could not be read", + ); + } else { + await expect(refreshed).resolves.toMatchObject({ + cost: { inputTokens: 200 }, + }); + } + + const completed = authoritative + ? tracker.stop({ input_tokens: 100, output_tokens: 10 }) + : tracker.stop(); + if (shouldReject) { + await expect(completed).rejects.toThrow( + "tracked session record could not be read", + ); + } else { + await expect(completed).resolves.toMatchObject({ + cost: { inputTokens: 200, outputTokens: 20 }, + }); + } + }, + ); + + test("rejects budget finalization while discovered worker metadata is incomplete", async () => { + const home = await codexHome(); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + await writeSession(home, "scan-thread", rootUsage); + const worker = join( + home, + "sessions", + "2026", + "07", + "26", + "rollout-pending-worker.jsonl", + ); + const metadata = JSON.stringify({ + type: "session_meta", + payload: { id: "worker-thread", parent_thread_id: "scan-thread" }, + }); + await writeFile(worker, metadata.slice(0, -1)); + const reportedCosts: number[] = []; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + onCost: (cost) => reportedCosts.push(cost.estimatedUsd), + }); + tracker.start("scan-thread"); + + await expect(tracker.stop(rootUsage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + expect(reportedCosts).toEqual([0.00032]); + + await appendFile( + worker, + `${metadata.slice(-1)}\n${JSON.stringify({ + type: "event_msg", + payload: { + type: "token_count", + info: { + total_token_usage: { + input_tokens: 10_000, + output_tokens: 1_000, + }, + }, + }, + })}\n${taskEvent("task_complete")}\n`, + ); + + await expect(tracker.stop(rootUsage)).resolves.toMatchObject({ + cost: { + inputTokens: 10_100, + outputTokens: 1_010, + estimatedUsd: 0.03232, + }, + }); + expect(reportedCosts).toEqual([0.00032, 0.03232]); + }); + + test("ignores unrelated incomplete rollouts confirmed by trusted session ownership", async () => { + const home = await codexHome(); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const root = await writeSession(home, "scan-thread", rootUsage); + await writeFile( + join( + home, + "sessions", + "2026", + "07", + "26", + "rollout-unrelated-incomplete.jsonl", + ), + '{"type":"session_meta","payload":{"id":"unrelated-thread"', + ); + let ownershipChecks = 0; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + resolveOwnedSessionPaths: async () => { + ownershipChecks += 1; + return new Set([root]); + }, }); - const usage = { input_tokens: 1_000, output_tokens: 20 }; tracker.start("scan-thread"); - expect(await tracker.stop(usage)).toEqual({ - usage, + await tracker.refresh(); + await tracker.refresh(); + expect(ownershipChecks).toBe(0); + await expect(tracker.stop(rootUsage)).resolves.toMatchObject({ cost: { - model: "gpt-5.6-luna", - inputTokens: 1_000, - cachedInputTokens: 0, - cacheWriteInputTokens: 0, - outputTokens: 20, - estimatedUsd: 0.000224, + inputTokens: 100, + outputTokens: 10, + estimatedUsd: 0.00032, + }, + }); + expect(ownershipChecks).toBe(1); + }); + + test.each([ + ["delegated worker", "delegated", true, "partial"], + ["independent Deep Scan worker", "independent", true, "partial"], + ["unrelated session", "unrelated", false, "partial"], + ["empty delegated worker", "delegated", true, "empty"], + ["empty independent Deep Scan worker", "independent", true, "empty"], + ["empty unrelated session", "unrelated", false, "empty"], + ["malformed delegated worker", "delegated", true, "malformed"], + [ + "malformed independent Deep Scan worker", + "independent", + true, + "malformed", + ], + ["malformed unrelated session", "unrelated", false, "malformed"], + ["unattributed delegated worker usage", "delegated", true, "usage"], + [ + "unattributed independent Deep Scan worker usage", + "independent", + true, + "usage", + ], + ["unattributed unrelated usage", "unrelated", false, "usage"], + ["legacy delegated worker metadata", "delegated", true, "session-id"], + [ + "legacy independent Deep Scan worker metadata", + "independent", + true, + "session-id", + ], + ["legacy unrelated session metadata", "unrelated", false, "session-id"], + [ + "deleted delegated worker rollout", + "missing-file-delegated", + true, + "missing", + ], + [ + "deleted independent Deep Scan worker rollout", + "missing-file-independent", + true, + "missing", + ], + [ + "unobserved delegated worker rollout", + "unobserved-delegated", + true, + "unobserved", + ], + [ + "unobserved independent Deep Scan worker rollout", + "unobserved-independent", + true, + "unobserved", + ], + ["complete root-only scan", "root-only", false, "missing"], + ["incomplete ownership graph", "missing-worker", true, "partial"], + ["missing scan root", "missing-root", true, "partial"], + ["missing Codex state database", "missing-database", true, "partial"], + ["missing workbench scan", "missing-scan", true, "partial"], + ] as const)( + "checks incomplete metadata against trusted SQLite ownership for a %s", + async (_description, relationship, shouldReject, contents) => { + const home = await codexHome(); + const stateDirectory = join(home, "workbench"); + await mkdir(stateDirectory); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const root = await writeSession(home, "scan-thread", rootUsage); + const incomplete = relationship.startsWith("unobserved-") + ? join(home, "unobserved-owned-worker.jsonl") + : join( + home, + "sessions", + "2026", + "07", + "26", + "rollout-unrelated-incomplete.jsonl", + ); + const workerUsage = JSON.stringify({ + type: "event_msg", + payload: { + type: "token_count", + info: { + total_token_usage: { input_tokens: 10_000, output_tokens: 1_000 }, + }, + }, + }); + const unidentifiedContents = + contents === "empty" + ? "" + : contents === "malformed" + ? '{"type":"session_meta","payload":\n' + : contents === "usage" + ? `${workerUsage}\n` + : contents === "session-id" || contents === "unobserved" + ? `${JSON.stringify({ + type: "session_meta", + payload: { + ...(contents === "session-id" + ? { session_id: "pending-thread" } + : { id: "pending-thread" }), + ...(relationship === "delegated" || + relationship === "unobserved-delegated" + ? { parent_thread_id: "scan-thread" } + : {}), + }, + })}\n${workerUsage}\n${taskEvent("task_complete")}\n` + : '{"type":"session_meta","payload":{"id":"unrelated-thread"'; + await writeFile(incomplete, unidentifiedContents); + const { resolvePluginPython, resolveScanSessionPaths } = await import( + "../src/runtime.js" + ); + const python = await resolvePluginPython({ environment: process.env }); + const fixture = spawnSync( + python, + [ + "-I", + "-B", + "-c", + [ + "import json, sqlite3, sys", + "from pathlib import Path", + "home, root, incomplete, relationship = json.loads(sys.argv[1])", + "workbench = sqlite3.connect(Path(home) / 'workbench' / 'workbench.sqlite3')", + "workbench.execute('PRAGMA journal_mode = WAL')", + "workbench.execute('CREATE TABLE scans (id TEXT PRIMARY KEY, workspace_id TEXT NOT NULL, mode TEXT NOT NULL)')", + "workbench.execute('CREATE TABLE workspaces (id TEXT PRIMARY KEY, thread_id TEXT)')", + "workbench.execute('CREATE TABLE deep_scan_workers (scan_id TEXT NOT NULL, sdk_thread_id TEXT)')", + "workbench.execute('INSERT INTO scans VALUES (?, ?, ?)', ('fixture-scan', 'fixture-workspace', 'deep'))", + "workbench.execute('INSERT INTO workspaces VALUES (?, ?)', ('fixture-workspace', 'scan-thread'))", + "if relationship in ('independent', 'missing-file-independent', 'unobserved-independent'):", + " workbench.execute('INSERT INTO deep_scan_workers VALUES (?, ?)', ('fixture-scan', 'pending-thread'))", + "workbench.commit()", + "workbench.close()", + "state = sqlite3.connect(Path(home) / 'state_7.sqlite')", + "state.execute('PRAGMA journal_mode = WAL')", + "state.execute('CREATE TABLE threads (id TEXT PRIMARY KEY, rollout_path TEXT NOT NULL)')", + "state.execute('CREATE TABLE thread_spawn_edges (parent_thread_id TEXT NOT NULL, child_thread_id TEXT NOT NULL)')", + "if relationship != 'missing-root':", + " state.execute('INSERT INTO threads VALUES (?, ?)', ('scan-thread', root))", + "if relationship not in ('missing-worker', 'root-only'):", + " state.execute('INSERT INTO threads VALUES (?, ?)', ('pending-thread', incomplete))", + "if relationship in ('delegated', 'missing-worker', 'missing-file-delegated', 'unobserved-delegated'):", + " state.execute('INSERT INTO thread_spawn_edges VALUES (?, ?)', ('scan-thread', 'pending-thread'))", + "state.commit()", + "state.close()", + ].join("\n"), + JSON.stringify([home, root, incomplete, relationship]), + ], + { encoding: "utf8" }, + ); + expect(fixture.status, fixture.stderr).toBe(0); + if (contents === "missing") { + await rm(incomplete); + } + if (relationship === "missing-database") { + await rm(join(home, "state_7.sqlite")); + } + const { PLUGIN_ROOT } = await import("./plugin-root.js"); + let pluginRoot = PLUGIN_ROOT; + if (relationship === "unrelated") { + pluginRoot = join(home, "ownership-plugin"); + await mkdir(join(pluginRoot, "scripts"), { recursive: true }); + const helper = join(PLUGIN_ROOT, "scripts", "workbench_scan_usage.py"); + await writeFile( + join(pluginRoot, "scripts", "workbench_scan_usage.py"), + [ + "import os, sys", + "assert sys.flags.isolated and sys.dont_write_bytecode", + "assert not any(os.environ.get(key) for key in ('OPENAI_API_KEY', 'CODEX_API_KEY', 'OPENROUTER_API_KEY', 'FIREWORKS_API_KEY'))", + `exec(compile(open(${JSON.stringify(helper)}).read(), ${JSON.stringify(helper)}, 'exec'), globals())`, + ].join("\n"), + ); + } + let ownershipChecks = 0; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + resolveOwnedSessionPaths: async (threadId) => { + ownershipChecks += 1; + return await resolveScanSessionPaths( + { + python, + pluginRoot, + environment: { + PATH: process.env["PATH"], + SystemRoot: + process.env["SystemRoot"] ?? process.env["SYSTEMROOT"], + CODEX_HOME: home, + CODEX_SECURITY_STATE_DIR: stateDirectory, + OPENAI_API_KEY: "synthetic-never-forwarded", + CODEX_API_KEY: "synthetic-never-forwarded", + OPENROUTER_API_KEY: "synthetic-never-forwarded", + FIREWORKS_API_KEY: "synthetic-never-forwarded", + }, + }, + relationship === "missing-scan" ? "missing-scan" : "fixture-scan", + threadId, + ); + }, + }); + tracker.start("scan-thread"); + await tracker.refresh(); + expect(ownershipChecks).toBe(0); + + const result = tracker.stop(rootUsage); + if (shouldReject) { + await expect(result).rejects.toThrow(/could not be verified/u); + } else { + await expect(result).resolves.toMatchObject({ + cost: { inputTokens: 100, outputTokens: 10, estimatedUsd: 0.00032 }, + }); + } + expect(ownershipChecks).toBe(1); + }, + ); + + test("does not let an unrelated partial session hide incomplete owned usage", async () => { + const home = await codexHome(); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const root = await writeSession(home, "scan-thread", rootUsage); + const worker = await writeSession( + home, + "worker-thread", + { input_tokens: 10, output_tokens: 1 }, + "scan-thread", + undefined, + undefined, + true, + ); + await appendIncompleteTokenUsage(worker); + await writeFile( + join( + home, + "sessions", + "2026", + "07", + "26", + "rollout-unrelated-incomplete.jsonl", + ), + '{"type":"session_meta","payload":{"id":"unrelated-thread"', + ); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + resolveOwnedSessionPaths: async () => new Set([root, worker]), + }); + tracker.start("scan-thread"); + + await expect(tracker.stop(rootUsage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + }); + + test.each([ + ["empty", ""], + ["malformed", '{"type":"session_meta","payload":\n'], + [ + "unattributed billable", + `${JSON.stringify({ + type: "event_msg", + payload: { + type: "token_count", + info: { + total_token_usage: { input_tokens: 10_000, output_tokens: 1_000 }, + }, + }, + })}\n`, + ], + [ + "legacy session metadata", + `${JSON.stringify({ + type: "session_meta", + payload: { session_id: "worker-thread" }, + })}\n`, + ], + ] as const)( + "rejects an unidentified %s rollout without trusted ownership", + async (_description, contents) => { + const home = await codexHome(); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + await writeSession(home, "scan-thread", rootUsage); + await writeFile( + join( + home, + "sessions", + "2026", + "07", + "26", + "rollout-unidentified.jsonl", + ), + contents, + ); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + }); + tracker.start("scan-thread"); + + await expect(tracker.stop(rootUsage)).rejects.toThrow( + "The scan cost limit could not be verified", + ); + }, + ); + + test.each([ + ["optional incomplete metadata", undefined, "partial"], + ["optional empty metadata", undefined, "empty"], + ["optional malformed metadata", undefined, "malformed"], + ["empty discovered rollout", 0.001, "empty"], + ["completed unrelated rollout", 0.001, "unrelated"], + ["removed incomplete rollout", 0.001, "removed"], + ] as const)( + "preserves accounting compatibility for %s", + async (_description, maxCostUsd, scenario) => { + const home = await codexHome(); + const rootUsage = { input_tokens: 100, output_tokens: 10 }; + const root = await writeSession(home, "scan-thread", rootUsage); + const path = join( + home, + "sessions", + "2026", + "07", + "26", + "rollout-unidentified.jsonl", + ); + if (scenario === "unrelated") { + await writeSession(home, "unrelated-thread", { + input_tokens: 10_000, + output_tokens: 1_000, + }); + } else { + await writeFile( + path, + scenario === "empty" + ? "" + : scenario === "malformed" + ? '{"type":"session_meta","payload":\n' + : '{"type":"session_meta","payload":{"id":"pending-worker"', + ); + } + let ownershipChecks = 0; + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd, + resolveOwnedSessionPaths: async () => { + ownershipChecks += 1; + return new Set([root]); + }, + }); + tracker.start("scan-thread"); + + if (scenario === "removed") { + await tracker.refresh(); + await rm(path); + } + + await expect(tracker.stop(rootUsage)).resolves.toMatchObject({ + cost: { + inputTokens: 100, + outputTokens: 10, + estimatedUsd: 0.00032, + }, + }); + expect(ownershipChecks).toBe(maxCostUsd === undefined ? 0 : 1); + }, + ); + + test("rejects an incomplete final root event without delegated workers", async () => { + const home = await codexHome(); + const root = await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: 0.001, + }); + tracker.start("scan-thread"); + await appendIncompleteTokenUsage(root); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + + await expect( + tracker.stop({ input_tokens: 100, output_tokens: 10 }), + ).rejects.toThrow("The scan cost limit could not be verified"); + }); + + test("keeps final budget enforcement stable when cleanup stops tracking twice", async () => { + const budget = 0.001; + const exceeded = new AbortController(); + const reportedCosts: number[] = []; + const { tracker } = await workerScan({ + maxCostUsd: budget, + onCost: (cost) => { + reportedCosts.push(cost.estimatedUsd); + if (cost.estimatedUsd > budget) exceeded.abort(); + }, + }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + + expect( + (await tracker.stop({ input_tokens: 1_000, output_tokens: 100 })).cost, + ).toMatchObject({ inputTokens: 1_100, outputTokens: 110 }); + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 1_100, + outputTokens: 110, + }); + expect(reportedCosts).toEqual([0.00064, 0.00352]); + expect(exceeded.signal.aborted).toBe(true); + }); + + test.each(["incomplete", "unfinished", "unreadable"] as const)( + "preserves a definitive overage when final worker evidence is %s", + async (evidence) => { + const reportedCosts: number[] = []; + const { worker, tracker } = await workerScan({ + workerCompleted: evidence !== "unfinished", + maxCostUsd: 0.001, + onCost: (cost) => reportedCosts.push(cost.estimatedUsd), + }); + const refresh = tracker.refresh.bind(tracker); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + if (evidence === "incomplete") { + await appendIncompleteTokenUsage(worker); + } else if (evidence === "unreadable") { + tracker.refresh = async () => { + throw new Error("session read failed"); + }; + } + + await expect( + tracker.stop({ input_tokens: 1_000, output_tokens: 100 }), + ).rejects.toThrow( + evidence === "unreadable" + ? "session read failed" + : "The scan cost limit could not be verified", + ); + expect(reportedCosts).toEqual([0.00064, 0.00352]); + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 1_100, + outputTokens: 110, + estimatedUsd: 0.00352, + }); + tracker.refresh = refresh; + if (evidence !== "incomplete") { + await appendIncompleteTokenUsage(worker); + } + await appendFile(worker, `}\n${taskEvent("task_complete")}\n`); + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 11_000, + outputTokens: 1_100, + estimatedUsd: 0.0352, + }); + expect(reportedCosts.at(-1)).toBe(0.0352); + }, + ); + + test("returns a definitive overage first discovered during failure cleanup", async () => { + const { worker, tracker } = await workerScan({ + rootUsage: { input_tokens: 1_000, output_tokens: 100 }, + maxCostUsd: 0.001, + }); + await appendIncompleteTokenUsage(worker); + + expect((await tracker.stop()).cost).toMatchObject({ + inputTokens: 1_100, + outputTokens: 110, + estimatedUsd: 0.00352, + }); + }); + + test("preserves higher observed root usage when completed-turn usage is stale", async () => { + const { tracker } = await workerScan({ + rootUsage: { input_tokens: 1_000, output_tokens: 100 }, + maxCostUsd: 1, + }); + + expect( + (await tracker.stop({ input_tokens: 100, output_tokens: 10 })).cost, + ).toMatchObject({ inputTokens: 1_100, outputTokens: 110 }); + }); + + test("combines worker-only observations with a valid completed root", async () => { + const budget = 0.001; + const exceeded = new AbortController(); + const { tracker } = await workerScan({ + rootUsage: null, + maxCostUsd: budget, + onCost: (cost) => { + if (cost.estimatedUsd > budget) exceeded.abort(); }, }); + expect((await tracker.refresh()).cost?.inputTokens).toBe(100); + + expect( + (await tracker.stop({ input_tokens: 1_000, output_tokens: 100 })).cost, + ).toMatchObject({ inputTokens: 1_100, outputTokens: 110 }); + expect(exceeded.signal.aborted).toBe(true); + }); + + test("uses completed-turn usage when the final session refresh fails", async () => { + const home = await codexHome(); + await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const reportedCosts: number[] = []; + const reportedErrors: unknown[] = []; + const refreshError = new Error("session read failed"); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + onCost: (cost) => reportedCosts.push(cost.estimatedUsd), + onError: (error) => reportedErrors.push(error), + }); + tracker.start("scan-thread"); + expect((await tracker.refresh()).cost?.estimatedUsd).toBe(0.00032); + tracker.refresh = async () => { + throw refreshError; + }; + const usage = { input_tokens: 1_000, output_tokens: 100 }; + + expect(await tracker.stop(usage)).toMatchObject({ + usage, + cost: { inputTokens: 1_000, estimatedUsd: 0.0032 }, + }); + expect(reportedCosts).toEqual([0.00032, 0.0032]); + expect(reportedErrors).toEqual([refreshError]); + }); + + test("adds observed worker usage to the completed root after a failed refresh", async () => { + const { tracker } = await workerScan(); + expect((await tracker.refresh()).cost?.inputTokens).toBe(200); + tracker.refresh = async () => { + throw new Error("session read failed"); + }; + + expect( + (await tracker.stop({ input_tokens: 1_000, output_tokens: 100 })).cost, + ).toMatchObject({ inputTokens: 1_100, outputTokens: 110 }); }); + + testPosix( + "enforces a budget with completed-turn usage when only its root session is unreadable", + async () => { + const home = await codexHome(); + const active = await writeSession(home, "scan-thread", { + input_tokens: 100, + output_tokens: 10, + }); + const budget = 0.001; + const exceeded = new AbortController(); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-terra", + maxCostUsd: budget, + onCost: (cost) => { + if (cost.estimatedUsd > budget) exceeded.abort(); + }, + }); + tracker.start("scan-thread"); + await tracker.refresh(); + await chmod(active, 0o000); + + try { + await expect(tracker.refresh()).rejects.toThrow(); + expect( + (await tracker.stop({ input_tokens: 1_000, output_tokens: 100 })) + .cost, + ).toMatchObject({ estimatedUsd: 0.0032 }); + expect(exceeded.signal.aborted).toBe(true); + } finally { + await chmod(active, 0o600); + } + }, + ); + + testPosix( + "uses completed-turn usage when an unrelated session cannot be opened", + async () => { + const home = await codexHome(); + const unreadable = await writeSession(home, "unrelated-thread", { + input_tokens: 99, + output_tokens: 1, + }); + await chmod(unreadable, 0o000); + const tracker = new ScanCostTracker({ + codexHome: home, + model: "gpt-5.6-luna", + }); + tracker.start("scan-thread"); + const usage = { input_tokens: 1_000, output_tokens: 20 }; + + try { + expect(await tracker.stop(usage)).toMatchObject({ + usage, + cost: { inputTokens: 1_000, estimatedUsd: 0.000224 }, + }); + } finally { + await chmod(unreadable, 0o600); + } + }, + ); }); diff --git a/sdk/typescript/tests-ts/custom-validation.test.ts b/sdk/typescript/tests-ts/custom-validation.test.ts index 6da81a99e..9253c9bdf 100644 --- a/sdk/typescript/tests-ts/custom-validation.test.ts +++ b/sdk/typescript/tests-ts/custom-validation.test.ts @@ -141,7 +141,9 @@ async function* responseEvents( activity?: string, ): AsyncGenerator { for await (const event of completedEvents()) { - if ( + if (event.type === "thread.started") { + yield { ...event, thread_id: "validation-thread" }; + } else if ( event.type === "item.completed" && event.item.type === "agent_message" ) { @@ -402,7 +404,10 @@ describe("custom validation", () => { expect(threadOptions.threadSource).toBe("security_scan"); workingDirectories.push(threadOptions.workingDirectory); return { - id: "thread-1", + id: + threadOptions.workingDirectory === scanDir + ? "thread-1" + : "validation-thread", async runStreamed(prompt, turnOptions) { turns += 1; if (turns === 1) { diff --git a/sdk/typescript/tests-ts/deep-scan-worker-shutdown.test.ts b/sdk/typescript/tests-ts/deep-scan-worker-shutdown.test.ts index 594ce173f..98919a8a1 100644 --- a/sdk/typescript/tests-ts/deep-scan-worker-shutdown.test.ts +++ b/sdk/typescript/tests-ts/deep-scan-worker-shutdown.test.ts @@ -112,10 +112,13 @@ test("does not start a bundled worker when its permission profile check fails", expect(started).toBe(false); }); -test("settles completed bundled Deep Scan workers during coordinator cancellation", async () => { +test("drains completed bundled Deep Scan workers during coordinator cancellation", async () => { const parentController = new AbortController(); + const draining = Promise.withResolvers(); + const releaseDrain = Promise.withResolvers(); let workerSignal: AbortSignal | undefined; let iteratorClosed = false; + let settled = false; const WorkerExecutor = await bundledWorkerExecutor(async function* ( signal: AbortSignal, ) { @@ -127,34 +130,56 @@ test("settles completed bundled Deep Scan workers during coordinator cancellatio item: { type: "agent_message", text: "worker completed" }, }; yield { type: "turn.completed" }; - await new Promise(() => {}); + draining.resolve(); + await releaseDrain.promise; } finally { iteratorClosed = true; - parentController.abort( - "coordinator canceled its remaining workers during cleanup", - ); } }); - const timeout = setTimeout(() => { - parentController.abort("completed bundled worker remained pending"); - }, 1_000); + const outcome = runWorker(WorkerExecutor, parentController.signal).finally( + () => { + settled = true; + }, + ); try { - const result = await runWorker(WorkerExecutor, parentController.signal); - - expect(result).toEqual({ + await Promise.race([ + draining.promise, + outcome.then(() => { + throw new Error("Worker settled before its SDK stream drained."); + }), + ]); + expect(settled).toBe(false); + expect(iteratorClosed).toBe(false); + parentController.abort( + "coordinator canceled its remaining workers during cleanup", + ); + expect(workerSignal).not.toBe(parentController.signal); + expect(workerSignal?.aborted).toBe(false); + releaseDrain.resolve(); + expect(await outcome).toEqual({ finalResponse: "worker completed", threadId: "fixture-worker-thread", }); expect(iteratorClosed).toBe(true); expect(parentController.signal.aborted).toBe(true); - expect(workerSignal).not.toBe(parentController.signal); - expect(workerSignal?.aborted).toBe(false); } finally { - clearTimeout(timeout); + releaseDrain.resolve(); + await outcome.catch(() => {}); } }); +test("propagates bundled Deep Scan worker shutdown failures", async () => { + const WorkerExecutor = await bundledWorkerExecutor(async function* () { + yield { type: "turn.completed" }; + throw new Error("Synthetic worker shutdown failure"); + }); + + await expect( + runWorker(WorkerExecutor, new AbortController().signal), + ).rejects.toThrow("Synthetic worker shutdown failure"); +}); + test("forwards coordinator cancellation to active bundled Deep Scan workers", async () => { const parentController = new AbortController(); const cancellation = new Error("coordinator canceled an active worker"); diff --git a/sdk/typescript/tests-ts/runtime.test.ts b/sdk/typescript/tests-ts/runtime.test.ts index 66fc9bff7..f05334560 100644 --- a/sdk/typescript/tests-ts/runtime.test.ts +++ b/sdk/typescript/tests-ts/runtime.test.ts @@ -1952,12 +1952,12 @@ describe("plugin runtime preparation", () => { ]); }); - test("upgrades a cached 0.1.37 plugin with the real bundled Codex executable", async () => { + test("upgrades the prior bundled plugin through the real Codex cache", async () => { const root = await temporaryDirectory(); - const previous = await plugin(join(root, "previous"), "0.1.37"); + const previous = await plugin(join(root, "previous"), "0.1.60"); await writeFile( - join(previous, ".mcp.json"), - JSON.stringify({ mcpServers: { "codex-security": { env_vars: [] } } }), + join(previous, "scripts", "workbench_scan_usage.py"), + "LEGACY_SCAN_USAGE_READER = True\n", ); const home = join(root, "home"); await mkdir(home, { mode: 0o700 }); @@ -1983,9 +1983,89 @@ describe("plugin runtime preparation", () => { const credentials = await readFile(join(home, "auth.json"), "utf8"); const options = { codexCommand: command, environment }; - const first = await bootstrapPlugin(home, previous, options); - expect(first.version).toBe("0.1.37"); + const installedPrevious = await bootstrapPlugin(home, previous, options); + expect(installedPrevious.version).toBe("0.1.60"); + expect( + await readFile( + join( + installedPrevious.installedRoot, + "scripts", + "workbench_scan_usage.py", + ), + "utf8", + ), + ).toContain("LEGACY_SCAN_USAGE_READER"); const upgraded = await bootstrapPlugin(home, PLUGIN_ROOT, options); + + expect(upgraded.version).toBe(BUNDLED_PLUGIN_VERSION); + expect(upgraded.installedRoot).toEndWith( + join("codex-security", BUNDLED_PLUGIN_VERSION), + ); + const rollout = join(root, "rollout.jsonl"); + const sample = (input: number, cached: number, output: number) => + JSON.stringify({ + type: "event_msg", + timestamp: "2026-07-26T12:02:00Z", + payload: { + type: "token_count", + info: { + total_token_usage: { + input_tokens: input, + cached_input_tokens: cached, + output_tokens: output, + total_tokens: input + output, + }, + }, + }, + }); + await writeFile( + rollout, + [ + JSON.stringify({ + type: "session_meta", + payload: { id: "scan-thread" }, + }), + sample(100, 100, 100), + sample(101, 99, 99), + sample(101, 99, 99), + "", + ].join("\n"), + ); + const python = await resolvePluginPython({ environment }); + const persisted = spawnSync( + python, + [ + "-I", + "-B", + "-c", + [ + "import json, sys", + "from datetime import datetime, timezone", + "from pathlib import Path", + "sys.path.insert(0, sys.argv[1])", + "import workbench_scan_usage as usage", + "from workbench_validation import parse_scan_cost", + "session = usage.RolloutSession('scan-thread', None, Path(sys.argv[2]))", + "measured, warnings = usage._read_rollout_usage(session, started_at=datetime(2026, 7, 26, tzinfo=timezone.utc), completed_at=None)", + "parse_scan_cost(usage.measured_scan_cost_json({'coverage': 'complete', 'source': 'codex_rollout', 'threadCount': 1, **measured}))", + "print(json.dumps({'usage': measured, 'warnings': sorted(warnings)}))", + ].join("\n"), + join(upgraded.installedRoot, "scripts"), + rollout, + ], + { encoding: "utf8", env: environment, windowsHide: true }, + ); + expect(persisted.status, persisted.stderr).toBe(0); + expect(JSON.parse(persisted.stdout)).toMatchObject({ + usage: { + inputTokens: 101, + cachedInputTokens: 101, + outputTokens: 199, + totalTokens: 300, + }, + warnings: [], + }); + const configuration = JSON.parse( await readFile(join(upgraded.installedRoot, ".mcp.json"), "utf8"), ) as { @@ -1994,8 +2074,8 @@ describe("plugin runtime preparation", () => { const server = configuration.mcpServers["codex-security"]; expect(upgraded.version).toBe(BUNDLED_PLUGIN_VERSION); - expect(upgraded.version).not.toBe(first.version); - expect(upgraded.installedRoot).not.toBe(first.installedRoot); + expect(upgraded.version).not.toBe(installedPrevious.version); + expect(upgraded.installedRoot).not.toBe(installedPrevious.installedRoot); expect(server?.command).toBe("./scripts/launch_codex_security_mcp"); expect(server?.env_vars).toContain("CODEX_MANAGED_PACKAGE_ROOT"); expect(server?.env_vars).toContain("CODEX_MCP_NODE_PATH"); diff --git a/sdk/typescript/tests-ts/support/api-client.ts b/sdk/typescript/tests-ts/support/api-client.ts index 851a440cd..ef4d5c2de 100644 --- a/sdk/typescript/tests-ts/support/api-client.ts +++ b/sdk/typescript/tests-ts/support/api-client.ts @@ -68,6 +68,7 @@ export class TestClient extends CodexSecurity { throw new Error("Unexpected Codex invocation in test"); }, environment: {}, + resolveScanSessionPaths: async () => new Set(), runWorkbench: async (_options, args, input) => mockWorkbench(args, input), ...dependencies,