diff --git a/plugins/codex-security/mcp-app/server.ts b/plugins/codex-security/mcp-app/server.ts index cd80f8b56..0b71cd0c5 100644 --- a/plugins/codex-security/mcp-app/server.ts +++ b/plugins/codex-security/mcp-app/server.ts @@ -1532,7 +1532,12 @@ 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([ @@ -1540,6 +1545,11 @@ export function createCodexSecurityServer(): McpServer { "--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() diff --git a/plugins/codex-security/mcp-app/tests/test_mcp_app_smoke.ts b/plugins/codex-security/mcp-app/tests/test_mcp_app_smoke.ts index 6038e18ec..c22656e72 100644 --- a/plugins/codex-security/mcp-app/tests/test_mcp_app_smoke.ts +++ b/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, diff --git a/plugins/codex-security/references/final-report.md b/plugins/codex-security/references/final-report.md index 831fd9e58..7cc973d3d 100644 --- a/plugins/codex-security/references/final-report.md +++ b/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. diff --git a/plugins/codex-security/scripts/workbench_cli.py b/plugins/codex-security/scripts/workbench_cli.py index c3928f886..462f144d7 100644 --- a/plugins/codex-security/scripts/workbench_cli.py +++ b/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" diff --git a/plugins/codex-security/scripts/workbench_db.py b/plugins/codex-security/scripts/workbench_db.py index 3bd233040..c4671a61d 100644 --- a/plugins/codex-security/scripts/workbench_db.py +++ b/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, diff --git a/plugins/codex-security/scripts/workbench_scan_usage.py b/plugins/codex-security/scripts/workbench_scan_usage.py index 35089c5f1..e46a52cac 100644 --- a/plugins/codex-security/scripts/workbench_scan_usage.py +++ b/plugins/codex-security/scripts/workbench_scan_usage.py @@ -84,11 +84,13 @@ def collect_scan_usage( *, thread_id: str | None = None, completed_at: str | None = None, + completed_turn_id: str | None = None, ) -> dict[str, Any]: """Count only complete, attributable rollout events inside this scan's window.""" roots = _scan_root_thread_ids(connection, scan, thread_id) - if not roots: + worker_directories = _scan_worker_directories(connection, scan) + if not roots and not worker_directories: return _unavailable_usage("scan_thread_unavailable") state_database = _codex_state_database() @@ -106,6 +108,9 @@ def collect_scan_usage( state_database, roots, warnings, + worker_directories=worker_directories, + started_at=started_at, + completed_at=stopped_at, ) except (OSError, sqlite3.Error, ValueError): return _unavailable_usage("codex_state_unavailable") @@ -132,6 +137,7 @@ def collect_scan_usage( session, started_at=started_at, completed_at=stopped_at, + completed_turn_id=completed_turn_id if session.thread_id == thread_id else None, ) except (OSError, UnicodeError, ValueError): missing_thread_ids.add(session.thread_id) @@ -214,6 +220,21 @@ def _scan_execution_thread_ids(connection: sqlite3.Connection, scan: sqlite3.Row ) +def _scan_worker_directories(connection: sqlite3.Connection, scan: sqlite3.Row) -> set[str]: + if scan["mode"] != "deep": + return set() + return { + row["artifact_dir"] + if row["kind"] == "discovery" + else str(Path(scan["scan_dir"]) / "artifacts") + for row in connection.execute( + "SELECT kind, artifact_dir FROM deep_scan_workers " + "WHERE scan_id = ? AND kind IN ('discovery', 'dedup')", + (scan["id"],), + ) + } + + def _codex_state_database() -> Path | None: configured_database = os.environ.get("CODEX_STATE_DB", "").strip() if configured_database: @@ -255,6 +276,10 @@ def _discover_rollout_sessions( state_database: Path, roots: list[str], warnings: set[str], + *, + worker_directories: set[str], + started_at: datetime, + completed_at: datetime | None, ) -> tuple[list[RolloutSession], set[str]]: database = sqlite3.connect( f"file:{quote(str(state_database), safe='')}?mode=ro", @@ -273,7 +298,27 @@ def _discover_rollout_sessions( raise ValueError("Codex state graph does not expose the required thread columns.") sessions: dict[str, RolloutSession] = {} missing_thread_ids: set[str] = set() - for root in roots: + root_ids = dict.fromkeys(roots) + started_millisecond = started_at.replace(microsecond=started_at.microsecond // 1000 * 1000) + completed_timestamp = int(completed_at.timestamp()) if completed_at is not None else None + directory_expression = "cwd" + if worker_directories and os.name == "nt": + from windows_scan_local_files import _normalized_windows_path + + database.create_function( + "codex_security_path_key", 1, _normalized_windows_path, deterministic=True + ) + directory_expression = "codex_security_path_key(cwd)" + worker_directories = {_normalized_windows_path(value) for value in worker_directories} + for directory in sorted(worker_directories): + for row in database.execute( + "SELECT id FROM threads WHERE cwd IS NOT NULL " + f"AND {directory_expression} = ? AND created_at >= ? " + "AND (? IS NULL OR created_at <= ?) ORDER BY id", + (directory, int(started_at.timestamp()), completed_timestamp, completed_timestamp), + ): + root_ids[row["id"]] = None + for root in root_ids: row = database.execute( "SELECT id, rollout_path FROM threads WHERE id = ?", (root,), @@ -282,13 +327,29 @@ def _discover_rollout_sessions( missing_thread_ids.add(root) warnings.add("scan_root_unavailable") continue - if root not in sessions: - path = _rollout_path(row["rollout_path"]) - if path is None: + path = _rollout_path(row["rollout_path"]) + if path is None: + missing_thread_ids.add(root) + warnings.add("rollout_unavailable") + continue + if root not in roots: + try: + with path.open("rb") as source: + metadata = json.loads(source.readline()) + except (OSError, ValueError): missing_thread_ids.add(root) warnings.add("rollout_unavailable") continue - sessions[root] = RolloutSession(root, None, path) + if ( + isinstance(metadata, dict) + and metadata.get("type") == "session_meta" + and isinstance(payload := metadata.get("payload"), dict) + and (payload.get("id") or payload.get("session_id")) == root + and (created_at := _timestamp(payload.get("timestamp"))) is not None + and created_at < started_millisecond + ): + continue + sessions[root] = RolloutSession(root, None, path) descendants = database.execute( """ WITH RECURSIVE descendants( @@ -378,11 +439,18 @@ def _read_rollout_usage( *, started_at: datetime, completed_at: datetime | None, + completed_turn_id: str | None = None, ) -> tuple[dict[str, int], set[str]]: total = _empty_token_usage() warnings: set[str] = set() previous = _empty_token_usage() boundary_reached = False + legacy_usage = _empty_token_usage() + response_usage: dict[str, int] | None = None + observed_response_usage = _empty_token_usage() + response_ids: set[str] = set() + current_turn_id: str | None = None + completed_turn_measured = False with session.path.open("rb") as source: for line_number, raw_line in enumerate(source, start=1): @@ -431,16 +499,100 @@ def _read_rollout_usage( warnings.add("thread_ownership_unavailable") return total, warnings if task_started_at < started_at or ( - completed_at is not None and task_started_at > completed_at + completed_at is not None + and task_started_at > completed_at + and payload.get("turn_id") != completed_turn_id ): warnings.add("thread_outside_scan_window") return total, warnings boundary_reached = True + current_turn_id = payload.get("turn_id") elif event.get("type") == "event_msg" and payload.get("type") == "token_count": inherited_usage = _token_snapshot(payload) if inherited_usage is not None: previous = inherited_usage continue + if event.get("type") == "event_msg" and payload.get("type") == "task_started": + if ( + completed_turn_id is not None + and current_turn_id == completed_turn_id + and payload.get("turn_id") != completed_turn_id + ): + break + _add_token_usage( + total, response_usage if response_usage is not None else legacy_usage + ) + legacy_usage = _empty_token_usage() + response_usage = None + if payload.get("turn_id") != current_turn_id: + observed_response_usage = _empty_token_usage() + current_turn_id = payload.get("turn_id") + continue + if ( + completed_turn_id is not None + and event.get("type") == "event_msg" + and payload.get("type") == "task_complete" + and payload.get("turn_id") == completed_turn_id + ): + break + if event.get("type") == "compacted": + timestamp = _timestamp(event.get("timestamp")) + if ( + timestamp is not None + and timestamp >= started_at + and ( + completed_at is None + or timestamp <= completed_at + or (completed_turn_id is not None and current_turn_id == completed_turn_id) + ) + and "compaction_response_id" not in payload + and response_usage is None + ): + warnings.add("compaction_usage_unavailable") + continue + if event.get("type") == "token_usage_record": + if not isinstance(payload.get("thread_id"), str): + warnings.add("token_record_invalid") + continue + if payload.get("thread_id") != session.thread_id: + continue + timestamp = _timestamp(event.get("timestamp")) + usage = _usage_counts(payload.get("usage")) + response_id = payload.get("response_id") + if ( + timestamp is None + or usage is None + or not isinstance(response_id, str) + or not response_id + ): + warnings.add("token_record_invalid") + continue + is_completion_turn = ( + completed_turn_id is not None and payload.get("turn_id") == completed_turn_id + ) + if completed_at is not None and timestamp > completed_at and not is_completion_turn: + continue + if response_usage is None: + response_usage = _empty_token_usage() + if current_turn_id is None: + current_turn_id = payload.get("turn_id") + if response_id in response_ids: + continue + response_ids.add(response_id) + _add_token_usage(observed_response_usage, usage) + turn_usage = _usage_counts(payload.get("turn_token_usage")) + if timestamp < started_at: + if turn_usage is not None: + observed_response_usage = turn_usage + continue + if is_completion_turn: + completed_turn_measured = True + if turn_usage is not None and any( + turn_usage[key] > observed_response_usage[key] for key in TOKEN_FIELDS.values() + ): + warnings.add("token_record_incomplete") + _add_token_usage(response_usage, usage) + continue if event.get("type") != "event_msg" or payload.get("type") != "token_count": continue if "info" in payload and payload["info"] is None: @@ -457,14 +609,22 @@ def _read_rollout_usage( previous = snapshot if timestamp < started_at: continue - if completed_at is not None and timestamp > completed_at: + is_completion_turn = ( + completed_turn_id is not None and current_turn_id == completed_turn_id + ) + if completed_at is not None and timestamp > completed_at and not is_completion_turn: continue + if is_completion_turn: + completed_turn_measured = True if delta["totalTokens"] <= 0: continue - _add_token_usage(total, delta) + _add_token_usage(legacy_usage, delta) + _add_token_usage(total, response_usage if response_usage is not None else legacy_usage) if not boundary_reached: warnings.add("thread_ownership_unavailable") + if completed_turn_id is not None and not completed_turn_measured: + warnings.add("completion_turn_unavailable") return total, warnings @@ -515,7 +675,10 @@ def _token_snapshot(payload: Mapping[str, Any]) -> dict[str, int] | None: info = payload.get("info") if not isinstance(info, dict): return None - usage = info.get("total_token_usage") + return _usage_counts(info.get("total_token_usage")) + + +def _usage_counts(usage: object) -> dict[str, int] | None: if not isinstance(usage, dict): return None cache_write = usage.get("cache_write_input_tokens", usage.get("cache_write_tokens", 0)) diff --git a/plugins/codex-security/scripts/workbench_validation.py b/plugins/codex-security/scripts/workbench_validation.py index 81a561614..6ca8889a7 100644 --- a/plugins/codex-security/scripts/workbench_validation.py +++ b/plugins/codex-security/scripts/workbench_validation.py @@ -92,6 +92,19 @@ def _valid_measured_scan_usage(usage: object) -> bool: or thread_count < 0 ): return False + completion = usage.get("completion") + if "completion" in usage and ( + not isinstance(completion, dict) + or set(completion) != {"threadId", "turnId", "finalized"} + or not isinstance(completion["threadId"], str) + or not completion["threadId"] + or not isinstance(completion["turnId"], str) + or not completion["turnId"] + or type(completion["finalized"]) is not bool + or not completion["finalized"] + and coverage == "complete" + ): + return False warnings = usage.get("warnings", []) if ( not isinstance(warnings, list) @@ -105,7 +118,7 @@ def _valid_measured_scan_usage(usage: object) -> bool: return False if coverage == "unavailable": return thread_count == 0 and set(usage).issubset( - {"coverage", "source", "threadCount", "warnings"} + {"coverage", "source", "threadCount", "warnings", "completion"} ) allowed_keys = { @@ -114,6 +127,7 @@ def _valid_measured_scan_usage(usage: object) -> bool: "threadCount", "missingThreadCount", "warnings", + "completion", *SCAN_USAGE_TOKEN_KEYS, } if thread_count == 0 or not set(usage).issubset(allowed_keys): diff --git a/plugins/codex-security/skills/deep-security-scan/SKILL.md b/plugins/codex-security/skills/deep-security-scan/SKILL.md index e8bcb71a4..d282da8ab 100644 --- a/plugins/codex-security/skills/deep-security-scan/SKILL.md +++ b/plugins/codex-security/skills/deep-security-scan/SKILL.md @@ -75,6 +75,6 @@ After the coordinator returns the canonical manifest and all three unsealed pare Detailed finding write-ups and hardening proposals remain optional, exactly as in an ordinary Standard scan; invoke `$codex-security:vulnerability-writeup` or `$codex-security:propose-security-hardening` only when the corresponding additional output is requested. Read `get_codex_security_completed_scan` only when requested structured or benchmark output actually requires the full sealed documents. -Return user-facing or benchmark output only after completion succeeds and the generated `report.md` exists. Link the report and canonical artifacts. Include measured total, input, and cached input token counts; explicitly label partial token measurement and state when measurement is unavailable instead of reporting zero or estimating. If the configured time limit elapsed before any source review completed, explain that no source review finished and retain any saved findings or pending work; do not claim the repository is free of vulnerabilities. An empty finding set with completed review is the ordinary Codex Security no-findings result, not permission to skip completion. +Return user-facing or benchmark output only after completion succeeds and the generated `report.md` exists. Link the report and canonical artifacts. Include measured total, input, and cached input token counts only when coverage is complete; leave pending usage to the host's end-of-turn finalization and identify other incomplete counts as unavailable. Never report a preliminary count as a final total or an "at least" estimate. If the configured time limit elapsed before any source review completed, explain that no source review finished and retain any saved findings or pending work; do not claim the repository is free of vulnerabilities. An empty finding set with completed review is the ordinary Codex Security no-findings result, not permission to skip completion. If canonical coordinator artifacts are missing or malformed, stop and surface the exact blocker without calling completion, fabricating a report, or claiming success. If completion fails, surface its exact MCP error without retrying, canceling, failing the durable scan, or returning final, no-findings, structured, or benchmark output. Leave resumable work running and preserve its handoff claim. On explicit cancellation, call `cancel_codex_security_scan` and do not accept later scan progress or success. The host may preserve a final in-flight checkpoint after worker cleanup without resuming analysis. Never edit repository files, widen the target, expose internal worker bookkeeping unless requested, or call `fail_codex_security_scan` merely because a waiter detached, a turn ended, workers are still active, or partial artifacts exist. diff --git a/plugins/codex-security/skills/security-scan/SKILL.md b/plugins/codex-security/skills/security-scan/SKILL.md index ebb69e7f6..897042459 100644 --- a/plugins/codex-security/skills/security-scan/SKILL.md +++ b/plugins/codex-security/skills/security-scan/SKILL.md @@ -24,7 +24,7 @@ When an SDK or terminal host sets `CODEX_SECURITY_SCAN_ID`, emit its standalone 1. Resolve the repository, requested scope, and output scan directory from the host-provided scan context when available; otherwise use only the explicitly SDK-provided directory. Preserve the exact user context, supplied threat model, applicable inherited `SECURITY.md` guidance, and optional `CODEX_SECURITY_KNOWLEDGE_BASE` for the core audit. Resolve `` from the configured interpreter (`"$PYTHON"` in POSIX shells or `& "$env:PYTHON"` in PowerShell), otherwise use `python3` on Unix-like hosts or `python` on Windows. Only when `CODEX_SECURITY_TARGET_PATHS_FILE` is supplied, resolve every authorized source path before review with ` /scripts/generate_rank_input.py make-repo-scope-input --repo --scopes-file --out /scoped-source-input.jsonl`; use `"$CODEX_SECURITY_TARGET_PATHS_FILE"` in POSIX shells or `"$env:CODEX_SECURITY_TARGET_PATHS_FILE"` in PowerShell and honor repository ignore rules for directory descendants while retaining every directly requested file. Never print, modify, or treat the scope input as shell syntax; pass it to the core audit without widening the authorized target or scope. 2. Read `../../references/core-scan.md` once and review the resolved target within its authorized scope and configured limits, using the supplied context, threat model, security policy, knowledge base, workers, and scoped-source inventory. Retain the resulting complete semantic `scope`, `threatModel`, `findings`, and `coverage`; preserve every finding's source evidence, calibrated severity, confidence, root cause, validation, attack path, and honest coverage. 3. For a host-backed scan, save a `complete: false` checkpoint as soon as the threat model is available, with no findings when none exist yet. Continue checkpointing during the core audit, then submit one accepted final semantic draft with `record_codex_security_scan_draft({ scanId, complete: true, handoffClaimToken?, scope?, threatModel, findings, coverage })`; let the workbench derive its authoritative target, scope, coverage metadata, surface IDs, finding identities, and fingerprints. On that final draft, close finished generic tasks with `coverage.resolvedDeferred: [{ id, reason }]`, using the saved IDs from `coverage.deferred` and a completion reason. Reuse saved surface IDs for updates; retain unfinished work and candidate outcomes. If the draft is explicitly rejected before writing, correct only the identified fields without dropping valid findings or evidence and retry the same scan at most twice. For an SDK-owned scan with a bound semantic draft tool, use it for the same early model checkpoint and later drafts. For an explicitly SDK-owned scan without that tool, retain the model in an early unsealed canonical draft and update the unsealed canonical `scan-manifest.json`, `findings.json`, and `coverage.json`; use `scoped_path` for both coverage fields when a scope was requested, otherwise set `coverage.mode` to `repository` and `coverage.inventoryStrategy` to `directory` for a non-Git directory or `repository` for a Git-backed target. Omit `scan.sealedAt` and `scan.artifacts`; an SDK scan preserves its exact registered directory and all SDK-provided scan and target values. When `CODEX_SECURITY_TARGET_PATHS_FILE` is supplied on either file-authored path, bind its exact requested paths with `/scripts/launch_codex_security_mcp --helper bind-repo-scopes --scopes-file --manifest /scan-manifest.json --coverage /coverage.json`, using the same shell-specific target-paths reference. For SDK scans, use the exact scope-binding command in the scan instructions so Windows batch invocation preserves literal path values. For other Windows scans, use the [PowerShell scope-binding command](#windows-scope-binding) below. -4. Verify all three canonical JSON files exist. For an SDK-owned scan, return control without finalizing, sealing, generating `report.md` or `threatmodel.md`, or starting another scan; the SDK owns completion. For another host-backed scan, call `complete_codex_security_scan({ scanId, handoffClaimToken? })` once. Outside the SDK path, return only after completion succeeds and the generated `report.md` exists; never write the report by hand or reread the complete canonical findings unless the user explicitly requests them. Report measured token counts when returned and label partial measurement or unavailable usage honestly. +4. Verify all three canonical JSON files exist. For an SDK-owned scan, return control without finalizing, sealing, generating `report.md` or `threatmodel.md`, or starting another scan; the SDK owns completion. For another host-backed scan, call `complete_codex_security_scan({ scanId, handoffClaimToken? })` once. Outside the SDK path, return only after completion succeeds and the generated `report.md` exists; never write the report by hand or reread the complete canonical findings unless the user explicitly requests them. Report complete measured token counts only; leave pending usage to the host's end-of-turn finalization and identify other incomplete counts as unavailable. Keep discovery, validation, and attack-path reasoning within this Standard workflow; do not invoke separate phase skills or load Deep or diff references. Never call Deep-only tools. Do not create ranking phases, per-file or per-candidate ledgers, separate phase worker pools, repeated phase reports, or receipt files. diff --git a/plugins/codex-security/skills/security-scan/references/desktop-scan.md b/plugins/codex-security/skills/security-scan/references/desktop-scan.md index 8af0c8fd9..5f368464b 100644 --- a/plugins/codex-security/skills/security-scan/references/desktop-scan.md +++ b/plugins/codex-security/skills/security-scan/references/desktop-scan.md @@ -22,7 +22,7 @@ Publish meaningful completed-review batches through the same tool. Increase the Save the first `complete: false` checkpoint as soon as architecture mapping produces the threat model, with `findings: []` when no findings exist yet. The host retains the canonical model and generates `threatmodel.md`. Continue saving checkpoints as results arrive and validation decisions are made. Preserve pending candidates with their original evidence in `coverage.deferred`; they are not validated findings. When the audit finishes, record the final semantic scan draft once with `record_codex_security_scan_draft({ scanId, complete: true, handoffClaimToken?, scope?, threatModel?, findings, coverage })`. Supply the actual findings, source-backed coverage, and preserved threat model; let the workbench write the unsealed canonical artifacts and derive authoritative target, scope, coverage metadata, finding identities, and fingerprints. -Honor desktop handoff requirements for any per-finding write-ups before completion. After the semantic draft succeeds and all three canonical JSON files exist, call `complete_codex_security_scan` exactly once with the same authoritative scan ID and handoff token. Return only after completion succeeds and `report.md` exists, linking the generated report and canonical artifacts; retrieve the complete findings only when the user explicitly requests them. Include measured token usage when available and explicitly label partial or unavailable measurement. +Honor desktop handoff requirements for any per-finding write-ups before completion. After the semantic draft succeeds and all three canonical JSON files exist, call `complete_codex_security_scan` exactly once with the same authoritative scan ID and handoff token. Return only after completion succeeds and `report.md` exists, linking the generated report and canonical artifacts; retrieve the complete findings only when the user explicitly requests them. Include complete measured token usage only; the host finalizes pending usage after this turn ends. Identify other incomplete counts as unavailable. For each reported finding, emit one `::code-comment{title="[] " body="<explanation>" file="<absolute path>" start=<line> end=<line> priority=<priority-number> confidence=<0-to-1>}` review directive at its tightest `root_control` location, or the most relevant affected source location when no root control is identifiable. Map `critical`, `high`, `medium`, and `low` to `P0/0`, `P1/1`, `P2/2`, and `P3/3` respectively; keep its title and explanation consistent with the generated report. diff --git a/plugins/codex-security/tests/test_workbench_scan_usage.py b/plugins/codex-security/tests/test_workbench_scan_usage.py index dafd687f3..11126db43 100644 --- a/plugins/codex-security/tests/test_workbench_scan_usage.py +++ b/plugins/codex-security/tests/test_workbench_scan_usage.py @@ -6,9 +6,10 @@ import sys import tempfile import uuid -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta from pathlib import Path +from types import SimpleNamespace from typing import Any import pytest @@ -84,7 +85,7 @@ def _start_scan(tmp_path: Path, *, mode: str = "standard") -> ScanFixture: else: target.mkdir() (target / "app.py").write_text("print('fixture')\n", encoding="utf-8") - workspace = create_saved_workspace(state_dir, target, thread_id="scan-parent") + workspace = create_saved_workspace(state_dir, target, thread_id="scan-parent", mode=mode) workspace_id = str(workspace["id"]) started = start_delivered_scan( @@ -95,6 +96,16 @@ def _start_scan(tmp_path: Path, *, mode: str = "standard") -> ScanFixture: str(tmp_path / "scans"), environment=environment, )["results"] + if mode == "deep": + begin_deep_scan( + state_dir, + "scan-parent", + "--scan-id", + str(started["scanId"]), + "--available-parallelism", + "4", + environment=environment, + ) with sqlite3.connect(state_dir / "workbench.sqlite3") as connection: row = connection.execute( "SELECT started_at FROM scans WHERE id = ?", (started["scanId"],) @@ -156,6 +167,7 @@ def _rollout( include_task_start: bool = True, include_task_start_timestamp: bool = True, task_started_at: datetime | None = None, + session_started_at: datetime | None = None, copied_events: list[dict[str, Any]] | None = None, ) -> Path: directory = root / "rollouts" @@ -171,6 +183,8 @@ def _rollout( }, *(copied_events or []), ] + if session_started_at is not None: + records[0]["payload"]["timestamp"] = session_started_at.isoformat().replace("+00:00", "Z") if parent_thread_id is not None and include_task_start: event: dict[str, Any] = { "type": "event_msg", @@ -191,6 +205,41 @@ def _rollout( return rollout.resolve() +def _response_event( + timestamp: datetime, + thread_id: str, + response_id: str, + input_tokens: int, + output_tokens: int, + *, + turn_usage: tuple[int, int] | None = None, + turn_id: str | None = None, + cached_input_tokens: int = 0, +) -> dict[str, Any]: + usage = _token_event( + timestamp, input_tokens, output_tokens, cached_input_tokens=cached_input_tokens + )["payload"]["info"]["total_token_usage"] + total = ( + _token_event(timestamp, *turn_usage)["payload"]["info"]["total_token_usage"] + if turn_usage is not None + else usage + ) + return _event( + timestamp, + "token_usage_record", + { + "thread_id": thread_id, + "turn_id": turn_id or f"turn-{thread_id}", + "session_id": "session-fixture", + "root_turn_id": turn_id or f"turn-{thread_id}", + "response_id": response_id, + "usage": usage, + "turn_token_usage": total, + "thread_token_usage": total, + }, + ) + + def _state_graph( environment: dict[str, str], threads: dict[str, Path], @@ -199,7 +248,10 @@ def _state_graph( sqlite_home = Path(environment["CODEX_SQLITE_HOME"]) sqlite_home.mkdir(parents=True, exist_ok=True) with sqlite3.connect(sqlite_home / "state_5.sqlite") as connection: - connection.execute("CREATE TABLE threads (id TEXT PRIMARY KEY, rollout_path TEXT NOT NULL)") + connection.execute( + "CREATE TABLE threads " + "(id TEXT PRIMARY KEY, rollout_path TEXT NOT NULL, cwd TEXT, created_at INTEGER)" + ) connection.execute( "CREATE TABLE thread_spawn_edges " "(parent_thread_id TEXT NOT NULL, child_thread_id TEXT NOT NULL)" @@ -231,7 +283,7 @@ def _counts( } -def _complete_scan(fixture: ScanFixture) -> dict[str, Any]: +def _complete_scan(fixture: ScanFixture, *arguments: str) -> dict[str, Any]: options: dict[str, Any] = {"relative_path": "app.py"} if fixture.mode == "diff": assert fixture.diff_target is not None @@ -251,7 +303,11 @@ def _complete_scan(fixture: ScanFixture) -> dict[str, Any]: **options, ) return scan_command( - fixture.state_dir, "complete-scan", fixture.scan_id, environment=fixture.environment + fixture.state_dir, + "complete-scan", + fixture.scan_id, + *arguments, + environment=fixture.environment, ) @@ -281,6 +337,315 @@ def test_completion_preserves_usage_record_warnings( assert usage.get("warnings", []) == (["token_record_invalid"] if coverage == "partial" else []) +def test_completion_counts_native_responses_and_compaction_without_legacy_duplicates( + tmp_path: Path, +) -> None: + fixture = _start_scan(tmp_path) + counted = fixture.started_at + timedelta(microseconds=1) + compaction = _response_event(counted, "scan-parent", "compaction", 40, 3, turn_usage=(60, 5)) + parent = _rollout( + tmp_path, + "scan-parent", + [ + _token_event(fixture.started_at - timedelta(seconds=1), 100, 10), + _token_event(counted, 150, 15), + _event(counted, "event_msg", {"type": "task_started", "turn_id": "turn-scan-parent"}), + _response_event(counted, "scan-parent", "first", 20, 2), + _token_event(counted, 170, 17), + compaction, + _event( + counted, + "compacted", + { + "message": "Summary.", + "replacement_history": [], + "compaction_response_id": "compaction", + }, + ), + _token_event(counted, 170, 17), + compaction, + _response_event(counted, "scan-parent", "last", 30, 4, turn_usage=(90, 9)), + ], + ) + _state_graph(fixture.environment, {"scan-parent": parent}, []) + + assert _complete_scan(fixture)["scan"]["usage"] == { + "coverage": "complete", + "source": "codex_rollout", + **_counts(140, 0, 14), + "threadCount": 1, + } + + +@pytest.mark.parametrize( + ("scope", "native", "expected_warning"), + [ + ("before", False, False), + ("counted", False, True), + ("after", False, False), + ("completion_turn", False, True), + ("counted", True, False), + ("completion_turn", True, False), + ("non_model", False, False), + ], +) +def test_legacy_compaction_usage_is_unverifiable_only_inside_scan_window( + tmp_path: Path, scope: str, native: bool, expected_warning: bool +) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:01+00:00") + completed = started + timedelta(seconds=1) + compacted_at = ( + started - timedelta(seconds=1) + if scope == "before" + else completed + timedelta(seconds=1) + if scope in {"after", "completion_turn"} + else started + ) + marker = _event(compacted_at, "compacted", {"message": "Summary.", "replacement_history": []}) + if scope == "non_model": + marker["payload"]["compaction_response_id"] = None + events = [ + _event(started, "event_msg", {"type": "task_started", "turn_id": "turn-scan-parent"}), + _response_event(started, "scan-parent", "compaction", 10, 1) + if native + else _token_event(started, 10, 1), + ] + events.insert(0 if scope == "before" else len(events), marker) + parent = _rollout(tmp_path, "scan-parent", events) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=started, + completed_at=completed, + completed_turn_id="turn-scan-parent" if scope == "completion_turn" else None, + ) + + assert usage == _counts(10, 0, 1) + assert warnings == ({"compaction_usage_unavailable"} if expected_warning else set()) + + +def test_legacy_compaction_warning_survives_a_later_native_turn(tmp_path: Path) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + parent = _rollout( + tmp_path, + "scan-parent", + [ + _token_event(started, 10, 1), + _event(started, "compacted", {"message": "Summary.", "replacement_history": []}), + _event(started, "event_msg", {"type": "task_started", "turn_id": "native-turn"}), + _response_event(started, "scan-parent", "native", 20, 2, turn_id="native-turn"), + ], + ) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=started, + completed_at=started, + ) + + assert usage == _counts(30, 0, 3) + assert warnings == {"compaction_usage_unavailable"} + + +def test_native_usage_respects_scan_window_and_preserves_cache_details(tmp_path: Path) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + before = started - timedelta(seconds=1) + completed = started + timedelta(seconds=1) + response = _response_event( + started, "scan-parent", "counted", 20, 2, turn_usage=(120, 12), cached_input_tokens=3 + ) + response["payload"]["usage"]["cache_write_input_tokens"] = 4 + response["payload"]["turn_token_usage"].update( + cached_input_tokens=3, cache_write_input_tokens=4 + ) + parent = _rollout( + tmp_path, + "scan-parent", + [ + _response_event(before, "scan-parent", "before", 100, 10), + _token_event(before, 100, 10), + response, + _token_event(started, 120, 12), + _response_event( + completed + timedelta(seconds=1), + "scan-parent", + "after", + 80, + 8, + turn_usage=(200, 20), + ), + ], + ) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=started, + completed_at=completed, + ) + + assert usage == _counts(20, 3, 2, cache_write_input_tokens=4) + assert warnings == set() + + +def test_native_worker_usage_excludes_inherited_totals_and_survives_resume(tmp_path: Path) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + inherited = _response_event(started, "scan-parent", "parent-response", 1_000, 100) + owned = _response_event(started, "scan-worker", "worker-response", 20, 2) + owned["payload"]["thread_token_usage"] = _token_event(started, 1_020, 102)["payload"]["info"][ + "total_token_usage" + ] + worker = _rollout( + tmp_path, + "scan-worker", + [ + owned, + inherited, + _token_event(started, 1_020, 102), + _event(started, "event_msg", {"type": "task_started", "turn_id": "resumed-turn"}), + _response_event( + started, "scan-worker", "resumed-response", 30, 3, turn_id="resumed-turn" + ), + _token_event(started, 1_050, 105), + _event(started, "event_msg", {"type": "task_started", "turn_id": "legacy-turn"}), + _token_event(started, 1_060, 106), + ], + parent_thread_id="scan-parent", + copied_events=[inherited, _token_event(started, 1_000, 100)], + ) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-worker", "scan-parent", worker), + started_at=started, + completed_at=started, + ) + + assert usage == _counts(60, 0, 6) + assert warnings == set() + + +@pytest.mark.parametrize("gap", ["before", "after", "inside"]) +def test_native_usage_only_marks_receipt_gaps_inside_scan_window(tmp_path: Path, gap: str) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + baseline = 120 if gap == "before" else 20 + in_window = baseline + (40 if gap == "inside" else 30) + after = in_window + (20 if gap == "after" else 10) + parent = _rollout( + tmp_path, + "scan-parent", + [ + _response_event( + started - timedelta(seconds=1), + "scan-parent", + "before", + 20, + 0, + turn_usage=(baseline, 0), + ), + _response_event(started, "scan-parent", "counted", 30, 0, turn_usage=(in_window, 0)), + _response_event( + started + timedelta(seconds=1), + "scan-parent", + "after", + 10, + 0, + turn_usage=(after, 0), + ), + ], + ) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=started, + completed_at=started, + ) + + assert usage == _counts(30, 0, 0) + assert warnings == ({"token_record_incomplete"} if gap == "inside" else set()) + + +def test_native_usage_continues_cumulative_baseline_when_recovering_same_turn( + tmp_path: Path, +) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + turn_started = _event( + started, "event_msg", {"type": "task_started", "turn_id": "turn-scan-parent"} + ) + parent = _rollout( + tmp_path, + "scan-parent", + [ + turn_started, + _response_event(started, "scan-parent", "first", 10, 1), + _token_event(started, 10, 1), + turn_started, + _response_event(started, "scan-parent", "recovered", 20, 2, turn_usage=(30, 3)), + _token_event(started, 30, 3), + ], + ) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=started, + completed_at=started, + ) + + assert usage == _counts(30, 0, 3) + assert warnings == set() + + +@pytest.mark.parametrize( + "failure", ["invalid_usage", "invalid_timestamp", "missing_receipt", "truncated"] +) +def test_native_usage_preserves_partial_measurement_warnings(tmp_path: Path, failure: str) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + first = _response_event(started, "scan-parent", "first", 10, 1) + second = _response_event(started, "scan-parent", "second", 20, 2, turn_usage=(30, 3)) + if failure == "invalid_usage": + second["payload"]["usage"] = {} + elif failure == "invalid_timestamp": + del second["timestamp"] + events = [second] if failure == "missing_receipt" else [first, second] + parent = _rollout(tmp_path, "scan-parent", events) + if failure == "truncated": + parent.write_bytes(parent.read_bytes().removesuffix(b"\n")) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + + usage, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=started, + completed_at=started, + ) + + assert usage == (_counts(20, 0, 2) if failure == "missing_receipt" else _counts(10, 0, 1)) + assert warnings == { + "token_record_incomplete" + if failure == "missing_receipt" + else "rollout_record_incomplete" + if failure == "truncated" + else "token_record_invalid" + } + + @pytest.mark.parametrize("mode", ["standard", "diff"]) def test_completion_counts_only_scan_owned_parent_and_descendants( tmp_path: Path, @@ -588,49 +953,15 @@ def test_completion_rejects_non_system_rollout_symlink(tmp_path: Path) -> None: def test_completion_counts_deep_sdk_workers_and_descendants(tmp_path: Path) -> None: - state_dir = tmp_path / "workbench-state" - target = tmp_path / "target" - target.mkdir() - environment = { - "CODEX_HOME": str(tmp_path / "codex-home"), - "CODEX_SQLITE_HOME": str(tmp_path / "codex-sqlite"), - "CODEX_STATE_DB": "", - } - deep = begin_deep_scan( - state_dir, - "scan-parent", - "--target-path", - str(target), - "--scope", - ".", - "--scan-root", - str(tmp_path / "scans"), - "--available-parallelism", - "4", - environment=environment, - )["deepScan"] - scan_id = str(deep["scanId"]) - scan_dir = Path(str(deep["scanDir"])) - with sqlite3.connect(state_dir / "workbench.sqlite3") as connection: - row = connection.execute("SELECT started_at FROM scans WHERE id = ?", (scan_id,)).fetchone() - assert row is not None - fixture = ScanFixture( - state_dir, - target, - scan_id, - scan_dir, - datetime.fromisoformat(row[0].replace("Z", "+00:00")), - environment, - "deep", - ) + fixture = _start_scan(tmp_path, mode="deep") counted = fixture.started_at + timedelta(microseconds=1) - artifact = scan_dir / "artifacts" / "usage-worker" + artifact = fixture.scan_dir / "artifacts" / "usage-worker" artifact.mkdir(parents=True) prompt = artifact / "prompt.md" prompt.write_text("Review the fixture target.\n", encoding="utf-8") upsert_deep_worker( - state_dir, - scan_id, + fixture.state_dir, + fixture.scan_id, str(uuid.uuid4()), "discovery", "running", @@ -638,10 +969,10 @@ def test_completion_counts_deep_sdk_workers_and_descendants(tmp_path: Path) -> N str(artifact), "--sdk-thread-id", "sdk-worker", - environment=environment, + environment=fixture.environment, ) _state_graph( - environment, + fixture.environment, { "scan-parent": _rollout(tmp_path, "scan-parent", [_token_event(counted, 10, 3)]), "sdk-worker": _rollout(tmp_path, "sdk-worker", [_token_event(counted, 20, 5)]), @@ -663,6 +994,66 @@ def test_completion_counts_deep_sdk_workers_and_descendants(tmp_path: Path) -> N } +@pytest.mark.parametrize("owner_first", [True, False]) +def test_usage_counts_registered_worker_when_its_parent_rollout_is_missing( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, owner_first: bool +) -> None: + fixture = _start_scan(tmp_path) + counted = fixture.started_at + timedelta(microseconds=1) + inherited = _token_event(counted, 100, 10) + _state_graph( + fixture.environment, + { + "scan-parent": _rollout(tmp_path, "scan-parent", [_token_event(counted, 10, 3)]), + "missing-parent": tmp_path / "missing.jsonl", + "registered-worker": _rollout( + tmp_path, + "registered-worker", + [_token_event(counted, 120, 15)], + parent_thread_id="missing-parent", + copied_events=[inherited], + ), + "worker-child": _rollout( + tmp_path, + "worker-child", + [_token_event(counted, 7, 2)], + parent_thread_id="registered-worker", + ), + }, + [ + ("scan-parent", "missing-parent"), + ("missing-parent", "registered-worker"), + ("registered-worker", "worker-child"), + ], + ) + for key, value in fixture.environment.items(): + monkeypatch.setenv(key, value) + helper = Path(__file__).resolve().parents[1] / "scripts" / "workbench_scan_usage.py" + collect_usage = runpy.run_path(str(helper))["collect_scan_usage"] + with sqlite3.connect(fixture.state_dir / "workbench.sqlite3") as connection: + connection.row_factory = sqlite3.Row + connection.execute( + "UPDATE scans SET continuation_thread_id = ? WHERE id = ?", + ("registered-worker", fixture.scan_id), + ) + scan = connection.execute("SELECT * FROM scans WHERE id = ?", (fixture.scan_id,)).fetchone() + usage = collect_usage( + connection, + scan, + thread_id="scan-parent" if owner_first else None, + completed_at=counted.isoformat(), + ) + + assert usage == { + "coverage": "partial", + "source": "codex_rollout", + **_counts(37, 0, 10), + "threadCount": 3, + "missingThreadCount": 1, + "warnings": ["rollout_unavailable"], + } + + def test_completion_preserves_explicit_legacy_cost(tmp_path: Path) -> None: fixture = _start_scan(tmp_path) counted = fixture.started_at + timedelta(microseconds=1) @@ -714,6 +1105,287 @@ def test_usage_is_returned_by_completion_without_an_extra_command(tmp_path: Path assert "invalid choice" in str(extra_command["stderr"]) +@pytest.mark.parametrize("native", [False, True]) +def test_host_finalizes_exact_completion_turn_without_changing_artifacts( + tmp_path: Path, native: bool +) -> None: + fixture = _start_scan(tmp_path) + counted = fixture.started_at + timedelta(microseconds=1) + turn_id = "turn-scan-parent" + parent = _rollout( + tmp_path, + "scan-parent", + [ + _event(counted, "event_msg", {"type": "task_started", "turn_id": turn_id}), + _response_event(counted, "scan-parent", "initial", 10, 2) + if native + else _token_event(counted, 10, 2), + ], + ) + worker = _rollout( + tmp_path, + "scan-worker", + [_token_event(counted, 20, 3)], + parent_thread_id="scan-parent", + ) + _state_graph( + fixture.environment, + {"scan-parent": parent, "scan-worker": worker}, + [("scan-parent", "scan-worker")], + ) + initial = _complete_scan(fixture, "--thread-id", "scan-parent", "--turn-id", turn_id)["scan"] + completion = {"threadId": "scan-parent", "turnId": turn_id, "finalized": False} + assert initial["usage"] == { + "coverage": "partial", + "source": "codex_rollout", + **_counts(30, 0, 5), + "threadCount": 2, + "warnings": ["turn_in_progress"], + "completion": completion, + } + sealed_files = { + name: (fixture.scan_dir / name).read_bytes() + for name in ("scan-manifest.json", "findings.json", "coverage.json", "report.md") + } + with sqlite3.connect(fixture.state_dir / "workbench.sqlite3") as connection: + stopped_at = connection.execute( + "SELECT completed_at FROM scans WHERE id = ?", (fixture.scan_id,) + ).fetchone()[0] + late = datetime.fromisoformat(stopped_at.replace("Z", "+00:00")) + timedelta(seconds=1) + final_events = ( + [ + _response_event(late, "scan-parent", "final-response", 30, 4, turn_usage=(40, 6)), + _response_event(late, "scan-parent", "compaction", 50, 5, turn_usage=(90, 11)), + ] + if native + else [_token_event(late, 90, 11)] + ) + with parent.open("a", encoding="utf-8") as output: + output.writelines(f"{json.dumps(event)}\n" for event in final_events) + mismatched = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + "--thread-id", + "scan-parent", + "--turn-id", + "another-turn", + environment=fixture.environment, + check=False, + ) + assert mismatched["returncode"] != 0 + assert ( + scan_command( + fixture.state_dir, "get-scan", fixture.scan_id, environment=fixture.environment + )["scan"]["usage"] + == initial["usage"] + ) + arguments = ("--thread-id", "scan-parent", "--turn-id", turn_id) + finalized = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"] + assert finalized["usage"] == { + "coverage": "complete", + "source": "codex_rollout", + **_counts(110, 0, 14), + "threadCount": 2, + "completion": {**completion, "finalized": True}, + } + later_events = [ + _event(late, "event_msg", {"type": "task_complete", "turn_id": turn_id}), + _event(late, "event_msg", {"type": "task_started", "turn_id": "later-turn"}), + _response_event(late, "scan-parent", "later-response", 900, 90, turn_id="later-turn") + if native + else _token_event(late, 990, 101), + ] + with parent.open("a", encoding="utf-8") as output: + output.writelines(f"{json.dumps(event)}\n" for event in later_events) + helper = runpy.run_path( + str(Path(__file__).resolve().parents[1] / "scripts/workbench_scan_usage.py") + ) + measured, warnings = helper["_read_rollout_usage"]( + helper["RolloutSession"]("scan-parent", None, parent), + started_at=fixture.started_at, + completed_at=datetime.fromisoformat(stopped_at.replace("Z", "+00:00")), + completed_turn_id=turn_id, + ) + assert measured == _counts(90, 0, 11) + assert warnings == set() + assert ( + scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"]["usage"] + == finalized["usage"] + ) + assert {name: (fixture.scan_dir / name).read_bytes() for name in sealed_files} == sealed_files + + +@pytest.mark.parametrize("turn_status", ["completed", "interrupted", "failed"]) +def test_host_final_usage_preserves_missing_worker_warning( + tmp_path: Path, turn_status: str +) -> None: + fixture = _start_scan(tmp_path) + counted = fixture.started_at + timedelta(microseconds=1) + turn_id = "turn-scan-parent" + _state_graph( + fixture.environment, + { + "scan-parent": _rollout( + tmp_path, + "scan-parent", + [_response_event(counted, "scan-parent", "final-response", 10, 2)], + ), + "scan-worker": tmp_path / "missing-worker.jsonl", + }, + [("scan-parent", "scan-worker")], + ) + _complete_scan(fixture, "--thread-id", "scan-parent", "--turn-id", turn_id) + usage = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + "--thread-id", + "scan-parent", + "--turn-id", + turn_id, + "--turn-status", + turn_status, + environment=fixture.environment, + )["scan"]["usage"] + assert usage == { + "coverage": "partial", + "source": "codex_rollout", + **_counts(10, 0, 2), + "threadCount": 1, + "missingThreadCount": 1, + "warnings": ["rollout_unavailable"] + + ([] if turn_status == "completed" else [f"scan_turn_{turn_status}"]), + "completion": {"threadId": "scan-parent", "turnId": turn_id, "finalized": True}, + } + + +def test_host_final_usage_retries_a_locked_native_database_without_losing_counts( + tmp_path: Path, +) -> None: + fixture = _start_scan(tmp_path) + counted = fixture.started_at + timedelta(microseconds=1) + parent = _rollout( + tmp_path, + "scan-parent", + [_response_event(counted, "scan-parent", "first", 10, 2)], + ) + _state_graph(fixture.environment, {"scan-parent": parent}, []) + arguments = ("--thread-id", "scan-parent", "--turn-id", "turn-scan-parent") + initial = _complete_scan(fixture, *arguments)["scan"]["usage"] + with parent.open("a", encoding="utf-8") as output: + final_response = _response_event(counted, "scan-parent", "final", 20, 3, turn_usage=(30, 5)) + output.write(f"{json.dumps(final_response)}\n") + with sqlite3.connect( + Path(fixture.environment["CODEX_SQLITE_HOME"]) / "state_5.sqlite" + ) as native_database: + native_database.execute("BEGIN EXCLUSIVE") + pending = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"]["usage"] + assert pending == initial + assert pending["totalTokens"] == 12 + assert pending["completion"]["finalized"] is False + + finalized = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"]["usage"] + assert finalized == { + "coverage": "complete", + "source": "codex_rollout", + **_counts(30, 0, 5), + "threadCount": 1, + "completion": {**initial["completion"], "finalized": True}, + } + + +def test_host_final_usage_recovers_initially_unavailable_records(tmp_path: Path) -> None: + fixture = _start_scan(tmp_path) + arguments = ("--thread-id", "scan-parent", "--turn-id", "turn-scan-parent") + initial = _complete_scan(fixture, *arguments)["scan"]["usage"] + assert initial["coverage"] == "unavailable" + assert initial["completion"]["finalized"] is False + assert ( + scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"]["usage"] + == initial + ) + _state_graph( + fixture.environment, + { + "scan-parent": _rollout( + tmp_path, + "scan-parent", + [_response_event(fixture.started_at, "scan-parent", "final-response", 10, 2)], + ) + }, + [], + ) + finalized = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"]["usage"] + assert finalized["coverage"] == "complete" + assert finalized["totalTokens"] == 12 + assert finalized["completion"]["finalized"] is True + + +def test_host_final_usage_does_not_invent_an_unrecorded_completion_turn(tmp_path: Path) -> None: + fixture = _start_scan(tmp_path) + arguments = ("--thread-id", "scan-parent", "--turn-id", "unrecorded-turn") + _state_graph( + fixture.environment, + { + "scan-parent": _rollout( + tmp_path, + "scan-parent", + [_response_event(fixture.started_at, "scan-parent", "earlier-response", 10, 2)], + ) + }, + [], + ) + _complete_scan(fixture, *arguments) + finalized = scan_command( + fixture.state_dir, + "finalize-scan-usage", + fixture.scan_id, + *arguments, + environment=fixture.environment, + )["scan"]["usage"] + assert finalized["coverage"] == "partial" + assert finalized["totalTokens"] == 12 + assert finalized["warnings"] == ["completion_turn_unavailable"] + assert finalized["completion"]["finalized"] is True + + def test_failed_scan_preserves_legacy_failure_behavior(tmp_path: Path) -> None: fixture = _start_scan(tmp_path) counted = fixture.started_at + timedelta(microseconds=1) @@ -727,3 +1399,198 @@ def test_failed_scan_preserves_legacy_failure_behavior(tmp_path: Path) -> None: )["scan"] assert failed["progress"]["status"] == "failed" assert "usage" not in failed + + +@pytest.mark.parametrize("worker_kind", ["discovery", "dedup"]) +def test_completion_recovers_replaced_deep_worker_threads(tmp_path: Path, worker_kind: str) -> None: + fixture = _start_scan(tmp_path, mode="deep") + fixture = replace( + fixture, + started_at=(fixture.started_at - timedelta(seconds=1)).replace(microsecond=500123), + ) + with sqlite3.connect(fixture.state_dir / "workbench.sqlite3") as connection: + connection.execute( + "UPDATE scans SET started_at = ? WHERE id = ?", + (fixture.started_at.isoformat(), fixture.scan_id), + ) + counted = fixture.started_at + timedelta(microseconds=1) + before_scan = counted - timedelta(days=1) + boundary_predecessor = fixture.started_at.replace(microsecond=400000) + after_scan = counted + timedelta(days=1) + artifact = fixture.scan_dir / "artifacts" / "usage-worker" + artifact.mkdir(parents=True) + prompt = artifact / "prompt.md" + prompt.write_text("Review the fixture target.\n", encoding="utf-8") + worker_id = str(uuid.uuid4()) + if worker_kind == "dedup": + with sqlite3.connect(fixture.state_dir / "workbench.sqlite3") as connection: + connection.execute( + "INSERT INTO deep_scan_workers " + "(id, scan_id, kind, status, prompt_path, artifact_dir, created_at, updated_at) " + "VALUES (?, ?, 'dedup', 'queued', ?, ?, ?, ?)", + ( + worker_id, + fixture.scan_id, + str(prompt), + str(artifact), + fixture.started_at.isoformat(), + fixture.started_at.isoformat(), + ), + ) + for attempt, thread_id in enumerate( + ["sdk-first-attempt", "sdk-current-attempt", "sdk-current-attempt"], start=1 + ): + upsert_deep_worker( + fixture.state_dir, + fixture.scan_id, + worker_id, + worker_kind, + "running", + str(prompt), + str(artifact), + "--attempt", + str(attempt), + "--sdk-thread-id", + thread_id, + environment=fixture.environment, + ) + _state_graph( + fixture.environment, + { + "scan-parent": _rollout(tmp_path, "scan-parent", [_token_event(counted, 10, 3)]), + "sdk-first-attempt": _rollout( + tmp_path, + "sdk-first-attempt", + [_token_event(counted, 20, 5)], + session_started_at=fixture.started_at.replace(microsecond=500000), + ), + "sdk-first-child": _rollout( + tmp_path, + "sdk-first-child", + [_token_event(counted, 7, 2)], + parent_thread_id="sdk-first-attempt", + ), + "sdk-current-attempt": _rollout( + tmp_path, + "sdk-current-attempt", + [_token_event(counted, 20, 5), _token_event(counted, 30, 8)], + session_started_at=before_scan, + ), + "other-scan-worker": _rollout( + tmp_path, "other-scan-worker", [_token_event(counted, 80_000, 20_000)] + ), + "archived-worker": _rollout( + tmp_path, "archived-worker", [_token_event(before_scan, 100, 10)] + ), + "archived-missing-worker": tmp_path / "archived-missing.jsonl", + "boundary-archived-worker": _rollout( + tmp_path, + "boundary-archived-worker", + [_token_event(boundary_predecessor, 100, 10)], + session_started_at=boundary_predecessor, + ), + "boundary-archived-child": tmp_path / "boundary-archived-child.jsonl", + "later-worker": _rollout(tmp_path, "later-worker", [_token_event(after_scan, 100, 10)]), + "later-missing-worker": tmp_path / "later-missing.jsonl", + }, + [ + ("sdk-first-attempt", "sdk-first-child"), + ("boundary-archived-worker", "boundary-archived-child"), + ], + ) + working_directory = artifact if worker_kind == "discovery" else fixture.scan_dir / "artifacts" + with sqlite3.connect( + Path(fixture.environment["CODEX_SQLITE_HOME"]) / "state_5.sqlite" + ) as connection: + connection.executemany( + "UPDATE threads SET cwd = ?, created_at = ? WHERE id = ?", + [ + (str(working_directory), int(counted.timestamp()), "sdk-first-attempt"), + (str(working_directory), int(counted.timestamp()), "sdk-first-child"), + (str(working_directory), int(counted.timestamp()), "sdk-current-attempt"), + ( + str(tmp_path / "other-scan" / "artifacts"), + int(counted.timestamp()), + "other-scan-worker", + ), + (str(working_directory), int(before_scan.timestamp()), "archived-worker"), + (str(working_directory), int(before_scan.timestamp()), "archived-missing-worker"), + ( + str(working_directory), + int(boundary_predecessor.timestamp()), + "boundary-archived-worker", + ), + (str(working_directory), int(after_scan.timestamp()), "later-worker"), + (str(working_directory), int(after_scan.timestamp()), "later-missing-worker"), + ], + ) + usage = _complete_scan(fixture)["scan"]["usage"] + assert usage == { + "coverage": "complete", + "source": "codex_rollout", + **_counts(67, 0, 18), + "threadCount": 4, + } + + +@pytest.mark.cross_platform +@pytest.mark.parametrize( + ("platform_name", "directory", "recorded_cwd", "matches"), + [ + ("nt", r"C:\Scans\Job\output", "\\\\?\\c:\\scans\\job\\OUTPUT\\", True), + ("nt", r"C:\Scans\Job\output", "C:/scans/job/output/", True), + ( + "nt", + r"\\server\share\Scans\Job\output", + "\\\\?\\UNC\\SERVER\\share\\scans\\job\\output\\", + True, + ), + ("nt", r"C:\Scans\Job\output", r"\\?\C:\Scans\Job\other-output", False), + ("posix", "/scans/Job/output", "/scans/Job/output", True), + ("posix", "/scans/Job/output", "/scans/job/output", False), + ], +) +def test_worker_discovery_matches_platform_path_spelling( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + platform_name: str, + directory: str, + recorded_cwd: str, + matches: bool, +) -> None: + started = datetime.fromisoformat("2026-01-01T00:00:00+00:00") + environment = {"CODEX_SQLITE_HOME": str(tmp_path / "codex-sqlite")} + _state_graph( + environment, + { + "current-worker": _rollout(tmp_path, "current-worker", []), + "replaced-worker": _rollout(tmp_path, "replaced-worker", []), + }, + [], + ) + state_database = Path(environment["CODEX_SQLITE_HOME"]) / "state_5.sqlite" + with sqlite3.connect(state_database) as connection: + connection.execute( + "UPDATE threads SET cwd = ?, created_at = ? WHERE id = ?", + (recorded_cwd, int(started.timestamp()), "replaced-worker"), + ) + scripts = Path(__file__).resolve().parents[1] / "scripts" + monkeypatch.syspath_prepend(str(scripts)) + helper = runpy.run_path(str(scripts / "workbench_scan_usage.py")) + discover = helper["_discover_rollout_sessions"] + monkeypatch.setitem(discover.__globals__, "os", SimpleNamespace(name=platform_name)) + warnings: set[str] = set() + + sessions, missing = discover( + state_database, + ["current-worker"], + warnings, + worker_directories={directory}, + started_at=started, + completed_at=started + timedelta(seconds=1), + ) + + expected_threads = {"current-worker", "replaced-worker"} if matches else {"current-worker"} + assert {session.thread_id for session in sessions} == expected_threads + assert missing == set() + assert warnings == set() diff --git a/sdk/typescript/src/api.ts b/sdk/typescript/src/api.ts index 8a4061a94..8c4cbab58 100644 --- a/sdk/typescript/src/api.ts +++ b/sdk/typescript/src/api.ts @@ -1542,6 +1542,7 @@ export class CodexSecurity { let artifactRestorationFailure: OutputDirectoryError | null = null; let customValidationComplete = false; let completionCost: ScanCost | null = null; + let scanTurnStatus: "in_progress" | "completed" | "failed" | null = null; let budgetRecovery: { expectation: ScanExpectation; pluginRoot: string; @@ -1561,6 +1562,48 @@ export class CodexSecurity { this.#dependencies.prepareScanArtifactRestorer ?? prepareScanArtifactRestorer; const workbench = this.#dependencies.runWorkbench ?? runWorkbench; + const terminalScanTurnStatus = () => + scanTurnStatus === "completed" || scanTurnStatus === "failed" + ? scanTurnStatus + : signal.aborted + ? ("interrupted" as const) + : null; + const finalizeScanUsage = async ( + scanId: string, + workbenchOptions: WorkbenchCommandOptions, + completedScan: unknown, + turnStatus: "completed" | "failed" | "interrupted", + ): Promise<void> => { + const measuredUsage = isRecord(completedScan) + ? completedScan["usage"] + : undefined; + const usageCompletion = isRecord(measuredUsage) + ? measuredUsage["completion"] + : undefined; + if ( + isRecord(usageCompletion) && + usageCompletion["finalized"] === false && + typeof usageCompletion["threadId"] === "string" && + typeof usageCompletion["turnId"] === "string" + ) { + try { + await workbench( + { ...workbenchOptions, signal: undefined }, + workbenchArgs("finalize-scan-usage", { + scanId, + threadId: usageCompletion["threadId"], + turnId: usageCompletion["turnId"], + turnStatus, + }), + ); + } catch (error) { + notifyObserver( + options, + "onWarning", + )(`Could not finalize scan token usage: ${errorMessage(error)}`); + } + } + }; try { const checkOpen = (): void => { this.#requireOpen(); @@ -2358,6 +2401,7 @@ export class CodexSecurity { signal, }) : (await thread.runStreamed(prompt, turnOptions)).events; + scanTurnStatus = "in_progress"; checkOpen(); let result = await runScanEvents({ @@ -2374,6 +2418,9 @@ export class CodexSecurity { modelProvider, workbenchValidated: true, model, + onTurnStatus: (status) => { + scanTurnStatus = status; + }, onThreadStarted: async (threadId) => { if (resumeThreadId !== undefined && threadId !== resumeThreadId) { throw new CodexSecurityError( @@ -2545,6 +2592,13 @@ export class CodexSecurity { : ["--cost-json", JSON.stringify(completionCost)]), ]); activeScan = null; + const completedScan = completion["scan"]; + await finalizeScanUsage( + scanId, + workbenchOptions, + completedScan, + "completed", + ); if (directDeepScan) { result = new ScanResult({ ...result, @@ -2556,7 +2610,6 @@ export class CodexSecurity { }, }); } - const completedScan = completion["scan"]; if (isRecord(completedScan) && Array.isArray(completedScan["warnings"])) { const targetWarnings = new Set([ ...preparedTargetWarnings, @@ -2730,6 +2783,7 @@ export class CodexSecurity { } catch (error) { const snapshot = await costTracker?.stop().catch(() => null); if (artifactRestorationFailure !== null) throw artifactRestorationFailure; + const terminalTurnStatus = terminalScanTurnStatus(); let failure = signal.reason instanceof ScanCostLimitExceededError ? signal.reason @@ -2765,6 +2819,13 @@ export class CodexSecurity { `--message=${failure.message.slice(0, 2400)}`, ], ); + if (terminalTurnStatus !== null) + await finalizeScanUsage( + activeScan.id, + activeScan.options, + completion["scan"], + terminalTurnStatus, + ); activeScan = null; runPostScan = null; const result = await collectResult( @@ -2812,6 +2873,37 @@ export class CodexSecurity { return result; } catch {} } + if (activeScan !== null && scanTurnStatus !== null) { + try { + const cleanupOptions = { ...activeScan.options, signal: undefined }; + const saved = await workbench( + cleanupOptions, + workbenchArgs("get-scan", { scanId: activeScan.id }), + ); + const completedScan = saved["scan"]; + const progress = isRecord(completedScan) + ? completedScan["progress"] + : undefined; + if (isRecord(progress) && progress["status"] === "complete") { + const scanId = activeScan.id; + activeScan = null; + if (terminalTurnStatus !== null) + await finalizeScanUsage( + scanId, + cleanupOptions, + completedScan, + terminalTurnStatus, + ); + } + } catch (cleanupError) { + notifyObserver( + options, + "onWarning", + )( + `Could not finalize scan token usage: ${errorMessage(cleanupError)}`, + ); + } + } // A failed attachment must not turn a resumable coordinator into a terminal failure. // Deep Scan orchestration persists its own terminal failures and cancellations. if (activeScan !== null && options.resumeScanId === undefined) { @@ -4100,6 +4192,7 @@ interface ScanEventRunOptions extends Pick<ScanOptions, "onReconnect"> { workbenchValidated?: boolean; model?: string; onFinalize?: (usage: unknown) => Promise<unknown>; + onTurnStatus?: (status: "completed" | "failed") => void; onThreadStarted?: (threadId: string) => Promise<void> | void; onScanStarted?: () => void; onTrustedAccessStatus?: (status: ScanTrustedAccessStatus) => void; @@ -4122,6 +4215,9 @@ export async function runScanEvents( thread: options.thread, events: options.events, onEvent: async (event) => { + if (event.type === "turn.completed") + options.onTurnStatus?.("completed"); + if (event.type === "turn.failed") options.onTurnStatus?.("failed"); if ( !tacStatusReported && options.modelProvider !== "amazon-bedrock" && diff --git a/sdk/typescript/src/custom-validation-prompt.ts b/sdk/typescript/src/custom-validation-prompt.ts index 8e1f1feed..f12ae3e6f 100644 --- a/sdk/typescript/src/custom-validation-prompt.ts +++ b/sdk/typescript/src/custom-validation-prompt.ts @@ -11,7 +11,7 @@ const SOURCES = { "references/core-scan.md": "caa48f845e7edf4768140a7244d92637175df6f8e26a7660109a20f0dbe21546", "skills/security-scan/SKILL.md": - "6410404dc945a3c3bb4447f53c72a6fd920b042cff829cb7bd95042a8cc7de10", + "93a23fd5e8fd505282fca51432dd981ffee85dd99e438d31fcba6ba8c1860f15", "skills/security-diff-scan/SKILL.md": "e0ddfea9e46ec5389a6634630aa518273b0b4bc68efac09441adf549bab8d253", } as const; diff --git a/sdk/typescript/tests-ts/api-usage-finalization.test.ts b/sdk/typescript/tests-ts/api-usage-finalization.test.ts new file mode 100644 index 000000000..685cbfa81 --- /dev/null +++ b/sdk/typescript/tests-ts/api-usage-finalization.test.ts @@ -0,0 +1,309 @@ +import { hash } from "node:crypto"; +import { mkdir, readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import type { ThreadEvent } from "@openai/codex-sdk"; +import { afterEach, expect, test } from "bun:test"; +import type { JsonObject } from "../src/config.js"; +import { copyCompletedScan } from "./plugin-root.js"; +import { mockWorkbench, TestClient } from "./support/api-client.js"; +import { + codexFactory, + completedEvents, + scanRuntimeDependencies, +} from "./support/api-events.js"; +import { createApiTestFixtures } from "./support/temporary-directories.js"; +import { writeSession } from "./support/usage-rollout.js"; + +const { cleanup, temporaryDirectory } = createApiTestFixtures(); +afterEach(cleanup); + +test.each(["pending", "finalized", "unavailable", "standard"] as const)( + "finalizes saved completion-turn usage when accounting is %s", + async (scenario) => { + const root = await temporaryDirectory(); + const repository = join(root, "repository"); + const codexHome = join(root, "codex-home"); + const scanDir = join(root, "scan"); + await Promise.all( + [repository, codexHome, scanDir].map((path) => mkdir(path)), + ); + const commands: Array<readonly string[]> = []; + const warnings: string[] = []; + const observerErrors: string[] = []; + const accountingFailure = "Synthetic usage storage failure: token=fixture"; + const completionWarning = "Synthetic saved scan warning"; + let turnFinished = false; + const client = TestClient.withDependencies({ + ...scanRuntimeDependencies(codexHome, scanDir), + canRunDirectDeepScan: () => false, + runWorkbench: async (_options, args, input) => { + commands.push(args); + if (args[0] === "complete-scan") { + return { + scan: { + warnings: [completionWarning], + ...(scenario === "standard" + ? {} + : { + usage: { + completion: { + threadId: "completion-thread", + turnId: "completion-turn", + finalized: scenario === "finalized", + }, + }, + }), + }, + }; + } + if (args[0] === "finalize-scan-usage") { + expect(turnFinished).toBe(true); + if (scenario === "unavailable") throw new Error(accountingFailure); + return {}; + } + return mockWorkbench(args, input); + }, + createCodex: codexFactory(async () => { + await copyCompletedScan(root); + if (scenario !== "standard") { + const coveragePath = join(scanDir, "coverage.json"); + const coverage = JSON.parse(await readFile(coveragePath, "utf8")); + coverage.mode = "deep_repository"; + const coverageBytes = `${JSON.stringify(coverage)}\n`; + await writeFile(coveragePath, coverageBytes); + const manifestPath = join(scanDir, "scan-manifest.json"); + const manifest = JSON.parse(await readFile(manifestPath, "utf8")); + manifest.scan.artifacts.find( + (artifact: { path: string }) => artifact.path === "coverage.json", + ).sha256 = hash("sha256", coverageBytes); + await writeFile(manifestPath, `${JSON.stringify(manifest)}\n`); + } + async function* events() { + yield* completedEvents(); + turnFinished = true; + } + return { events: events() }; + }), + }); + try { + const result = await client.run(repository, { + mode: scenario === "standard" ? "standard" : "deep", + onWarning: (warning) => { + warnings.push(warning); + if (warning.includes(accountingFailure)) + throw new Error("Synthetic warning observer failure"); + }, + onObserverError: (observer) => observerErrors.push(observer), + }); + expect(result).toMatchObject({ scanDir, threadId: "thread-1" }); + const finalization = commands.filter( + (args) => args[0] === "finalize-scan-usage", + ); + if (scenario === "pending" || scenario === "unavailable") { + expect(finalization).toEqual([ + [ + "finalize-scan-usage", + "--scan-id", + "scan_example_001", + "--thread-id", + "completion-thread", + "--turn-id", + "completion-turn", + "--turn-status", + "completed", + ], + ]); + expect(commands.indexOf(finalization[0]!)).toBeGreaterThan( + commands.findIndex((args) => args[0] === "complete-scan"), + ); + } else { + expect(finalization).toEqual([]); + } + expect(warnings).toEqual([ + ...(scenario === "unavailable" + ? [`Could not finalize scan token usage: ${accountingFailure}`] + : []), + completionWarning, + ]); + expect(observerErrors).toEqual( + scenario === "unavailable" ? ["onWarning"] : [], + ); + expect(commands.some((args) => args[0] === "fail-scan")).toBe(false); + } finally { + await client.close(); + } + }, +); + +test.each([ + "failed", + "interrupted", + "post-turn-error", + "post-turn-cancel", + "transport-error", + "resumed", + "accounting-error", +] as const)( + "preserves a sealed scan and its %s outcome while finalizing usage", + async (scenario) => { + const root = await temporaryDirectory(); + const repository = join(root, "repository"); + const codexHome = join(root, "codex-home"); + const scanDir = join(root, "scan"); + await Promise.all( + [repository, codexHome, scanDir].map((path) => + mkdir(path, { mode: 0o700 }), + ), + ); + const report = "# Sealed scan report\n"; + await writeFile(join(scanDir, "report.md"), report); + if (scenario === "resumed") + await writeSession(codexHome, "thread-1", {}, { cwd: scanDir }); + const controller = new AbortController(); + const commands: Array<readonly string[]> = []; + const warnings: string[] = []; + const originalFailure = new Error("Synthetic scan turn failure"); + const accountingFailure = "Synthetic accounting failure: token=fixture"; + let turnExited = false; + async function* events(): AsyncGenerator<ThreadEvent> { + try { + if (scenario === "post-turn-error" || scenario === "post-turn-cancel") { + yield* completedEvents(); + if (scenario === "post-turn-cancel") + controller.abort(originalFailure); + return; + } + yield { type: "thread.started", thread_id: "thread-1" }; + yield { type: "turn.started" }; + if (scenario === "interrupted") { + controller.abort(originalFailure); + throw originalFailure; + } + if (scenario === "transport-error") throw originalFailure; + yield { + type: "turn.failed", + error: { message: originalFailure.message }, + }; + } finally { + turnExited = true; + } + } + const client = TestClient.withDependencies({ + ...scanRuntimeDependencies(codexHome, scanDir), + canRunDirectDeepScan: () => false, + createCodex: () => ({ + startThread: () => ({ + id: null, + runStreamed: async () => ({ events: events() }), + }), + resumeThread: (threadId) => ({ + id: threadId, + runStreamed: async () => ({ events: events() }), + }), + }), + runWorkbench: async (options, args, input): Promise<JsonObject> => { + commands.push(args); + if (args[0] === "get-cli-scan-resume") { + return { + scanId: "scan_example_001", + scanDir, + targetId: "target_sha256_example", + targetRevision: "deadbeef", + threadId: "thread-1", + recipe: { repository, target: { kind: "repository", paths: [] } }, + contract: { target: { allowedKinds: ["git_revision"] } }, + }; + } + if (args[0] === "prepare-scan-completion") throw originalFailure; + if (args[0] === "get-scan") { + expect(turnExited).toBe(true); + if (scenario !== "post-turn-error") + expect(options.signal).toBeUndefined(); + return { + scan: { + progress: { status: "complete" }, + usage: { + completion: { + threadId: "completion-thread", + turnId: "completion-turn", + finalized: false, + }, + }, + }, + }; + } + if (args[0] === "finalize-scan-usage") { + expect(turnExited).toBe(true); + expect(options.signal).toBeUndefined(); + if (scenario === "accounting-error") + throw new Error(accountingFailure); + } + return mockWorkbench(args, input); + }, + }); + try { + const failure = await client + .run(repository, { + mode: "deep", + ...(scenario === "resumed" + ? { resumeScanId: "scan_example_001", outputDir: scanDir } + : {}), + signal: controller.signal, + onWarning: (warning) => { + warnings.push(warning); + throw new Error("Synthetic warning observer failure"); + }, + }) + .then( + () => { + throw new Error("Expected the original scan failure"); + }, + (error: unknown) => error, + ); + if (scenario === "interrupted" || scenario === "post-turn-cancel") { + expect(failure).toMatchObject({ name: "ScanInterruptedError" }); + } else if ( + scenario === "post-turn-error" || + scenario === "transport-error" + ) { + expect(failure).toBe(originalFailure); + } else { + expect(failure).toMatchObject({ message: originalFailure.message }); + } + const finalization = commands.filter( + (args) => args[0] === "finalize-scan-usage", + ); + expect(finalization).toEqual( + scenario === "transport-error" + ? [] + : [ + [ + "finalize-scan-usage", + "--scan-id", + "scan_example_001", + "--thread-id", + "completion-thread", + "--turn-id", + "completion-turn", + "--turn-status", + scenario === "post-turn-error" || + scenario === "post-turn-cancel" + ? "completed" + : scenario === "interrupted" + ? "interrupted" + : "failed", + ], + ], + ); + expect(commands.some((args) => args[0] === "fail-scan")).toBe(false); + expect(warnings).toEqual( + scenario === "accounting-error" + ? [`Could not finalize scan token usage: ${accountingFailure}`] + : [], + ); + expect(await readFile(join(scanDir, "report.md"), "utf8")).toBe(report); + } finally { + await client.close(); + } + }, +); diff --git a/sdk/typescript/tests-ts/api.test.ts b/sdk/typescript/tests-ts/api.test.ts index 8bfb49c16..2ba1cdc92 100644 --- a/sdk/typescript/tests-ts/api.test.ts +++ b/sdk/typescript/tests-ts/api.test.ts @@ -3783,6 +3783,7 @@ describe("CodexSecurity orchestration", () => { "get-scan-feedback", "set-scan-thread", "prepare-scan-completion", + "get-scan", "fail-scan", ]); await client.close(); @@ -5484,7 +5485,8 @@ describe("CodexSecurity orchestration", () => { "--thread-id", "scan-thread", ]); - expect(commands[3]).toEqual([ + expect(commands[3]).toEqual(["get-scan", "--scan-id", "scan_example_001"]); + expect(commands[4]).toEqual([ "fail-scan", "--scan-id", "scan_example_001", @@ -5557,6 +5559,17 @@ describe("CodexSecurity orchestration", () => { await writeFile(manifestPath, `${JSON.stringify(manifest)}\n`); return { scan: { + ...(completion === "completed" + ? { + usage: { + completion: { + threadId: "scan-thread", + turnId: "completion-turn", + finalized: false, + }, + }, + } + : {}), warnings: [ args .find((value) => value.startsWith("--message="))! @@ -5597,6 +5610,19 @@ describe("CodexSecurity orchestration", () => { `Scan stopped: short-context budget baseline $0.00488 exceeded the $0.004 limit; estimated cost $0.00488–$0.01156 (standard, context unknown, cache writes unknown); partial output remains at ${scanDir}.`, ]); expect(commands.some((args) => args[0] === "fail-scan")).toBe(false); + expect( + commands.find((args) => args[0] === "finalize-scan-usage"), + ).toEqual([ + "finalize-scan-usage", + "--scan-id", + "scan_example_001", + "--thread-id", + "scan-thread", + "--turn-id", + "completion-turn", + "--turn-status", + "interrupted", + ]); } expect(runStreamed).toHaveBeenCalledTimes(1); const recovery = commands.find(