refactor(plugin): simplify workbench persistence

plugins/codex-security/scripts/deep_scan_workbench.py

@@ -17,16 +17,10 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 from deep_scan_config import resolve_deep_scan_config
-from filesystem_identity import serialize_filesystem_identity
 from finalize_scan_contract import _read_scan_local_json
 from workbench.handoff import require_current_continuation
 from workbench.storage import create_private_directory
-from workbench_target import (
-    directory_content_digest,
-    directory_snapshot_regular_file_count,
-    git_revision,
-    worktree_content_digest,
-)
+from workbench_target import directory_snapshot_regular_file_count
 from workbench_validation import optional_text, require_uuid, user_context_argument
 
 DEEP_SCAN_WORKER_KINDS = ("setup", "discovery", "dedup")
@@ -671,18 +665,14 @@ def begin_deep_scan_for_scan(
     if workflow_version is None:
         raise SystemExit("workflow-version is required.")
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         scan, _ = require_owned_scan(connection, scan_id, thread_id)
         require_current_continuation(
             scan,
             args.claim_token,
             error_message="Deep Scan orchestration is owned by another continuation.",
         )
         ensure_deep_scan_run(connection, scan, config, workflow_version, dependencies().now())
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return deep_scan_result(connection, scan_id, start_disposition="created")
 
 
@@ -697,19 +687,13 @@ def begin_deep_scan_for_target(
     if existing is not None:
         return begin_deep_scan_for_scan(connection, existing["id"], thread_id, args)
     target_metadata = target.stat()
-    revision = git_revision(target)
-    target_snapshot_digest = (
-        directory_content_digest(target)
-        if revision == "unversioned"
-        else worktree_content_digest(target)
-    )
-    target_device = serialize_filesystem_identity(target_metadata.st_dev)
-    target_inode = serialize_filesystem_identity(target_metadata.st_ino)
+    target_identity = dependencies().scan_target_identity(target, None, metadata=target_metadata)
+    revision, target_snapshot_digest, target_device, target_inode = target_identity
     scope_file_count = directory_snapshot_regular_file_count(
         target if scope == "." else target / scope
     )
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         existing = existing_deep_scan_for_target(connection, thread_id, target_path, scope)
         if existing is not None:
             existing_run = connection.execute(
@@ -759,11 +743,7 @@ def begin_deep_scan_for_target(
         workflow_version = optional_text(args.workflow_version, maximum=256)
         if workflow_version is None:
             raise SystemExit("workflow-version is required.")
-        root = (
-            Path(args.scan_root).expanduser().resolve()
-            if args.scan_root
-            else dependencies().state_dir() / "scans"
-        )
+        root = dependencies().resolve_scan_root(args.scan_root)
         target_root = (root / dependencies().safe_segment(target.name)).resolve()
         if target_root == target or target in target_root.parents:
             raise SystemExit("The scan artifact directory must be outside the selected target.")
@@ -800,55 +780,28 @@ def begin_deep_scan_for_target(
                 timestamp,
             ),
         )
-        connection.execute(
-            """
-            INSERT INTO scans (
-                id, workspace_id, target_id, target_path, target_revision, target_snapshot_digest,
-                target_device, target_inode, scope, mode, user_context,
-                deep_scan_owner_thread_id, scan_dir, model, reasoning_effort, status, phase,
-                handoff_status, started_at, created_at, updated_at
-            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'deep', ?, ?, ?, ?, ?,
-                'running', 'preflight', 'delivered', ?, ?, ?)
-            """,
-            (
-                scan_id,
-                workspace_id,
-                target_id,
-                target_path,
-                revision,
-                target_snapshot_digest,
-                target_device,
-                target_inode,
-                scope,
-                user_context,
-                thread_id,
-                str(scan_dir),
-                model,
-                reasoning_effort,
-                timestamp,
-                timestamp,
-                timestamp,
-            ),
-        )
-        connection.execute(
-            """
-            INSERT INTO scan_progress (
-                scan_id, scope_file_count, review_items_total, review_items_completed,
-                reportable_findings_count, updated_at
-            ) VALUES (?, ?, 0, 0, 0, ?)
-            """,
-            (scan_id, scope_file_count, timestamp),
-        )
-        connection.execute(
-            "UPDATE workspaces SET active_scan_id = ?, updated_at = ? WHERE id = ?",
-            (scan_id, timestamp, workspace_id),
+        workspace = connection.execute(
+            "SELECT * FROM workspaces WHERE id = ?", (workspace_id,)
+        ).fetchone()
+        dependencies().insert_running_scan(
+            connection,
+            scan_id=scan_id,
+            workspace=workspace,
+            target=target,
+            scope=scope,
+            diff_target=None,
+            target_identity=target_identity,
+            target_root=target_root,
+            target_summary=None,
+            scope_file_count=scope_file_count,
+            timestamp=timestamp,
+            handoff_status="delivered",
+            model=model,
+            reasoning_effort=reasoning_effort,
+            scan_dir=scan_dir,
         )
         scan = dependencies().require_scan(connection, scan_id)
         ensure_deep_scan_run(connection, scan, config, workflow_version, timestamp)
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return deep_scan_result(connection, scan_id, start_disposition="created")
 
 
@@ -931,7 +884,7 @@ def claim_deep_scan_coordinator_locked(
     connection: sqlite3.Connection, args: argparse.Namespace, scan_id: str
 ) -> dict[str, Any]:
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         scan, _ = require_owned_scan(connection, scan_id, args.thread_id)
         require_current_continuation(
             scan,
@@ -963,10 +916,6 @@ def claim_deep_scan_coordinator_locked(
             """,
             (int(args.coordinator_generation != run["coordinator_generation"]), timestamp, scan_id),
         )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return {
         **deep_scan_result(connection, scan_id),
         "coordinatorDisposition": disposition,
@@ -1128,7 +1077,7 @@ def upsert_deep_scan_worker(
     scan_id = require_uuid(args.scan_id, "scan-id")
     worker_id = require_uuid(args.worker_id, "worker-id")
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         run = require_deep_scan_run(connection, scan_id)
         require_current_coordinator(run, args)
         scan = dependencies().require_scan(connection, scan_id)
@@ -1349,10 +1298,6 @@ def upsert_deep_scan_worker(
                 worker_id,
             ),
         )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return deep_scan_result(connection, scan_id)
 
 
@@ -1365,7 +1310,7 @@ def claim_deep_scan_dedup(
     if len(set(input_ids)) != len(input_ids):
         raise SystemExit("Dedup input worker IDs must be unique.")
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         run, scan = require_running_deep_scan(connection, scan_id)
         require_current_coordinator(run, args)
         prompt_path = deep_scan_path(scan, args.prompt_path, "Dedup prompt path", kind="file")
@@ -1494,10 +1439,6 @@ def claim_deep_scan_dedup(
             """,
             (timestamp, scan_id),
         )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return deep_scan_result(connection, scan_id)
 
 
@@ -2011,7 +1952,7 @@ def record_deep_scan_publication_failure(
         raise SystemExit("message is required.")
     with dependencies().scan_completion_lock(scan_id):
         connection.execute("BEGIN IMMEDIATE")
-        try:
+        with connection:
             run = require_deep_scan_run(connection, scan_id)
             require_current_coordinator(run, args)
             scan = dependencies().require_scan(connection, scan_id)
@@ -2035,10 +1976,6 @@ def record_deep_scan_publication_failure(
                     """,
                     (message, timestamp, scan_id),
                 )
-            connection.commit()
-        except BaseException:
-            connection.rollback()
-            raise
     return deep_scan_result(connection, scan_id)
 
 

plugins/codex-security/scripts/workbench_cli.py

@@ -9,7 +9,6 @@
 # Some plugin hosts launch Python with safe-path isolation enabled.
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 import deep_scan_workbench as deep_scan
-import workbench_remediation as remediation
 from deep_scan_workbench import non_negative_int
 from workbench_constants import (
     DIFF_TARGET_KINDS,
@@ -342,7 +341,10 @@ def parse_args(description: str) -> argparse.Namespace:
     release_finding_remediation_claim.add_argument("--request-id", required=True)
     release_finding_remediation_claim.add_argument("--action-token", required=True)
 
-    remediation.register_cancel_finding_remediation_request(subparsers)
+    cancel_remediation = subparsers.add_parser("cancel-finding-remediation-request")
+    cancel_remediation.add_argument("--occurrence-id", required=True)
+    cancel_remediation.add_argument("--request-id", required=True)
+    cancel_remediation.add_argument("--action-token", required=True)
 
     set_finding_remediation = subparsers.add_parser("set-finding-remediation")
     set_finding_remediation.add_argument("--occurrence-id", required=True)

plugins/codex-security/scripts/workbench_constants.py

@@ -16,18 +16,7 @@
 FINDING_SEVERITIES = ("critical", "high", "medium", "low", "informational")
 FINDING_STATUSES = ("open", "closed")
 FINDING_CLOSE_REASONS = ("already_fixed", "wont_fix", "false_positive")
-REMEDIATION_STATES = (
-    "idle",
-    "requested",
-    "generated",
-    "applied",
-    "verifying",
-    "verified",
-    "failed",
-    "superseded",
-)
 REMEDIATION_UPDATE_STATES = ("generated", "applied", "verifying", "verified", "failed")
-REMEDIATION_PENDING_ACTIONS = ("generate", "apply", "verify")
 EXPORT_FORMATS = ("csv", "json", "sarif")
 ARTIFACTS = {
     "coverage": "coverage.json",

plugins/codex-security/scripts/workbench_db.py

@@ -1096,7 +1096,7 @@ def pin_legacy_manifest_digest(
     connection: sqlite3.Connection, scan_id: str, manifest_digest: str
 ) -> None:
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         scan = require_scan(connection, scan_id)
         current = scan["seal_manifest_digest"]
         if current is not None and current != manifest_digest:
@@ -1106,10 +1106,6 @@ def pin_legacy_manifest_digest(
                 "UPDATE scans SET seal_manifest_digest = ? WHERE id = ?",
                 (manifest_digest, scan["id"]),
             )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
 
 
 def complete_scan(
@@ -1543,17 +1539,13 @@ def add_warning() -> None:
     manifest_digest = published_manifest_digest(scan_dir, manifest)
     if prepare_only:
         connection.execute("BEGIN IMMEDIATE")
-        try:
+        with connection:
             updated = connection.execute(
                 "UPDATE scans SET completion_warnings_json = ? WHERE id = ? AND status = 'running'",
                 (json.dumps(warnings), scan["id"]),
             )
             if updated.rowcount != 1:
                 raise SystemExit("Only a running scan can be prepared for completion.")
-            connection.commit()
-        except BaseException:
-            connection.rollback()
-            raise
         context = scan_context(connection, scan["id"])
         context["targetWarnings"] = target_warnings
         return context
@@ -1686,7 +1678,7 @@ def register_cli_scan(connection: sqlite3.Connection, args: argparse.Namespace)
     workspace_id = str(uuid.uuid4())
 
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         archive_scan(connection, args, scan_dir, timestamp, require_canonical_scan_directory)
         target_id = ensure_security_target(connection, str(repository))
         if parent_scan_id is not None:
@@ -1741,10 +1733,6 @@ def register_cli_scan(connection: sqlite3.Connection, args: argparse.Namespace)
         )
         if workflow_id is not None:
             register_workflow_scan(connection, workflow_id, scan_id, str(scan_dir), timestamp)
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     scan = require_scan(connection, scan_id)
     return {
         "contract": scan_contract(scan),
@@ -1877,7 +1865,7 @@ def set_finding_triage(connection: sqlite3.Connection, args: argparse.Namespace)
     note = optional_text(args.note, maximum=2400)
     require_close_note(close_reason, note)
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         timestamp = now()
         occurrence = require_occurrence(connection, args.occurrence_id)
         if args.status == "closed":
@@ -1949,10 +1937,6 @@ def set_finding_triage(connection: sqlite3.Connection, args: argparse.Namespace)
             """,
             (occurrence["id"], args.status, close_reason, note, timestamp),
         )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return scan_context(connection, occurrence["scan_id"])
 
 
@@ -2276,7 +2260,7 @@ def set_finding_remediation(
     action_token = require_uuid(args.action_token, "action-token")
     summary = optional_text(args.summary, maximum=2400)
     verification_summary = optional_text(args.verification_summary, maximum=2400)
-    try:
+    with connection:
         occurrence = require_occurrence(connection, args.occurrence_id)
         require_finding_open(connection, occurrence["id"])
         scan = require_scan(connection, occurrence["scan_id"])
@@ -2391,10 +2375,6 @@ def set_finding_remediation(
             raise SystemExit(
                 "This remediation request changed. Refresh it before recording an update."
             )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
     return scan_context(connection, occurrence["scan_id"])
 
 
@@ -2899,7 +2879,7 @@ def backfill_legacy_finding_details(connection: sqlite3.Connection, scan: sqlite
         return
 
     connection.execute("BEGIN IMMEDIATE")
-    try:
+    with connection:
         current = require_scan(connection, scan["id"])
         recorded_digest = current["seal_manifest_digest"]
         if recorded_digest is not None and recorded_digest != manifest_digest:
@@ -2917,10 +2897,6 @@ def backfill_legacy_finding_details(connection: sqlite3.Connection, scan: sqlite
                 "UPDATE scans SET seal_manifest_digest = ? WHERE id = ?",
                 (manifest_digest, scan["id"]),
             )
-        connection.commit()
-    except BaseException:
-        connection.rollback()
-        raise
 
 
 def legacy_finding_matches(row: sqlite3.Row, finding: Any) -> bool: