fix: correct underreported token totals in Codex Security scans

plugins/codex-security/mcp-app/server.ts

@@ -1532,14 +1532,24 @@ export function createCodexSecurityServer(): McpServer {
       annotations: writingAnnotations,
       _meta: modelActionMeta,
     },
-    async ({ scanId, handoffClaimToken }) => {
+    async ({ scanId, handoffClaimToken }, extra) => {
+      const turnMetadata = codexTurnMetadataFromExtra(extra);
+      const turnId =
+        turnMetadata?.thread_source === "security_scan"
+          ? turnMetadata.turn_id
+          : undefined;
       try {
         return scanActionResult(
           await runWorkbench([
             "complete-scan",
             "--scan-id",
             scanId,
             ...optionalArg("--claim-token", handoffClaimToken),
+            ...optionalArg("--thread-id", threadIdFromExtra(extra)),
+            ...optionalArg(
+              "--turn-id",
+              typeof turnId === "string" ? turnId : undefined,
+            ),
           ]),
           "Validated and indexed the completed Codex Security scan.",
         );
@@ -1549,6 +1559,35 @@ export function createCodexSecurityServer(): McpServer {
     },
   );
 
+  server.registerTool(
+    "finalize_codex_security_scan_usage",
+    {
+      title: "Finalize Codex Security Local Scan Usage",
+      description:
+        "App-only. Record final measured usage after the scan's completion turn has finished. Pass the exact thread and turn returned in usage.completion; do not call while that turn is running.",
+      inputSchema: {
+        ...scanSchema,
+        threadId: z.string().trim().min(1),
+        turnId: z.string().trim().min(1),
+        turnStatus: z.enum(["completed", "interrupted", "failed"]),
+      },
+      annotations: writingAnnotations,
+      _meta: appMeta,
+    },
+    async ({ scanId, threadId, turnId, turnStatus }) =>
+      scanActionResult(
+        await runWorkbench(
+          workbenchArgs("finalize-scan-usage", {
+            scanId,
+            threadId,
+            turnId,
+            turnStatus,
+          }),
+        ),
+        "Recorded final Codex Security scan usage.",
+      ),
+  );
+
   server.registerTool(
     "fail_codex_security_scan",
     {
@@ -2218,18 +2257,25 @@ function threadIdFromExtra(extra: unknown): string | undefined {
   const thread = isJsonObject(metadata?.thread) ? metadata.thread : undefined;
   if (typeof thread?.id === "string" && thread.id.trim())
     return thread.id.trim();
-  return undefined;
+  const threadId = codexTurnMetadataFromExtra(extra)?.thread_id;
+  return typeof threadId === "string" && threadId.trim()
+    ? threadId.trim()
+    : undefined;
+}
+
+function codexTurnMetadataFromExtra(extra: unknown): JsonObject | undefined {
+  const metadata = requestMetadataFromExtra(extra);
+  return isJsonObject(metadata?.["x-codex-turn-metadata"])
+    ? metadata["x-codex-turn-metadata"]
+    : undefined;
 }
 
 function codexModelSettingsFromExtra(extra: unknown): {
   model?: string;
   reasoningEffort?: string;
 } {
-  const metadata = requestMetadataFromExtra(extra);
   // Codex supplies the already-resolved runtime model and effective effort on every MCP call.
-  const turnMetadata = isJsonObject(metadata?.["x-codex-turn-metadata"])
-    ? metadata["x-codex-turn-metadata"]
-    : undefined;
+  const turnMetadata = codexTurnMetadataFromExtra(extra);
   const model =
     typeof turnMetadata?.model === "string"
       ? turnMetadata.model.trim()

plugins/codex-security/mcp-app/tests/test_mcp_app_smoke.ts

@@ -539,7 +539,9 @@ async function assertWorkspaceWorksWithoutUiCapability(scope: string) {
   }
 }
 
-async function assertHeadlessStandardScanWorksWithoutUiCapability() {
+async function assertHeadlessStandardScanWorksWithoutUiCapability(
+  threadSource?: string,
+) {
   const headlessStateDir = path.join(tmpdir(), randomUUID());
   const headlessScanRoot = path.join(tmpdir(), randomUUID());
   const headlessServer = startTestServer({
@@ -690,13 +692,48 @@ async function assertHeadlessStandardScanWorksWithoutUiCapability() {
         scanId: result.scanId,
         handoffClaimToken: result.handoffClaimToken,
       },
-      _meta: { "openai/threadId": ownerThread },
+      _meta: {
+        "x-codex-turn-metadata": {
+          thread_id: ownerThread,
+          turn_id: "fixture-completion-turn",
+          thread_source: threadSource,
+        },
+      },
     });
     assertNoError(completed);
     assert.equal(
       completed.result.structuredContent.scan.progress.status,
       "complete",
     );
+    if (threadSource === "security_scan") {
+      assert.deepEqual(
+        completed.result.structuredContent.scan.usage.completion,
+        {
+          threadId: ownerThread,
+          turnId: "fixture-completion-turn",
+          finalized: false,
+        },
+      );
+      const finalizedUsage = await headlessServer.callTool(10, {
+        name: "finalize_codex_security_scan_usage",
+        arguments: {
+          scanId: result.scanId,
+          threadId: ownerThread,
+          turnId: "fixture-completion-turn",
+          turnStatus: "completed",
+        },
+      });
+      assertNoError(finalizedUsage);
+      assert.deepEqual(
+        finalizedUsage.result.structuredContent.scan.usage,
+        completed.result.structuredContent.scan.usage,
+      );
+    } else {
+      assert.equal(
+        completed.result.structuredContent.scan.usage.completion,
+        undefined,
+      );
+    }
   } finally {
     await headlessServer.stop();
     await rm(headlessStateDir, { recursive: true, force: true });
@@ -1452,6 +1489,7 @@ try {
   ])
     await assertWorkspaceWorksWithoutUiCapability(scope);
   await assertHeadlessStandardScanWorksWithoutUiCapability();
+  await assertHeadlessStandardScanWorksWithoutUiCapability("security_scan");
   await assertDeepScanPersistsRetryableWorkerStartupError();
   await assertUserInputFailureLogging();
   if (process.platform !== "win32") {
@@ -1714,6 +1752,7 @@ try {
   );
   const progress = findTool("update_codex_security_scan_progress");
   const complete = findTool("complete_codex_security_scan");
+  const finalizeUsage = findTool("finalize_codex_security_scan_usage");
   const fail = findTool("fail_codex_security_scan");
   const setFindingTriage = findTool("set_codex_security_finding_triage");
   const requestFindingRemediation = findTool(
@@ -1912,6 +1951,7 @@ try {
   assert.deepEqual(attachHandoff._meta.ui.visibility, ["app"]);
   assert.deepEqual(progress._meta.ui.visibility, ["model"]);
   assert.deepEqual(complete._meta.ui.visibility, ["model"]);
+  assert.deepEqual(finalizeUsage._meta.ui.visibility, ["app"]);
   assert.deepEqual(fail._meta.ui.visibility, ["model"]);
   for (const tool of [
     launcher,

plugins/codex-security/references/final-report.md

@@ -30,7 +30,7 @@ During a host-backed scan, checkpoint saved findings and pending candidates with
 
 On a confirmed failure or explicit cancellation, the workbench preserves saved results without marking the scan successful. Report validated findings separately from pending candidates and describe scope limitations and remaining work. Use the retained report and exports; do not start additional model work after cancellation or replace saved results with a no-findings response.
 
-After `complete_codex_security_scan` succeeds, include its returned `usage.totalTokens`, `usage.inputTokens`, and `usage.cachedInputTokens` in the final response when `usage.coverage` is `complete` or `partial`; explicitly label a partial measurement. If coverage is `unavailable`, say that token usage could not be measured instead of reporting zero or estimating a cost. Report only measured completion metadata in a terminal/chat host. Token usage is workbench metadata, not a reason to modify sealed scan artifacts or the deterministic report.
+After `complete_codex_security_scan` succeeds, include its returned `usage.totalTokens`, `usage.inputTokens`, and `usage.cachedInputTokens` only when `usage.coverage` is `complete`. When `usage.completion.finalized` is `false`, leave the token total to the host: it finalizes usage after this turn ends, including the final response. Do not quote the preliminary count as a final total or an "at least" estimate. For other partial or unavailable measurements, say that the full token count is unavailable rather than reporting zero or estimating a cost. Token usage is workbench metadata, not a reason to modify sealed scan artifacts or the deterministic report.
 
 Before workbench-owned Standard or workbench-backed diff completion, require `record_codex_security_scan_draft` to succeed. Before Deep completion, require the coordinator to return the parent scan's already-authored canonical manifest; do not submit another draft or rerun worker phases. Explicitly SDK-owned workflows verify their canonical JSON before the SDK finalizes it. Completion validates and seals existing canonical artifacts and generates `report.md`; it does not create missing artifacts or run skipped scan phases.
 

plugins/codex-security/scripts/workbench_cli.py

@@ -208,6 +208,14 @@ def add_command(name: str, *required: str) -> argparse.ArgumentParser:
     complete_scan.add_argument("--claim-token")
     complete_scan.add_argument("--cost-json")
     complete_scan.add_argument("--thread-id")
+    complete_scan.add_argument("--turn-id")
+
+    finalize_scan_usage = add_command(
+        "finalize-scan-usage", "--scan-id", "--thread-id", "--turn-id"
+    )
+    finalize_scan_usage.add_argument(
+        "--turn-status", choices=("completed", "interrupted", "failed"), default="completed"
+    )
 
     complete_budget_exhausted_scan = add_command(
         "complete-budget-exhausted-scan", "--scan-id", "--cost-json"

plugins/codex-security/scripts/workbench_db.py

@@ -1162,9 +1162,43 @@ def complete_scan(
             cost_json,
             prepare_only=prepare_only,
             thread_id=getattr(args, "thread_id", None),
+            turn_id=getattr(args, "turn_id", None),
         )
 
 
+def finalize_scan_usage(connection: sqlite3.Connection, args: argparse.Namespace) -> dict[str, Any]:
+    scan_id = require_uuid(args.scan_id, "scan-id")
+    with scan_completion_lock(scan_id):
+        scan = require_scan(connection, scan_id)
+        stored = json.loads(scan["cost_json"]) if scan["cost_json"] is not None else {}
+        usage = stored.get("usage", {})
+        completion = usage.get("completion")
+        if scan["status"] != "complete" or not completion:
+            raise SystemExit("The scan has no pending completion usage.")
+        if completion["threadId"] != args.thread_id or completion["turnId"] != args.turn_id:
+            raise SystemExit("The scan usage belongs to a different completion turn.")
+        if not completion["finalized"]:
+            usage = scan_usage.collect_scan_usage(
+                connection,
+                scan,
+                thread_id=args.thread_id,
+                completed_turn_id=args.turn_id,
+            )
+            if usage["coverage"] == "unavailable":
+                return scan_context(connection, scan_id)
+            if args.turn_status != "completed":
+                if usage["coverage"] == "complete":
+                    usage["coverage"] = "partial"
+                usage["warnings"] = [*usage.get("warnings", []), f"scan_turn_{args.turn_status}"]
+            usage["completion"] = {**completion, "finalized": True}
+            cost_json = parse_scan_cost(json.dumps({**stored, "usage": usage}))
+            with connection:
+                connection.execute(
+                    "UPDATE scans SET cost_json = ? WHERE id = ?", (cost_json, scan_id)
+                )
+        return scan_context(connection, scan_id)
+
+
 def complete_budget_exhausted_scan(
     connection: sqlite3.Connection, args: argparse.Namespace
 ) -> dict[str, Any]:
@@ -1454,6 +1488,7 @@ def complete_scan_locked(
     *,
     prepare_only: bool = False,
     thread_id: str | None = None,
+    turn_id: str | None = None,
 ) -> dict[str, Any]:
     scan = require_scan(connection, scan_id)
     if scan["status"] == "complete":
@@ -1610,6 +1645,18 @@ def add_warning() -> None:
             thread_id=thread_id,
             completed_at=completion_timestamp,
         )
+        if thread_id is not None and turn_id is not None:
+            measured_usage["completion"] = {
+                "threadId": thread_id,
+                "turnId": turn_id,
+                "finalized": False,
+            }
+            if measured_usage["coverage"] == "complete":
+                measured_usage["coverage"] = "partial"
+            measured_usage["warnings"] = [
+                *measured_usage.get("warnings", []),
+                "turn_in_progress",
+            ]
         cost_json = parse_scan_cost(scan_usage.measured_scan_cost_json(measured_usage))
     connection.execute("BEGIN IMMEDIATE")
     with connection:
@@ -3386,6 +3433,7 @@ def main(*, before_archive: Callable[[], None] | None = None) -> None:
             "register-cli-scan": partial(register_cli_scan, before_archive=before_archive),
             "set-scan-thread": set_scan_thread,
             "set-scan-cost-limit": set_scan_cost_limit,
+            "finalize-scan-usage": finalize_scan_usage,
             "list-global-findings": native_indexes.list_global_findings,
             "list-repositories": native_indexes.list_repositories,
             "list-findings": list_findings,