scripts/workbench_finding_workflows.py

"""Durable state for the opt-in local scan, publication, and dedupe workflow."""

from __future__ import annotations

import argparse
import hashlib
import json
import sqlite3
import sys
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))
from workbench_target import directory_content_digest, git_output, git_revision

WORKFLOW_BINDINGS = {
    "repositoryPath": "repository_path",
    "scanRequestDigest": "scan_request_digest",
    "scanId": "scan_id",
    "scanDir": "scan_dir",
    "artifactDigest": "artifact_digest",
    "destination": "destination",
}
WORKFLOW_STAGES = ("scan", "publish", "dedupe")


def read_workflow(connection: sqlite3.Connection, workflow_id: str) -> dict[str, Any] | None:
    row = connection.execute(
        "SELECT * FROM finding_workflows WHERE id = ?", (workflow_id,)
    ).fetchone()
    if row is None:
        return None
    state = {"id": row["id"], "stages": {}}
    for field, column in WORKFLOW_BINDINGS.items():
        if row[column] is not None:
            state[field] = row[column]
    if row["scope_repository_id"] is not None:
        state["scope"] = {"repositoryId": row["scope_repository_id"]}
    elif row["scope_all_repositories"] is not None:
        state["scope"] = {"allRepositories": bool(row["scope_all_repositories"])}
    results = json.loads(row["results_json"])
    for stage in WORKFLOW_STAGES:
        current = {"status": row[f"{stage}_status"]}
        if row[f"{stage}_error"] is not None:
            current["error"] = row[f"{stage}_error"]
        if stage in results:
            current["result"] = results[stage]
        state["stages"][stage] = current
    if "dedupePendingWrite" in results:
        state["stages"]["dedupe"]["pendingWrite"] = results["dedupePendingWrite"]
    return state


def save_workflow(connection: sqlite3.Connection, state: dict[str, Any], timestamp: str) -> None:
    values = {"id": state["id"]}
    values.update({column: state.get(field) for field, column in WORKFLOW_BINDINGS.items()})
    scope = state.get("scope", {})
    values.update(
        scope_repository_id=scope.get("repositoryId"),
        scope_all_repositories=scope.get("allRepositories"),
    )
    results = {}
    for stage in WORKFLOW_STAGES:
        current = state["stages"][stage]
        values[f"{stage}_status"] = current["status"]
        values[f"{stage}_error"] = current.get("error")
        if "result" in current:
            results[stage] = current["result"]
    if "pendingWrite" in state["stages"]["dedupe"]:
        results["dedupePendingWrite"] = state["stages"]["dedupe"]["pendingWrite"]
    values.update(
        results_json=json.dumps(results, allow_nan=False),
        created_at=timestamp,
        updated_at=timestamp,
    )
    updates = ", ".join(
        f"{column} = excluded.{column}" for column in values if column not in {"id", "created_at"}
    )
    connection.execute(
        f"INSERT INTO finding_workflows ({', '.join(values)}) "
        f"VALUES ({', '.join('?' for _ in values)}) ON CONFLICT(id) DO UPDATE SET {updates}",
        tuple(values.values()),
    )


def bind_workflow(state: dict[str, Any], binding: dict[str, Any]) -> None:
    for field, value in binding.items():
        if field not in WORKFLOW_BINDINGS and field != "scope":
            raise SystemExit("Unknown workflow binding.")
        if field in state and state[field] != value:
            raise SystemExit(
                f"Workflow {state['id']} is already bound to a different {field}. "
                "Use another --workflow-id."
            )
        state[field] = value


def register_workflow_scan(
    connection: sqlite3.Connection, workflow_id: str, scan_id: str, scan_dir: str, timestamp: str
) -> None:
    """Called inside the scan-registration transaction, before model execution."""
    state = read_workflow(connection, workflow_id)
    if state is None:
        raise SystemExit("The workflow must be started before registering its scan.")
    if state["stages"]["scan"]["status"] == "completed":
        raise SystemExit("The workflow scan is already complete.")
    previous = state.get("scanId")
    if previous is not None:
        row = connection.execute("SELECT status FROM scans WHERE id = ?", (previous,)).fetchone()
        if row is not None and row["status"] == "complete":
            raise SystemExit("Reuse the workflow's completed scan instead of registering another.")
    state["scanId"] = scan_id
    state["scanDir"] = scan_dir
    save_workflow(connection, state, timestamp)


def finding_workflow(
    connection: sqlite3.Connection, payload: dict[str, Any], timestamp: str
) -> dict[str, Any]:
    workflow_id = payload["id"]
    if not isinstance(workflow_id, str) or not workflow_id.strip():
        raise SystemExit("workflowId must be a nonempty string.")
    if payload["action"] == "get":
        return {"workflow": read_workflow(connection, workflow_id)}
    if payload["action"] == "source":
        target = Path(payload["repository"]).resolve(strict=True)
        return {
            "source": {
                "repository": str(target),
                "revision": git_revision(target),
                "refsDigest": hashlib.sha256(
                    (git_output(target, "show-ref") or "").encode()
                ).hexdigest(),
                "content": directory_content_digest(target, include_ignored=True),
            }
        }
    if payload["action"] == "get-review":
        row = connection.execute(
            "SELECT result_json FROM finding_workflow_reviews WHERE workflow_id = ? AND review_key = ?",
            (workflow_id, payload["key"]),
        ).fetchone()
        return {"review": json.loads(row["result_json"]) if row is not None else None}
    if payload["action"] == "save-review":
        binding = payload["binding"]
        source = binding["source"]
        scope = binding["scope"]
        with connection:
            connection.execute(
                """INSERT INTO finding_workflow_reviews
                (workflow_id, review_key, review_contract_version, codex_version,
                 source_repository_path, source_revision, source_refs_digest, source_content_digest,
                 scope_repository_id, scope_all_repositories, model, effort, settings_digest,
                 prompt_digest, contract_digest, result_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(workflow_id, review_key) DO NOTHING""",
                (
                    workflow_id,
                    payload["key"],
                    binding["version"],
                    binding["codexVersion"],
                    source["repository"],
                    source["revision"],
                    source["refsDigest"],
                    source["content"],
                    scope.get("repositoryId"),
                    scope.get("allRepositories"),
                    binding["model"],
                    binding["effort"],
                    binding.get("settingsDigest"),
                    binding["promptDigest"],
                    binding["contractDigest"],
                    json.dumps(payload["result"], allow_nan=False),
                    timestamp,
                ),
            )
        return {}
    connection.execute("BEGIN IMMEDIATE")
    with connection:
        state = read_workflow(connection, workflow_id)
        if state is None:
            state = {
                "id": workflow_id,
                "stages": {stage: {"status": "pending"} for stage in WORKFLOW_STAGES},
            }
        bind_workflow(state, payload.get("binding", {}))
        action = payload["action"]
        if action != "bind":
            stage = payload["stage"]
            if stage not in state["stages"]:
                raise SystemExit("Unknown workflow stage.")
            current = state["stages"][stage]
            if current["status"] != "completed":
                if action == "begin":
                    current["status"] = "running"
                elif action == "complete":
                    state["stages"][stage] = {"status": "completed", "result": payload["result"]}
                elif action == "fail":
                    current.update(status="failed", error=payload["error"])
                elif action == "prepare-dedupe":
                    current.update(result=payload["result"], pendingWrite=payload["pendingWrite"])
                else:
                    raise SystemExit("Unknown workflow action.")
        save_workflow(connection, state, timestamp)
        return {"workflow": state}


if __name__ == "__main__":
    argparse.ArgumentParser(description=__doc__).parse_args()