fix(cost): enforce verifiable scan spending limits

plugins/codex-security/.codex-plugin/plugin.json

@@ -1,6 +1,6 @@
 {
   "name": "codex-security",
-  "version": "0.1.94",
+  "version": "0.1.99",
   "description": "Codex Security workflows for security scans, analysis, and investigation.",
   "author": {
     "name": "OpenAI"

plugins/codex-security/mcp-app/src/deep-scan/executor.ts

@@ -136,7 +136,7 @@ export class CodexSdkWorkerExecutor implements CodexWorkerExecutor {
           } else if (event.type === "turn.completed") {
             turnCompleted = true;
             request.signal.removeEventListener("abort", forwardAbort);
-            break;
+            // Drain the SDK stream so the worker can flush its session records.
           } else if (event.type === "turn.failed") {
             throw new Error(event.error.message);
           } else if (event.type === "error") {

plugins/codex-security/mcp-app/tests/test_deep_scan_executor.mjs

@@ -87,7 +87,7 @@ try {
     await testSandboxNamespaceDiagnosticIsSanitized();
     await testOwnedArtifactToolFailureDiagnosticIsSanitized();
     await testStreamTerminationWithoutTerminalEventFails();
-    await testCompletedWorkerSettlesWithoutWaitingForProcessExit();
+    await testCompletedWorkerFlushesBeforeSettling();
     await testAbortPropagation();
     await testConfigurationFailureIsNonRetryable();
     await testThreadStartConfigurationFailureIsNonRetryable();
@@ -1005,7 +1005,7 @@ async function testAbortPropagation() {
   }
 }
 
-async function testCompletedWorkerSettlesWithoutWaitingForProcessExit() {
+async function testCompletedWorkerFlushesBeforeSettling() {
   const fixture = await fakeCodexFixture();
   const previousPath = process.env.CODEX_CLI_PATH;
   process.env.CODEX_CLI_PATH = fixture.executablePath;
@@ -1021,7 +1021,7 @@ async function testCompletedWorkerSettlesWithoutWaitingForProcessExit() {
     const promptPath = path.join(fixture.root, "prompt.md");
     const workingDirectory = path.join(fixture.root, "artifacts");
     await mkdir(workingDirectory);
-    await writeFile(promptPath, "COMPLETE_THEN_HANG\n");
+    await writeFile(promptPath, "COMPLETE_THEN_FLUSH\n");
     execution = new CodexSdkWorkerExecutor({
       parentSandbox: trustedParentSandbox
     }).run({
@@ -1037,14 +1037,15 @@ async function testCompletedWorkerSettlesWithoutWaitingForProcessExit() {
       new Promise((_, reject) => {
         timeout = setTimeout(() => {
           controller.abort("completed worker fixture timed out");
-          reject(new Error("completed worker did not settle after turn.completed"));
+          reject(new Error("completed worker did not settle after flushing usage"));
         }, 1_000);
       })
     ]);
     clearTimeout(timeout);
     assert.equal(result.threadId, "fixture-thread-id");
     assert.equal(result.finalResponse, "fixture final response");
     childPid = JSON.parse(await readFile(fixture.markerPath, "utf8")).pid;
+    assert.equal(await readFile(fixture.completionMarkerPath, "utf8"), "flushed\n");
 
     controller.abort("coordinator immediately canceled its remaining workers");
     await new Promise((resolve) => setTimeout(resolve, 250));
@@ -1337,6 +1338,7 @@ async function fakeCodexFixture(
   const root = await mkdtemp(path.join(tmpdir(), "codex-security-sdk-executor-"));
   temporaryRoots.push(root);
   const markerPath = path.join(root, "invocation.json");
+  const completionMarkerPath = path.join(root, "completed-usage.txt");
   const preflightMarkerPath = path.join(root, "preflight.json");
   const scriptPath = path.join(root, "fake-codex.mjs");
   await writeFile(scriptPath, [
@@ -1345,6 +1347,7 @@ async function fakeCodexFixture(
     `const preflightProfile = ${JSON.stringify(preflightProfile)};`,
     `const preflightAllowed = ${JSON.stringify(preflightAllowed)};`,
     `const preflightMarkerPath = ${JSON.stringify(preflightMarkerPath)};`,
+    `const completionMarkerPath = ${JSON.stringify(completionMarkerPath)};`,
     "if (process.argv.includes('app-server')) {",
     "  const preflight = { cwd: process.cwd(), codexHome: process.env.CODEX_HOME, requests: [] };",
     "  writeFileSync(preflightMarkerPath, JSON.stringify(preflight));",
@@ -1385,8 +1388,7 @@ async function fakeCodexFixture(
     "let stdin = '';",
     "for await (const chunk of process.stdin) stdin += chunk;",
     "const bedrockAuthentication = stdin.includes('CAPTURE_SYNTHETIC_BEDROCK_AUTH') ? Object.fromEntries(JSON.parse(process.env.FAKE_CODEX_BEDROCK_ENV_KEYS).map((name) => [name, process.env[name]])) : undefined;",
-    "writeFileSync(process.env.FAKE_CODEX_MARKER, JSON.stringify({ argv: process.argv.slice(2), stdin, cwd: process.cwd(), codexHome: process.env.CODEX_HOME, originator: process.env.CODEX_INTERNAL_ORIGINATOR_OVERRIDE, ...(stdin.includes('COMPLETE_THEN_HANG') ? { pid: process.pid } : {}), ...(bedrockAuthentication ? { bedrockAuthentication } : {}) }));",
-    "if (stdin.includes('COMPLETE_THEN_HANG')) process.on('SIGTERM', () => setTimeout(() => process.exit(0), 100));",
+    "writeFileSync(process.env.FAKE_CODEX_MARKER, JSON.stringify({ argv: process.argv.slice(2), stdin, cwd: process.cwd(), codexHome: process.env.CODEX_HOME, originator: process.env.CODEX_INTERNAL_ORIGINATOR_OVERRIDE, ...(stdin.includes('COMPLETE_THEN_FLUSH') ? { pid: process.pid } : {}), ...(bedrockAuthentication ? { bedrockAuthentication } : {}) }));",
     "if (stdin.includes('THREAD_START_CONFIG_ERROR')) { console.error('Error: thread/start: thread/start failed: agents.max_threads cannot be set when features.multi_agent_v2 is enabled (code -32600)'); process.exit(1); }",
     "if (stdin.includes('CONFIG_ERROR')) { console.error('failed to load configuration: invalid value'); process.exit(2); }",
     "if (stdin.includes('MCP_STARTUP_TIMEOUT') || stdin.includes('CATALOG_AUTH_ONLY') || stdin.includes('SYNC_AUTH_ONLY')) {",
@@ -1426,13 +1428,13 @@ async function fakeCodexFixture(
     "}",
     "console.log(JSON.stringify({ type: 'item.completed', item: { id: 'message-1', type: 'agent_message', text: 'fixture final response' } }));",
     "console.log(JSON.stringify({ type: 'turn.completed', usage: { input_tokens: 1, cached_input_tokens: 0, output_tokens: 1 } }));",
-    "if (stdin.includes('COMPLETE_THEN_HANG')) { setInterval(() => {}, 1_000); await new Promise(() => {}); }",
+    "if (stdin.includes('COMPLETE_THEN_FLUSH')) { await new Promise((resolve) => setTimeout(resolve, 100)); writeFileSync(completionMarkerPath, 'flushed\\n'); }",
     "}",
     ""
   ].join("\n"));
   await chmod(scriptPath, 0o755);
   process.env.FAKE_CODEX_MARKER = markerPath;
-  return { root, markerPath, preflightMarkerPath, executablePath: scriptPath };
+  return { root, markerPath, completionMarkerPath, preflightMarkerPath, executablePath: scriptPath };
 }
 
 function assertFlagPair(args, flag, value) {

plugins/codex-security/scripts/workbench_scan_usage.py

@@ -456,17 +456,29 @@ 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
             if timestamp < started_at:
                 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:
@@ -501,19 +513,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:

sdk/typescript/src/api.ts

@@ -151,6 +151,7 @@ import {
   resolveCodexCommand,
   resolvePluginPath,
   resolvePluginPython,
+  resolveScanSessionPaths,
   runWorkbench,
   setCodexSecurityCredentialLogout,
   type CodexCommand,
@@ -410,6 +411,7 @@ interface ClientDependencies {
     signal?: AbortSignal,
   ) => Promise<PreparedRuntime>;
   resolvePluginPython?: typeof resolvePluginPython;
+  resolveScanSessionPaths?: typeof resolveScanSessionPaths;
   prepareOutputDir?: typeof prepareOutputDir;
   prepareScanArtifactRestorer?: typeof prepareScanArtifactRestorer;
   repositoryRevision?: typeof repositoryRevision;
@@ -789,6 +791,7 @@ export class CodexSecurity {
     let scanFailure = false;
     let customValidationComplete = false;
     let completionCost: ScanCost | null = null;
+    let signaledCostUsage: unknown;
     let budgetRecovery: {
       expectation: ScanExpectation;
       pluginRoot: string;
@@ -968,8 +971,8 @@ export class CodexSecurity {
         pluginVersion: runtime.plugin.version,
       };
       const { model } = scanModelConfiguration(effectiveConfig);
-      validateScanCostLimit(options.maxCostUsd, model);
-      if (mode === "deep" && options.maxCostUsd !== undefined) {
+      validateScanCostLimit(maxCostUsd, model);
+      if (mode === "deep" && maxCostUsd !== undefined) {
         budgetRecovery = {
           expectation,
           pluginRoot: runtime.plugin.installedRoot,
@@ -996,7 +999,7 @@ export class CodexSecurity {
         );
       };
       const reportTrackingError = (error: unknown): void => {
-        if (options.maxCostUsd !== undefined) {
+        if (maxCostUsd !== undefined) {
           costAbortController.abort(error);
           return;
         }
@@ -1012,7 +1015,22 @@ export class CodexSecurity {
         model,
         repository: repo,
         scanDirectory: scanDir,
-        maxCostUsd: options.maxCostUsd,
+        maxCostUsd: maxCostUsd,
+        resolveOwnedSessionPaths:
+          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
@@ -1036,9 +1054,9 @@ export class CodexSecurity {
         onProgress:
           options.onProgress === undefined ? undefined : reportProgress,
         onCost:
-          options.onCost === undefined && options.maxCostUsd === undefined
+          options.onCost === undefined && maxCostUsd === undefined
             ? undefined
-            : (cost) => {
+            : (cost, usage) => {
                 latestCost = cost;
                 notifyObserver(
                   "onCost",
@@ -1049,8 +1067,10 @@ export class CodexSecurity {
                 );
                 if (
                   maxCostUsd !== undefined &&
-                  cost.estimatedUsd > maxCostUsd
+                  cost.estimatedUsd > maxCostUsd &&
+                  !costAbortController.signal.aborted
                 ) {
+                  signaledCostUsage = usage;
                   costAbortController.abort(
                     new ScanCostLimitExceededError(maxCostUsd, cost, scanDir),
                   );
@@ -1101,6 +1121,7 @@ export class CodexSecurity {
                     );
                     if (budgetSignal.aborted) return;
                     maxCostUsd = next;
+                    tracker.setMaxCostUsd(next);
                     notifyObserver(
                       "onCost",
                       options.onCost,
@@ -1132,7 +1153,7 @@ export class CodexSecurity {
         { ...preflightConfig, approval_policy: approvalPolicy },
         options.failureSeverity,
         knowledgeBase?.sources,
-        options.maxCostUsd,
+        maxCostUsd,
         deepScanOptions(options),
       );
       if (options.validationPrompt !== undefined)
@@ -1290,7 +1311,7 @@ export class CodexSecurity {
         runtime.configPath !== undefined,
         knowledgeBase !== null,
         options.scanPrompt,
-        options.maxCostUsd !== undefined,
+        maxCostUsd !== undefined,
         discoveryPrompt,
       );
       checkOpen();
@@ -1447,6 +1468,8 @@ export class CodexSecurity {
           }
         },
         onFinalize: async (usage) => {
+          // Validation threads are accounted separately from the discovery turn.
+          const discoveryUsage = usage;
           if (options.validationPrompt !== undefined) {
             tracker.recordUsage(usage);
             await tracker.refresh().catch(reportTrackingError);
@@ -1506,13 +1529,24 @@ export class CodexSecurity {
             customValidationComplete = true;
           }
           budgetAbortController.abort();
-          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 (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) {
+          if (maxCostUsd !== undefined && snapshot.cost === null) {
             notifyObserver(
               "onWarning",
               options.onWarning,
@@ -1733,11 +1767,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);
-      let 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 (
+        maxCostUsd !== undefined &&
+        knownCost !== undefined &&
+        knownCost !== null &&
+        knownCost.estimatedUsd > maxCostUsd &&
+        (signaledOverage !== null ||
+          (!this.#abortController.signal.aborted &&
+            options.signal?.aborted !== true)) &&
+        (signaledOverage === null ||
+          knownCost.estimatedUsd > signaledOverage.cost.estimatedUsd)
+      ) {
+        failure = new ScanCostLimitExceededError(
+          maxCostUsd,
+          knownCost,
+          scanDir,
+        );
+      }
       if (
         failure instanceof ScanCostLimitExceededError &&
         snapshot?.cost &&