fix(cost): enforce verifiable scan spending limits

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"

sdk/typescript/_bundled_plugin/scripts/workbench_scan_usage.py

@@ -471,17 +471,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:
@@ -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:

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<PreparedRuntime>;
   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 &&