feat(cli): allow interactive scan budget increases

.github/workflows/node-ci.yml

@@ -437,6 +437,10 @@ jobs:
         uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           persist-credentials: false
+      - name: Set up Python
+        uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
+        with:
+          python-version: "3.12"
       - name: Set up Node.js
         uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # v6
         with:

plugins/codex-security/scripts/workbench_cli.py

@@ -160,6 +160,10 @@ def parse_args(description: str) -> argparse.Namespace:
     set_scan_thread.add_argument("--scan-id", required=True)
     set_scan_thread.add_argument("--thread-id", required=True)
 
+    set_scan_cost_limit = subparsers.add_parser("set-scan-cost-limit")
+    set_scan_cost_limit.add_argument("--scan-id", required=True)
+    set_scan_cost_limit.add_argument("--max-cost-usd", required=True, type=float)
+
     get_scan_recipe = subparsers.add_parser("get-scan-recipe")
     get_scan_recipe.add_argument("--scan-id", required=True)
 

plugins/codex-security/scripts/workbench_db.py

@@ -7,6 +7,7 @@
 import errno
 import hashlib
 import json
+import math
 import os
 import re
 import sqlite3
@@ -1789,6 +1790,31 @@ def set_scan_thread(connection: sqlite3.Connection, args: argparse.Namespace) ->
     return {"scanId": scan["id"], "threadId": args.thread_id}
 
 
+def set_scan_cost_limit(connection: sqlite3.Connection, args: argparse.Namespace) -> dict[str, Any]:
+    scan_id = require_uuid(args.scan_id, "scan-id")
+    limit = args.max_cost_usd
+    if not math.isfinite(limit) or limit <= 0:
+        raise SystemExit("The scan cost limit must be a positive finite USD amount.")
+    with scan_completion_lock(scan_id), connection:
+        scan = require_scan(connection, scan_id)
+        if scan["status"] != "running" or scan["recipe_json"] is None:
+            raise SystemExit("Only a running CLI scan can increase its cost limit.")
+        recipe = json.loads(scan["recipe_json"], parse_constant=reject_non_finite_json)
+        previous = recipe.get("maxCostUsd")
+        if (
+            not isinstance(previous, (int, float))
+            or isinstance(previous, bool)
+            or limit <= previous
+        ):
+            raise SystemExit("The new cost limit must exceed the current limit.")
+        recipe["maxCostUsd"] = limit
+        connection.execute(
+            "UPDATE scans SET recipe_json = ?, updated_at = ? WHERE id = ?",
+            (json.dumps(recipe, allow_nan=False), now(), scan["id"]),
+        )
+    return {"scanId": scan["id"], "maxCostUsd": limit}
+
+
 def parse_scan_recipe(value: str, repository: Path) -> dict[str, Any]:
     if len(value.encode("utf-8")) > SCAN_RECIPE_MAX_BYTES:
         raise SystemExit("Scan launch recipe must be no larger than 256 KiB.")
@@ -3468,6 +3494,8 @@ def main() -> None:
             result = register_cli_scan(connection, args)
         elif args.command == "set-scan-thread":
             result = set_scan_thread(connection, args)
+        elif args.command == "set-scan-cost-limit":
+            result = set_scan_cost_limit(connection, args)
         elif args.command == "get-scan-recipe":
             result = get_scan_recipe(connection, args)
         elif args.command == "compare-scans":

plugins/codex-security/tests/test_workbench_db.py

@@ -181,6 +181,54 @@ def complete_budget_scan(state_dir: Path, scan_id: str, *, check: bool = True) -
     )
 
 
+def test_cost_limit_increases_are_saved_without_replacing_the_scan_recipe(
+    tmp_path: Path,
+) -> None:
+    state_dir, _, _, scan_id, _ = budget_scan_fixture(tmp_path)
+    original = run_workbench(state_dir, "get-scan-recipe", "--scan-id", scan_id)["recipe"]
+    for limit in (0.0055, 0.006):
+        run_workbench(
+            state_dir,
+            "set-scan-cost-limit",
+            "--scan-id",
+            scan_id,
+            "--max-cost-usd",
+            str(limit),
+        )
+    saved = run_workbench(state_dir, "get-scan-recipe", "--scan-id", scan_id)["recipe"]
+    assert saved == {**original, "maxCostUsd": 0.006}
+    assert complete_budget_scan(state_dir, scan_id)["scan"]["progress"]["status"] == "complete"
+    stopped = run_workbench(
+        state_dir,
+        "set-scan-cost-limit",
+        "--scan-id",
+        scan_id,
+        "--max-cost-usd",
+        "1",
+        check=False,
+    )
+    assert stopped["returncode"] != 0
+
+
+@pytest.mark.parametrize("limit", ["0", "-1", "nan", "inf", "0.004", "0.005"])
+def test_cost_limit_rejects_invalid_or_nonincreasing_totals(tmp_path: Path, limit: str) -> None:
+    state_dir, _, _, scan_id, _ = budget_scan_fixture(tmp_path, mode="standard")
+    result = run_workbench(
+        state_dir,
+        "set-scan-cost-limit",
+        "--scan-id",
+        scan_id,
+        "--max-cost-usd",
+        limit,
+        check=False,
+    )
+    assert result["returncode"] != 0
+    assert (
+        run_workbench(state_dir, "get-scan-recipe", "--scan-id", scan_id)["recipe"]["maxCostUsd"]
+        == 0.005
+    )
+
+
 def test_budget_exhaustion_preserves_unvalidated_discovery_as_deferred_work(
     tmp_path: Path,
 ) -> None:

sdk/typescript/README.md

@@ -519,6 +519,27 @@ discovery has finished, the scan returns a sealed partial report without more
 model calls and lists unvalidated candidates as follow-up work. Bulk scans
 apply the limit per repository attempt.
 
+For a single scan in the interactive dashboard, reaching 80% of the limit
+offers a higher **total** USD limit. Enter a larger amount to approve it, or
+press Enter with an empty input or Escape to keep the current limit. The scan
+continues running while you decide, and the existing limit remains enforced
+until the increase is saved. Increases keep the same scan and accumulated cost;
+they do not restart work or extend time or discovery limits. CI, JSON/JSONL,
+`--headless`, and `--verbose` scans do not offer budget increases. If usage crosses the limit
+before an increase is approved, the scan still stops.
+
+SDK callers can supply `onBudgetApproaching({ maxCostUsd, cost, signal })` and
+return a higher total limit, or `undefined` to keep the current limit. The
+callback runs once per limit at 80% usage without blocking tracking or
+execution. Its signal aborts when the scan stops or finishes model work; late
+answers are ignored. Invalid increases or failures to save them leave the
+existing limit in place and report a warning. `onCost(cost, maxCostUsd)` reports
+the current limit, including after an approved increase.
+
+These amounts estimate API-equivalent model usage, not ChatGPT subscription
+allowance. Post-scan prompts run after scan cost tracking ends and are outside
+this limit.
+
 ### Bulk scans
 
 Run `gh auth login`, then `npx @openai/codex-security bulk-scan` to select

sdk/typescript/src/api.ts

@@ -60,7 +60,6 @@ import {
 import {
   estimateScanCost,
   ScanCostTracker,
-  sumTokenUsage,
   type ScanCost,
   type ScanSessionEvent,
 } from "./cost.js";
@@ -244,7 +243,10 @@ export interface ScanOptions extends DeepScanOptions {
   expectedPluginVersion?: string;
   failureSeverity?: SeverityLevel;
   maxCostUsd?: number;
-  onCost?: (cost: Readonly<ScanCost>) => void;
+  onCost?: (cost: Readonly<ScanCost>, maxCostUsd?: number) => void;
+  onBudgetApproaching?: (
+    budget: ScanBudget,
+  ) => number | undefined | Promise<number | undefined>;
   onOutputArchived?: (archiveDir: string) => void;
   onOutputDirReady?: (scanDir: string) => void;
   onAuthentication?: (authentication: ScanAuthentication) => void;
@@ -335,6 +337,12 @@ export interface ScanWarningDetails {
   kind: "target_changed";
 }
 
+export interface ScanBudget {
+  maxCostUsd: number;
+  cost: Readonly<ScanCost>;
+  signal: AbortSignal;
+}
+
 type ScanObserverName =
   | "onAuthentication"
   | "onCost"
@@ -750,6 +758,14 @@ export class CodexSecurity {
       costAbortController.signal,
       ...(options.signal === undefined ? [] : [options.signal]),
     ]);
+    const budgetAbortController = new AbortController();
+    const budgetSignal = AbortSignal.any([
+      signal,
+      budgetAbortController.signal,
+    ]);
+    let maxCostUsd = options.maxCostUsd;
+    let latestCost: Readonly<ScanCost> | null = null;
+    let notifiedLimit: number | undefined;
     let scanDir = "";
     let archivedScanDir: string | null = null;
     let targetPathsFile: string | null = null;
@@ -1010,24 +1026,86 @@ export class CodexSecurity {
           options.onCost === undefined && options.maxCostUsd === undefined
             ? undefined
             : (cost) => {
+                latestCost = cost;
                 notifyObserver(
                   "onCost",
                   options.onCost,
                   options.onObserverError,
                   cost,
+                  maxCostUsd,
                 );
                 if (
-                  options.maxCostUsd !== undefined &&
-                  cost.estimatedUsd > options.maxCostUsd
+                  maxCostUsd !== undefined &&
+                  cost.estimatedUsd > maxCostUsd
                 ) {
                   costAbortController.abort(
-                    new ScanCostLimitExceededError(
-                      options.maxCostUsd,
-                      cost,
-                      scanDir,
-                    ),
+                    new ScanCostLimitExceededError(maxCostUsd, cost, scanDir),
                   );
+                  return;
                 }
+                const request = options.onBudgetApproaching;
+                if (
+                  request === undefined ||
+                  maxCostUsd === undefined ||
+                  budgetSignal.aborted ||
+                  notifiedLimit === maxCostUsd ||
+                  cost.estimatedUsd < maxCostUsd * 0.8
+                )
+                  return;
+                const limit = maxCostUsd;
+                notifiedLimit = limit;
+                void Promise.resolve()
+                  .then(async () => {
+                    if (budgetSignal.aborted) return;
+                    const next = await request({
+                      maxCostUsd: limit,
+                      cost,
+                      signal: budgetSignal,
+                    });
+                    if (
+                      next === undefined ||
+                      budgetSignal.aborted ||
+                      activeScan === null
+                    )
+                      return;
+                    if (
+                      !Number.isFinite(next) ||
+                      next <= Math.max(limit, latestCost!.estimatedUsd)
+                    ) {
+                      throw new CodexSecurityError(
+                        "The new cost limit must exceed the current limit and estimated cost.",
+                      );
+                    }
+                    await workbench(
+                      { ...activeScan.options, signal: budgetSignal },
+                      [
+                        "set-scan-cost-limit",
+                        "--scan-id",
+                        activeScan.id,
+                        "--max-cost-usd",
+                        String(next),
+                      ],
+                    );
+                    if (budgetSignal.aborted) return;
+                    maxCostUsd = next;
+                    notifyObserver(
+                      "onCost",
+                      options.onCost,
+                      options.onObserverError,
+                      latestCost!,
+                      maxCostUsd,
+                    );
+                  })
+                  .catch((error: unknown) => {
+                    if (!budgetSignal.aborted) {
+                      notifyObserver(
+                        "onWarning",
+                        options.onWarning,
+                        options.onObserverError,
+                        `Could not increase scan cost limit: ${errorMessage(error)}`,
+                      );
+                    }
+                  });
               },
         onError: reportTrackingError,
       });
@@ -1357,6 +1435,9 @@ export class CodexSecurity {
         },
         onFinalize: async (usage) => {
           if (options.validationPrompt !== undefined) {
+            tracker.recordUsage(usage);
+            await tracker.refresh().catch(reportTrackingError);
+            checkOpen();
             await runCustomValidation({
               repository: repo,
               target: normalized,
@@ -1402,12 +1483,16 @@ export class CodexSecurity {
                     turn.lastStreamError ??
                       "The custom validation turn did not complete.",
                   );
-                usage = sumTokenUsage(usage, turn.usage);
+                budgetAbortController.abort();
+                tracker.recordUsage(turn.usage, turn.threadId);
+                await tracker.refresh().catch(reportTrackingError);
+                checkOpen();
                 return turn.finalResponse;
               },
             });
             customValidationComplete = true;
           }
+          budgetAbortController.abort();
           const snapshot = await tracker.stop(usage).catch((error: unknown) => {
             if (options.maxCostUsd !== undefined) throw error;
             reportTrackingError(error);
@@ -1637,10 +1722,21 @@ export class CodexSecurity {
       // 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 =
+      let failure =
         signal.reason instanceof ScanCostLimitExceededError
           ? signal.reason
           : error;
+      if (
+        failure instanceof ScanCostLimitExceededError &&
+        snapshot?.cost &&
+        snapshot.cost.estimatedUsd > failure.cost.estimatedUsd
+      ) {
+        failure = new ScanCostLimitExceededError(
+          failure.maxCostUsd,
+          snapshot.cost,
+          scanDir,
+        );
+      }
       if (
         failure instanceof ScanCostLimitExceededError &&
         budgetRecovery !== null &&
@@ -1763,6 +1859,7 @@ export class CodexSecurity {
       }
       throw failure;
     } finally {
+      budgetAbortController.abort();
       deepProgressTracker?.stop();
       // Removing the temporary scan inputs is best effort. A throw here would replace the
       // outcome the try and catch blocks already produced, so these failures are reported