diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 0000000..7823352 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,143 @@ +# CONTEXT + +Shared vocabulary for the Kitaru integration. Product boundary and roadmap +consequences live in GitHub issue #7; Phase 1 is #8; Phase 2 is #9. Decisions +are recorded in [`docs/adr/`](docs/adr/). + +## Product + +**tracegrad** is the evidence-gated optimization engine. It owns the prompt/text +artifact, artifact lineage and hash, deterministic distillation, attribution, +failure-theme aggregation, evidence verification, edit synthesis, token-budget +discipline, rejection memory, human approval, standalone JSONL source, +post-deployment trends, and the verification summary / decision UX. + +**Kitaru** owns provider-specific trace ingestion, trace normalization, the +session graph, execution storage, evaluators and judge execution, cohorts, +replay, the worker runtime, tool-history replay, and execution / replay +visualization. + +tracegrad is not an observability or replay platform. Kitaru is not a mandatory +core dependency. tracegrad does not duplicate Kitaru's execution UI. + +```text +without Kitaru: proposal → deploy → next batch → trends +with Kitaru: proposal → replay verify → deploy → trends +``` + +`trends` stays in core. For core-only users it is verification after +deployment, via the next evaluated batch. For Kitaru users it is post-deployment +confirmation that replay-verified improvements persist in real traffic. + +## Install modes + +**Core-only.** `uv tool install tracegrad`. No Kitaru package, server, worker, +or observability integration. `tracegrad run --traces …`, `apply`, and `trends` +are unchanged. `import tracegrad` never requires Kitaru. + +**tracegrad + Kitaru.** `uv tool install "tracegrad[kitaru]"`, extra pinned +`kitaru>=0.22,<0.23`. The extra installs a **client**. It does not install +verification (ADR 0001). Credentials and server URL come from Kitaru's own +config (`kitaru login`). tracegrad stores no Kitaru secrets. `.tracegradrc` +holds only non-secret selection (cohort name, evaluation name); CLI flags +override. + +## Source vs batch (Phase 1) + +`--source kitaru` is a **fetch-and-map** that writes the JSONL the existing +pipeline already reads (ADR 0004). `--traces` and `--source kitaru` are +mutually exclusive. The deterministic core does not learn Kitaru exists. There +is no `TraceSource` protocol, no `sources/` package, and no `ingest.py` rewrite +(ADR 0010). All Kitaru SDK code lives under `src/tracegrad/integrations/kitaru/`. + +**Session** — a Kitaru recording (`SessionResponse`). Durable id is the UUID; +`number` is display-only (`#4811`). + +**Trace** — a tracegrad `Trace`. One session maps to at most one trace. + +**Root LLM node** — an `llm_call` with no `subagent_call` anywhere in its +ancestry, following `parent_index` **and** `secondary_parent_indexes`. The +graph is a DAG. A node reachable from a subagent is not root even when one of +its parents is (ADR 0005). Subagent prompts and tool outputs cannot become the +artifact or `Trace.output`. Never guess a missing system prompt. + +**Source drop** — a Session could not become a Trace. Named kebab-case reasons +(`system-prompt-unavailable`, `multiple-system-prompts`, +`judge-rationale-missing`, `judge-score-out-of-range`, +`judge-score-unsupported`, `judge-score-unavailable`, `output-unavailable`, +`input-unavailable`, `ambiguous-evaluation`, …). + +**Batch drop** — a Trace is not part of this Batch (the four reasons +`ingest.py` already uses, including `prompt-hash-partition` and +`rationale-below-quality-floor`). + +Source drops and batch drops are two tables, never merged. Merging them lets a +mapping bug hide behind a legitimate partition. + +**Snapshot** — mapped JSONL plus source fingerprint under `.tracegrad/`, written +before ingest. Re-runs read it; `--refresh` refetches. Cohort version is +resolved **once** per run (`cohort_version_id` in the fingerprint). `engine=format` +manifests are refused with a named error (ADR 0002). + +**Judge mapping** — ADR 0003. `Score` stays `[0, 1]`; out-of-range and +non-numeric evaluations drop by name rather than being rescaled. +`judge_fingerprint` is derived from `evaluator_name` + `evaluator_version`. A +conflicting manifest value is an error. + +## Verification (Phase 2) + +`tracegrad verify --backend kitaru --run ` reuses the cohort, evaluator, +and agent metadata Phase 1 persisted. `VerificationBackend` lives in +`ports.py`; the implementation lives under `integrations/kitaru/`. + +**No backend.** `tracegrad verify` prints an actionable message and exits +non-zero. That does not block `run` / `apply` / `trends`. A verify that exits 0 +having done nothing would read as *verified* in CI. + +**Tool policy (hard invariant).** Every tracegrad-created replay sets +`HistoryConfig(scope=COHORT_VERSION, on_miss=FAIL)`. No passthrough through the +tracegrad path. A novel call becomes `TOOL_HISTORY_MISS`, not a live production +side effect. + +**Override scope (hard invariant).** Only the root LLM system prompt is +overridden (`ReplayOverride.system_prompt`). After replay, assert root LLM +nodes carry the candidate and non-root LLM nodes carry their baseline +counterpart. Violations are `OVERRIDE_SCOPE_DIVERGENCE`. A failed `select_evaluation` +is `SELECT_EVALUATION_FAILED` (drop reason in detail). A mismatched +`evaluator_version` is `EVALUATOR_VERSION_MISMATCH`. Scores that cannot +be classified are `SCORE_UNCLASSIFIED`. All of these are **incomparable** +— not improved, not regressed — and stay in the per-session buckets. +Headline aggregates still come from Kitaru (ADR 0006). + +**Cohort constraint.** Mixed-agent-version cohorts are refused with a +per-version breakdown (ADR 0007). Baseline and candidate use the same evaluator +version. + +**Apply gate.** With a backend configured (the originating run persisted Kitaru +source metadata), `apply` refuses unless a persisted verification exists whose +`candidate_prompt_hash` equals the hash of what is about to be written. +`--force` overrides. Core-only users are unaffected (ADR 0009). + +**Persistence / resume.** `.tracegrad/verification/.json`. +Persist the Kitaru run id immediately after creation. An interrupted +verification resumes and watches the existing experiment run rather than +creating a duplicate. + +**Inspection.** Hand off the dashboard base plus real identifiers. Do not build +`--open` or `tracegrad inspect` (ADR 0008). + +**Verdict.** The verification report never prints `SHIP`, and never applies or +reverts on its own. The human decides. + +**Aggregates.** Headline numbers come from +`/api/v1/ui/experiment-runs/{id}/evaluation-aggregates` so they match the +Kitaru UI. That `/api/v1/ui/` namespace is a UI-support contract, contained by +the `<0.23` pin. + +## Do not build + +Langfuse / LangSmith / Braintrust importers. A tracegrad-native replay engine. +Worker infrastructure. Tool mocking or a history engine. Cohort storage. Replay +experiment orchestration. A full eval runner. A competing trace viewer. +`--open` / `inspect`. An `approve`/`apply` rename (deferred, ADR 0009). +Multi-artifact editing (README "Beyond system prompts" is unaffected). diff --git a/README.md b/README.md index 4a07c5b..aa1efa3 100644 --- a/README.md +++ b/README.md @@ -78,9 +78,16 @@ tracegrad init tracegrad run --traces batch.jsonl --manifest manifest.json --estimate # cost preview tracegrad run --traces batch.jsonl --manifest manifest.json # analyze, propose tracegrad apply # review, accept/reject +tracegrad trends # last two runs tracegrad status # budget, trends, ledgers ``` +Core-only install has no Kitaru package, server, worker, or observability +integration. `run --traces`, `apply`, and `trends` work without it. +`tracegrad verify` without a backend prints an actionable message and exits +non-zero — a verify that exits 0 having done nothing would read as verified +in CI. + `run` prints one review card per proposed edit — the diff, the verbatim quotes behind it, and any flags — and writes the proposal to `.tracegrad/`. Nothing touches your prompt until `tracegrad apply`. `apply --revert` restores the @@ -94,6 +101,68 @@ the last two runs. Attribution is one model call per trace, so `run --jobs 8` is worth setting on any batch above a handful of traces. +### Optional: Kitaru as a trace source and replay backend + +Kitaru is an **optional extra**, pinned `kitaru>=0.22,<0.23`. The base package +has no Kitaru dependency; `import tracegrad` never requires it. The extra +installs a **client**. It does not install verification — replays run on a +worker in *your* agent virtualenv. tracegrad stores no Kitaru secrets; use +`kitaru login`. + +```sh +uv tool install "tracegrad[kitaru]" +# from this repo until PyPI: +# uv tool install "git+https://github.com/dnth/tracegrad[kitaru]" +kitaru login +``` + +`.tracegradrc` may hold only non-secret selection; flags override: + +```toml +[kitaru] +cohort = "support-production" +evaluation = "quality" +``` + +`--source kitaru` is a fetch-and-map: it writes JSONL the existing pipeline +already reads, then ingest runs unchanged. `--traces` and `--source kitaru` +are mutually exclusive. `engine = "format"` manifests are refused with a +named error. Mapped traces and a source fingerprint are snapshotted under +`.tracegrad/sources/kitaru/` before ingest; re-runs read the snapshot; +`--refresh` refetches. The cohort version is resolved once per run. + +```sh +tracegrad run \ + --source kitaru \ + --kitaru-cohort support-production \ + --kitaru-evaluation quality \ + --manifest manifest.json +``` + +`judge_fingerprint` is derived from the evaluator (`quality@3`). A conflicting +manifest value is an error. Set the manifest fingerprint to that derived +identity. + +After a proposal, verify the candidate against the same frozen cohort. This +needs a running Kitaru server, a worker in the agent's virtualenv, and a +registered agent version: + +```sh +tracegrad verify --backend kitaru --run run-0001 +tracegrad apply --all +``` + +`apply` then refuses unless a persisted verification exists whose +`candidate_prompt_hash` matches what is about to be written. `--force` +overrides. Core-only JSONL users are unaffected. The verification report +never prints `SHIP` and never applies or reverts on its own. + +Every tracegrad-created replay sets recorded tool history with `on_miss=fail`. +Passthrough is not reachable through the tracegrad path. + +`trends` stays in core either way: without Kitaru it is the next-batch check +after deploy; with Kitaru it confirms a replay-verified change in real traffic. + ### Try it on the bundled example `example/` holds a synthetic 13-trace batch for a support agent, its manifest, @@ -163,8 +232,10 @@ tracegrad reads an optional TOML file named `.tracegradrc` from the project root If it is absent, `neverDelete = []`, `minEffect = 0.05`, `minCoverage = 0.8`, and `convergenceRuns = 2` apply. The default attribution and synthesis harness providers are `openai` and `claude`. The supported top-level keys are -`neverDelete`, `minEffect`, `minCoverage`, `convergenceRuns`, and -`harness_presets`; see the package configuration model for the preset fields. +`neverDelete`, `minEffect`, `minCoverage`, `convergenceRuns`, +`harness_presets`, and `kitaru`; see the package configuration model for the +preset fields. The `[kitaru]` table holds only non-secret selection (`cohort`, +`evaluation`); credentials stay in `kitaru login`. ```toml neverDelete = ["prompt/identity"] @@ -200,11 +271,13 @@ measured it. ``` .tracegrad/ - distilled/ content-addressed distilled traces — the only text a quote may cite - ledgers/ append-only JSONL: runs, gaps, rejections, applied edits - reports/ per-run theme counts, the input to trend comparison - runs/ per-run proposal, resume checkpoint, and autopsy of dropped proposals - snapshots/ the prompt as it was before each apply + distilled/ content-addressed distilled traces — the only text a quote may cite + ledgers/ append-only JSONL: runs, gaps, rejections, applied edits + reports/ per-run theme counts, the input to trend comparison + runs/ per-run proposal, resume checkpoint, and autopsy of dropped proposals + snapshots/ the prompt as it was before each apply + sources/ optional mapped JSONL from `--source kitaru`, plus the source fingerprint + verification/ persisted replay-verification state (resume-safe) ``` Everything except `apply` only reads and appends. A killed run resumes from its diff --git a/docs/adr/0001-tracegrad-submits-verification-it-does-not-host-it.md b/docs/adr/0001-tracegrad-submits-verification-it-does-not-host-it.md new file mode 100644 index 0000000..1506be3 --- /dev/null +++ b/docs/adr/0001-tracegrad-submits-verification-it-does-not-host-it.md @@ -0,0 +1,49 @@ +# 0001. tracegrad submits verification; it does not host it + +## Status + +Accepted + +## Context + +Kitaru replay re-executes the user's real agent code. `ExperimentRunCreateRequest` +requires an `agent_version_id`, and `AgentVersionResponse` carries a `RunSpec` +with a shell command. Kitaru's README is explicit that replays run "on workers +in your environment: your virtualenv, your credentials, your network." + +tracegrad is installed via `uv tool install`, into an isolated virtualenv, which +by construction is not the virtualenv the agent runs in. The optional extra +`tracegrad[kitaru]` installs a client library. It does not install a worker, a +server, Docker, or the user's agent. + +Four capabilities an earlier draft assumed do not exist in the Kitaru 0.22 API, +and treating the extra as "gaining replay verification" would hide the real +prerequisites. + +## Decision + +tracegrad submits verification; it does not host it. + +Phase 2 requires the user to already have: + +1. a running Kitaru server (FastAPI + Postgres, via Docker), +2. a worker process running in the virtualenv where their agent code lives, +3. their agent instrumented with a Kitaru adapter and registered as an agent + version. + +`tracegrad verify` preflights those before spending anything: probe the server, +confirm a live worker claims this agent version (workers report `last_seen_at`), +and confirm the agent version and cohort version resolve. A replay experiment +is paid and slow; "no worker is polling" should surface in milliseconds. + +tracegrad never hosts a worker and does not try to. + +## Consequences + +- The extra is a client pin (`kitaru>=0.22,<0.23`), not a verification runtime. +- Core-only users keep a complete JSONL workflow with no Kitaru package, server, + or worker. +- Documentation must not claim that `tracegrad[kitaru]` "gains replay + verification" by itself. +- Preflight cannot check whether an adapter honours a system-prompt override; + that gap is closed by ADR 0006's post-replay assertion, not by prediction. diff --git a/docs/adr/0002-prompt-lineage-from-recorded-prompts.md b/docs/adr/0002-prompt-lineage-from-recorded-prompts.md new file mode 100644 index 0000000..69835b8 --- /dev/null +++ b/docs/adr/0002-prompt-lineage-from-recorded-prompts.md @@ -0,0 +1,44 @@ +# 0002. Prompt lineage comes from recorded system prompts + +## Status + +Accepted + +## Context + +tracegrad's ingest partitions a batch on `prompt_hash`. Rates are only +meaningful within one prompt version, so only the dominant partition survives. +The hash has to name the artifact under study. + +Kitaru sessions do not carry a tracegrad prompt hash. They do carry per-node +`system_prompt_selector` pointers into node inputs. Those recorded prompts are +the only honest lineage: guessing a missing system prompt would attribute +failures to an artifact the session never ran. + +`engine = "format"` manifests render differently per request. Hashing the +recorded (already rendered) prompts of a format template would give N partitions +of size one, and ingest would keep a batch of one. Reverse-templating those +recordings back onto a format template is a separate problem and is not solved +here. + +## Decision + +Hash the recorded system prompt extracted from root LLM nodes. The mapped batch +must be single-valued on that hash, which is what ingest already enforces. + +This supports `engine = "none"` only. When the manifest declares +`engine = "format"`, refuse the Kitaru source with a named error rather than +proceeding into a batch that collapses. + +Never infer or guess a missing system prompt. Zero unique recorded prompts is +`system-prompt-unavailable`. More than one unique recorded prompt on the root +LLM nodes of one session is `multiple-system-prompts`. + +Reverse-templating for `format` prompts is deferred, not rejected. + +## Consequences + +- Kitaru-sourced runs with a `format` manifest fail closed. +- `prompt_hash` on a mapped `Trace` is `text_hash` of the extracted system + prompt, so the rest of the pipeline does not learn Kitaru exists. +- Multi-prompt sessions drop at the source rather than poisoning the partition. diff --git a/docs/adr/0003-judge-scores-are-accepted-not-rescaled.md b/docs/adr/0003-judge-scores-are-accepted-not-rescaled.md new file mode 100644 index 0000000..eab2f6b --- /dev/null +++ b/docs/adr/0003-judge-scores-are-accepted-not-rescaled.md @@ -0,0 +1,47 @@ +# 0003. Judge scores are accepted, not rescaled + +## Status + +Accepted + +## Context + +`schema.Score` is bounded `[0, 1]`. Kitaru evaluations carry a `score` that may +be a bool, a float of unknown range, a string, or a categorical pair of score +plus value, plus an optional `passed` flag and an `explanation`. + +Rescaling an out-of-range float against an assumed range would invent a judge +the user did not run. A `--kitaru-score-range` flag would be a declared +normalization, which is a product decision of its own. + +tracegrad attribution is built on the rationale, not just the score. A missing +`explanation` is not a score of zero; it is a session that cannot be attributed. + +## Decision + +Map Kitaru evaluations onto `Judge` by name, never by rescaling: + +| Kitaru | tracegrad | +|---|---| +| `bool` score, or `passed` flag | 1.0 / 0.0 | +| `float` in [0, 1] | as-is | +| `float` outside [0, 1] | drop `judge-score-out-of-range` | +| `str` / `categorical` | drop `judge-score-unsupported` | +| missing `explanation` | drop `judge-rationale-missing` | + +`judge_fingerprint` is derived from the resolved `evaluator_name` + +`evaluator_version` and overrides the manifest value. A conflicting manifest +value is an error, not a tiebreak. Drift detection is only worth having if it +reads the thing that actually drifts. + +One evaluation name resolving to more than one `evaluator_version` across the +cohort is `ambiguous-evaluation`: refuse rather than mix. + +A declared `--kitaru-score-range` is deferred, not rejected. + +## Consequences + +- Out-of-range and non-numeric evaluations never become a `Trace`. +- Core ingest still sees only `[0, 1]` scores. +- Changing the Kitaru evaluator version changes the fingerprint, so trends + across that change are not comparable. diff --git a/docs/adr/0004-the-kitaru-source-snapshots-mapped-traces.md b/docs/adr/0004-the-kitaru-source-snapshots-mapped-traces.md new file mode 100644 index 0000000..b647efb --- /dev/null +++ b/docs/adr/0004-the-kitaru-source-snapshots-mapped-traces.md @@ -0,0 +1,54 @@ +# 0004. The Kitaru source snapshots mapped traces + +## Status + +Accepted + +## Context + +`--source kitaru` has to turn a frozen Kitaru cohort into the JSONL the existing +pipeline already reads. If the fetch lived inside ingest, the deterministic core +would learn Kitaru exists, and a run would be unreproducible the moment the +server was unreachable. + +Cohort names resolve to a moving "latest" version. Re-resolving `latest` inside +one run would mix two populations in one batch. + +## Decision + +`--source kitaru` is a fetch-and-map step that writes the JSONL the pipeline +already reads. The deterministic core does not learn Kitaru exists. + +The mapped `Trace` objects are written to `.tracegrad/` as JSONL alongside the +source fingerprint, before ingest. Re-runs read the snapshot; `--refresh` +refetches and rewrites. + +Prefer immutable cohort versions. Given only a cohort name: resolve the current +version once, persist the immutable `cohort_version_id`, use it for the whole +run, and include it in the source fingerprint. Never re-resolve `latest` within +a run. + +The source fingerprint is: + +```json +{ + "source": "kitaru", + "cohort_id": "...", + "cohort_version_id": "...", + "evaluation_name": "quality", + "evaluator_id": "...", + "evaluator_version": 3, + "agent_id": "...", + "mapping_version": 1 +} +``` + +A run is reproducible from the snapshot with the server unreachable. + +## Consequences + +- `ingest.py` does not change for Kitaru (see ADR 0010). +- Phase 2 reuses the cohort, evaluator, and agent metadata this snapshot + persists. +- `--traces` and `--source kitaru` are mutually exclusive: one batch, one + origin. diff --git a/docs/adr/0005-traces-are-extracted-from-root-nodes.md b/docs/adr/0005-traces-are-extracted-from-root-nodes.md new file mode 100644 index 0000000..c62a84c --- /dev/null +++ b/docs/adr/0005-traces-are-extracted-from-root-nodes.md @@ -0,0 +1,53 @@ +# 0005. Traces are extracted from root LLM nodes + +## Status + +Accepted + +## Context + +A Kitaru session is a DAG of nodes (`parent_index` and +`secondary_parent_indexes`). Node types include `llm_call`, `tool_call`, +`subagent_call`, and `span`. Subagents have their own instructions. Tool nodes +carry tool outputs, not the artifact's response. + +`SessionResponse.inputs` / `.outputs` are typed `Any` with no session-level +selector. Stringifying them as a fallback would let a tool payload or a nested +object become `Trace.input` / `Trace.output`. + +`Trace.input` is a `StrictStr`, so a multi-turn session cannot be represented +losslessly. + +A node reachable from a subagent is not a root LLM node even when one of its +parents is a root LLM node. Following only `parent_index` would miss that. + +## Decision + +A **root LLM node** is an `llm_call` with no `subagent_call` anywhere in its +ancestry, following `parent_index` **and** `secondary_parent_indexes`. + +- `input` is `input_text_selector` resolved on the **first** root LLM node. +- `output` is `output_text_selector` resolved on the **last** root LLM node. +- The system prompt is `system_prompt_selector` resolved against node inputs + for each root LLM node, uniqueness-checked per ADR 0002. +- `meta.model` comes from `SessionNodeResponse.model` on the root LLM nodes. +- `trace_id` is `SessionResponse.id` (UUID). `number` is carried for display + (`#4811`); the durable key stays the UUID. + +Session-level inputs/outputs are never stringified as a fallback. Tool outputs +cannot become the final output, because only `llm_call` nodes are consulted. A +subagent's system prompt can never become the artifact. + +Selector resolution failures drop by name (`input-unavailable`, +`output-unavailable`, `system-prompt-unavailable`). Multi-turn is lossy: a +session collapses to its first root input and last root output. Report it; do +not hide it. + +`trace.meta["trajectory"]` is out of scope: `TraceMeta` is `extra="forbid"` +with one field. Kitaru remains the system of record for trajectories. + +## Consequences + +- Mapping is provider-agnostic: it reads Kitaru's normalized session graph. +- Tests must cover DAG reachability via a secondary parent, not only trees. +- Display reports may print `#4811` while every persisted key stays the UUID. diff --git a/docs/adr/0006-override-scope-is-verified-not-predicted.md b/docs/adr/0006-override-scope-is-verified-not-predicted.md new file mode 100644 index 0000000..5c00582 --- /dev/null +++ b/docs/adr/0006-override-scope-is-verified-not-predicted.md @@ -0,0 +1,46 @@ +# 0006. Override scope is verified, not predicted + +## Status + +Accepted + +## Context + +`ReplayOverride.system_prompt` is a single unscoped string. The Kitaru adapters +disagree about where it lands: + +- OpenAI Agents scopes it to the starting agent — `if data.agent is not + starting_agent: return model_data`. Subagents keep their own instructions. +- LangGraph applies it in `_model_request` with no starting-agent check. + **Every** model call gets the candidate, subagents included. + +`AgentCapabilities` carries only tools, MCP servers, and skills. Whether an +adapter honours a system-prompt override, and where, is not exposed over the +wire. Preflight cannot predict it. + +On a multi-agent LangGraph app, a naive submission moves two variables and +reports one: the candidate prompt *and* every subagent's prompt. A report that +says "the prompt caused this" would be a lie. + +## Decision + +Change exactly one variable: the system prompt, passed as +`ReplayOverride.system_prompt`. Hold fixed: baseline inputs, agent code and +version, model, model params, cohort, evaluator version, recorded tool history. + +After each replay, fetch the result session's nodes and assert: + +- every root LLM node carries the candidate prompt, and +- every non-root LLM node carries what its baseline counterpart carried. + +A session failing either is `OVERRIDE_SCOPE_DIVERGENCE` and is reported +**incomparable** — not improved, not regressed. + +A tool-history miss is `TOOL_HISTORY_MISS` and likewise incomparable. + +## Consequences + +- Headline improved/regressed counts never include diverged sessions. +- Adapter disagreement becomes a typed, per-session fact instead of a silent + confounder. +- Preflight stays cheap; the assertion is the backstop. diff --git a/docs/adr/0007-verification-requires-a-single-agent-version-cohort.md b/docs/adr/0007-verification-requires-a-single-agent-version-cohort.md new file mode 100644 index 0000000..b6e509c --- /dev/null +++ b/docs/adr/0007-verification-requires-a-single-agent-version-cohort.md @@ -0,0 +1,43 @@ +# 0007. Verification requires a single-agent-version cohort + +## Status + +Accepted + +## Context + +`ExperimentRunCreateRequest.agent_version_id` applies to the whole run, while +`evaluate_baselines=True` scores the *stored* baselines under whatever version +each recorded. A mixed-agent-version cohort silently replays sessions under +code they never ran on. + +A report that says "the prompt caused this" when the agent code also moved is +worse than refusing the cohort. + +The same confounder exists on the judge: comparing a candidate scored under +evaluator version 3 against stored baseline scores from version 2 is not a +prompt comparison. + +Kitaru can derive a single-agent-version cohort version via `remove_session_ids`. +Doing that automatically would hide the population change from the user. + +## Decision + +Every session in the cohort version must share one `agent_version_id`. A mixed +cohort is refused, with the version breakdown and session count per version. + +This will refuse some real cohorts. That is acceptable. + +Baseline and candidate always use the same evaluator version. A comparison +against stale stored judge output from a different evaluator version is marked +incomparable. + +Deriving a single-agent-version cohort version via `remove_session_ids` is +deferred, not rejected. + +## Consequences + +- Phase 2 will not start an experiment run against a mixed cohort. +- Evaluator version is pinned from the Phase 1 source fingerprint, not + re-resolved to `latest`. +- Users who need a mixed population must build a new cohort version themselves. diff --git a/docs/adr/0008-inspection-is-handed-off-by-identifier.md b/docs/adr/0008-inspection-is-handed-off-by-identifier.md new file mode 100644 index 0000000..63ee40f --- /dev/null +++ b/docs/adr/0008-inspection-is-handed-off-by-identifier.md @@ -0,0 +1,37 @@ +# 0008. Inspection is handed off by identifier + +## Status + +Accepted + +## Context + +The Kitaru API exposes **no** UI URL for an experiment, run, session, or replay. +`client/dashboard_urls.py` has only `get_investigation_review_url`. +Reconstructing undocumented dashboard routes would break silently on the next +Kitaru UI change and would duplicate Kitaru's execution viewer inside +tracegrad, which the product boundary forbids. + +`--open` and `tracegrad inspect` were in an earlier draft. They cannot be +honestly implemented against this API. + +## Decision + +Inspection hands off the supported dashboard base plus real identifiers. +`--open` / `tracegrad inspect` are **not** built. + +Identifiers (experiment id, experiment-run id, session ids, replay ids) are +persisted on the verification record, so wiring a URL helper later is small. + +tracegrad builds no competing execution viewer. Proposal approval and +application stay in tracegrad, not the Kitaru UI. + +An upstream request for experiment / session / compare URL helpers is a +Kitaru-side follow-up, not a tracegrad feature. + +## Consequences + +- The verification report prints identifiers and the configured Kitaru server + URL; the user inspects in Kitaru's UI. +- No new CLI commands for browsing sessions. +- A later URL helper can read the persisted identifiers without a schema break. diff --git a/docs/adr/0009-apply-is-gated-on-a-matching-verification.md b/docs/adr/0009-apply-is-gated-on-a-matching-verification.md new file mode 100644 index 0000000..e6f12df --- /dev/null +++ b/docs/adr/0009-apply-is-gated-on-a-matching-verification.md @@ -0,0 +1,45 @@ +# 0009. Apply is gated on a matching verification + +## Status + +Accepted + +## Context + +With a replay backend, the point of `verify` is that a candidate is replayed +against the same frozen cohort before it can be applied. Gating `apply` on +"a verify ran for this run id" is not enough: the user can verify, then +hand-edit the proposal or accept a subset of edits, and the text that is about +to be written was never replayed. + +Core-only users have no backend. Blocking `apply` for them would break the +standalone workflow the umbrella issue guarantees. + +Renaming `apply` to `approve` was considered so that "approve for testing" +and "apply for real" could be different verbs. That is a broader CLI change +and is not required to make the gate real. + +## Decision + +With a backend configured — meaning the originating run persisted Kitaru source +metadata — `apply` refuses unless a persisted verification exists whose +`candidate_prompt_hash` equals the hash of what is about to be written. +`--force` overrides. + +Matching on the hash rather than the run id is what makes the gate real: +verify, hand-edit, and the gate correctly notices the text was never verified. + +Core-only users are unaffected: JSONL runs have no Kitaru source metadata, so +`apply` behaves as it does today. + +An `approve` / `apply` rename is deferred, not rejected. + +## Consequences + +- `apply --all` after a full-proposal verify is the matching happy path. +- Partial acceptance, hand-edits, and a stale candidate all refuse unless + `--force`. +- `run`, `apply` (core-only), and `trends` keep working without a backend. +- `tracegrad verify` with no backend prints an actionable message and exits + non-zero. A verify that exits 0 having done nothing would read as *verified* + to every downstream CI step. diff --git a/docs/adr/0010-no-tracesource-protocol-in-phase-1.md b/docs/adr/0010-no-tracesource-protocol-in-phase-1.md new file mode 100644 index 0000000..f2c9d7b --- /dev/null +++ b/docs/adr/0010-no-tracesource-protocol-in-phase-1.md @@ -0,0 +1,43 @@ +# 0010. No TraceSource protocol in Phase 1 + +## Status + +Accepted + +## Context + +An earlier draft introduced a `TraceSource` protocol and a `sources/` package, +and proposed splitting parsing from validation in `ingest.py`, so Kitaru and +JSONL could be two in-process sources behind one interface. + +With ADR 0004 there is no second in-process source to abstract over: Kitaru is +a fetch-and-map that writes the JSONL ingest already reads. +`ingest_traces` already accepts `Sequence[Trace]`, a path, or parsed records, +so the parsing/validation seam already exists. + +A protocol with one implementation is an extra indirection the deterministic +core would have to import, which is how Kitaru would leak across the boundary. + +`ports.py` exists so the orchestrator can hold a backend without becoming +backend-aware. That is the right place for `VerificationBackend` in Phase 2. +It is the wrong place for a trace source that does not enter the core. + +## Decision + +No `TraceSource` protocol. No `sources/` package. No change to `ingest.py` for +Kitaru. + +All Kitaru SDK imports live under `src/tracegrad/integrations/kitaru/`. +`import tracegrad` never requires Kitaru. Using a Kitaru path without the extra +returns an actionable install message rather than an `ImportError`. + +`VerificationBackend` goes in `ports.py`. The Kitaru implementation lives under +`src/tracegrad/integrations/kitaru/`. + +## Consequences + +- The core pipeline, ingest, and schema stay Kitaru-ignorant. +- Adding a second observability backend later is a new fetch-and-map, not a + new core protocol. +- Phase 2 can hold a backend through `ports.py` without teaching `pipeline.py` + about Kitaru. diff --git a/pyproject.toml b/pyproject.toml index 88fc4d0..cde55a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,9 @@ dependencies = [ "pydantic>=2.7", ] +[project.optional-dependencies] +kitaru = ["kitaru>=0.22,<0.23"] + [project.scripts] tracegrad = "tracegrad.cli:main" diff --git a/src/tracegrad/apply.py b/src/tracegrad/apply.py index 1288348..8716589 100644 --- a/src/tracegrad/apply.py +++ b/src/tracegrad/apply.py @@ -277,6 +277,47 @@ def snapshot_template(project_root: str | Path, run_id: str, template: Path) -> return target +def _candidate_text_and_rejections( + prompt: str, + proposal: Proposal, + accepted_indices: Iterable[int], +) -> tuple[str, tuple[str, ...]]: + """The text that would be written, plus resolve_edits rejection reasons. + + One path for verify, the apply gate, and apply_proposal so the ADR 0009 + hash is the bytes that land on disk, and ApplyResult still sees rejections. + """ + + selected = sorted({index for index in accepted_indices}) + for index in selected: + if index < 0 or index >= len(proposal.edits): + raise ApplyError(f"no such edit index: {index}") + accepted = [proposal.edits[index].edit for index in selected] + if not accepted: + return prompt, () + inventory = build_inventory(prompt) + resolution = resolve_edits(inventory, accepted) + return ( + apply_resolved(prompt, resolution.resolved), + tuple(item.reason for item in resolution.rejected), + ) + + +def candidate_prompt( + prompt: str, + proposal: Proposal, + accepted_indices: Iterable[int], +) -> str: + """The text that would be written if these indices were accepted. + + Used by verify, the apply gate, and apply_proposal so the ADR 0009 hash + gate hashes the same bytes that are written. + """ + + text, _rejections = _candidate_text_and_rejections(prompt, proposal, accepted_indices) + return text + + def apply_proposal( project_root: str | Path, proposal: Proposal, @@ -306,28 +347,17 @@ def apply_proposal( ) selected = sorted({index for index in accepted_indices}) - for index in selected: - if index < 0 or index >= len(proposal.edits): - raise ApplyError(f"no such edit index: {index}") - + # ADR 0009's hash gate is only sound if apply writes the same bytes verify + # hashed. _candidate_text_and_rejections is that shared path. + updated, resolution_rejections = _candidate_text_and_rejections( + current, proposal, selected + ) + selected_set = set(selected) accepted = [proposal.edits[index].edit for index in selected] rejected = [ - item.edit for index, item in enumerate(proposal.edits) if index not in set(selected) + item.edit for index, item in enumerate(proposal.edits) if index not in selected_set ] - if not accepted: - return ApplyResult( - template_file=template, - applied_prompt_hash=text_hash(current), - accepted=(), - rejected=tuple(rejected), - snapshot=None, - unchanged=True, - ) - - inventory = build_inventory(current) - resolution = resolve_edits(inventory, accepted) - updated = apply_resolved(current, resolution.resolved) if updated == current: return ApplyResult( template_file=template, @@ -336,7 +366,7 @@ def apply_proposal( rejected=tuple(rejected), snapshot=None, unchanged=True, - resolution_rejections=tuple(item.reason for item in resolution.rejected), + resolution_rejections=resolution_rejections, ) snapshot = snapshot_template(project_root, proposal.run_id, template) @@ -364,7 +394,7 @@ def apply_proposal( accepted=tuple(accepted), rejected=tuple(rejected), snapshot=snapshot, - resolution_rejections=tuple(item.reason for item in resolution.rejected), + resolution_rejections=resolution_rejections, ) diff --git a/src/tracegrad/cli.py b/src/tracegrad/cli.py index b9f3478..30e5552 100644 --- a/src/tracegrad/cli.py +++ b/src/tracegrad/cli.py @@ -19,6 +19,7 @@ ApplyError, Proposal, apply_proposal, + candidate_prompt, is_stale, latest_run_id, load_proposal, @@ -32,10 +33,16 @@ attribute_batch, resolve_attribution_backend, ) +from .canonical import text_hash from .config import ConfigError, load_config from .distill import DistillConfig, DistillError, distill_batch, render_manifest_prompt, store_batch from .gates import REJECTION_MEMORY_FILENAME, RejectionMemory, measure_tokens from .ingest import IngestError, ingest_traces +from .integrations.kitaru.errors import ( + NO_BACKEND_MESSAGE, + KitaruError, + KitaruSourceError, +) from .inventory import InventoryError, build_inventory from .pipeline import ( RUN_LEDGER_FILENAME, @@ -49,10 +56,60 @@ from .state import StateError, initialize, load_jsonl from .synthesize import SynthesisError from .trends import compare, convergence, format_trend, hysteresis +from .verify import VerifyError DEFAULT_RUN_ID_PREFIX = "run" +def _resolve_traces_argument( + args: argparse.Namespace, + out: TextIO, + *, + run_id: str | None = None, +) -> str | Path: + """JSONL path the existing pipeline reads. ``--source kitaru`` snapshots first.""" + + traces = getattr(args, "traces", None) + source = getattr(args, "source", None) + if traces and source: + raise KitaruSourceError("--traces and --source are mutually exclusive") + if not traces and not source: + raise KitaruSourceError("provide --traces or --source kitaru") + if source is None: + return traces + if source != "kitaru": + raise KitaruSourceError(f"unknown --source {source!r}; only 'kitaru' is supported") + + from .integrations.kitaru.require import require_kitaru + + require_kitaru() + config = load_config(args.project_root) + cohort = getattr(args, "kitaru_cohort", None) or config.kitaru.cohort + evaluation = getattr(args, "kitaru_evaluation", None) or config.kitaru.evaluation + if not cohort or not evaluation: + raise KitaruSourceError( + "--source kitaru requires --kitaru-cohort and --kitaru-evaluation " + "(or kitaru.cohort / kitaru.evaluation in .tracegradrc)" + ) + from .integrations.kitaru.source import prepare_kitaru_source + + prepared = prepare_kitaru_source( + project_root=args.project_root, + manifest=load_manifest(args.manifest), + cohort_name=cohort, + evaluation_name=evaluation, + cohort_version=getattr(args, "kitaru_cohort_version", None), + refresh=bool(getattr(args, "refresh", False)), + run_id=run_id, + ) + print(prepared.source_table, file=out) + if prepared.refreshed: + print("kitaru snapshot written; re-runs read it until --refresh", file=out) + else: + print("kitaru snapshot reused (pass --refresh to refetch)", file=out) + return prepared.traces_path + + def _next_run_id(project_root: str | Path) -> str: """Sequential, sortable run ids — no clock, so runs stay reproducible.""" @@ -92,9 +149,11 @@ def command_init(args: argparse.Namespace, out: TextIO) -> int: def command_run(args: argparse.Namespace, out: TextIO) -> int: + run_id = None if args.estimate else (args.run_id or _next_run_id(args.project_root)) + traces = _resolve_traces_argument(args, out, run_id=run_id) if args.estimate: estimate = estimate_run( - args.traces, + traces, args.manifest, project_root=args.project_root, base_directory=args.base_directory, @@ -102,9 +161,9 @@ def command_run(args: argparse.Namespace, out: TextIO) -> int: print(estimate.render(), file=out) return 0 - run_id = args.run_id or _next_run_id(args.project_root) + assert run_id is not None result = run_pipeline( - args.traces, + traces, args.manifest, run_id=run_id, project_root=args.project_root, @@ -277,6 +336,20 @@ def command_apply(args: argparse.Namespace, out: TextIO) -> int: selected = _selected_indices(args, proposal, out) if selected is None: return 1 + from .state import contained_path + + template = contained_path(args.base_directory, proposal.template_file) + current = template.read_text(encoding="utf-8") + if selected: + about_to_write = candidate_prompt(current, proposal, selected) + from .verify import refuse_ungated_apply + + refuse_ungated_apply( + args.project_root, + run_id=run_id, + candidate_prompt_hash=text_hash(about_to_write), + force=bool(args.force), + ) result = apply_proposal( args.project_root, proposal, @@ -353,6 +426,79 @@ def command_status(args: argparse.Namespace, out: TextIO) -> int: return 0 +def command_verify(args: argparse.Namespace, out: TextIO) -> int: + """Replay a candidate against the frozen cohort that produced the run.""" + + backend_name = getattr(args, "backend", None) + if not backend_name: + print(NO_BACKEND_MESSAGE, file=out) + return 1 + if backend_name != "kitaru": + print( + f"unknown verification backend {backend_name!r}; only 'kitaru' is supported", + file=out, + ) + return 1 + + from .verify import ( + build_request, + format_verification_report, + load_run_source_payload, + run_verification, + ) + + run_id = args.run_id or latest_run_id(args.project_root) + if run_id is None: + print("no run to verify; run tracegrad run first", file=out) + return 1 + source = load_run_source_payload(args.project_root, run_id) + if source is None: + raise VerifyError( + f"run {run_id} has no Kitaru source metadata. " + "verify reuses the cohort the originating --source kitaru run persisted." + ) + proposal = load_proposal(args.project_root, run_id) + if is_stale(proposal, base_directory=args.base_directory): + mark_stale(args.project_root, run_id) + print( + f"{proposal.template_file} changed since run {run_id}; " + "the proposal is stale — re-run tracegrad", + file=out, + ) + return 1 + request = build_request( + project_root=args.project_root, + run_id=run_id, + proposal=proposal, + base_directory=args.base_directory, + source=source, + ) + from .integrations.kitaru.backend import KitaruVerificationBackend + + result = run_verification(args.project_root, request, KitaruVerificationBackend()) + from .integrations.kitaru.client import configured_server_url + + print( + format_verification_report( + result, request, server_url=configured_server_url() + ), + file=out, + ) + return verify_exit_code(result.status) + + +def verify_exit_code(status: str) -> int: + """Process exit for a finished verify. + + 0 only when the replay ``completed``. Partial, canceled, failed, and + incomplete stay non-zero so ``verify && apply`` cannot write from a + half-run. Apply after REVIEW is a separate hash+cohort / ``--force`` + decision. + """ + + return 0 if status == "completed" else 1 + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="tracegrad", @@ -381,7 +527,29 @@ def build_parser() -> argparse.ArgumentParser: ("propose", command_propose, "propose edits, reusing cached attributions"), ): run_parser = subparsers.add_parser(name, parents=[common], help=help_text) - run_parser.add_argument("--traces", required=True, help="JSONL trace export") + run_parser.add_argument("--traces", default=None, help="JSONL trace export") + run_parser.add_argument( + "--source", + choices=["kitaru"], + default=None, + help="optional trace source; mutually exclusive with --traces", + ) + run_parser.add_argument("--kitaru-cohort", default=None, help="Kitaru cohort name") + run_parser.add_argument( + "--kitaru-evaluation", + default=None, + help="Kitaru evaluation name to map onto judge.score", + ) + run_parser.add_argument( + "--kitaru-cohort-version", + default=None, + help="immutable cohort version id, display version, or number", + ) + run_parser.add_argument( + "--refresh", + action="store_true", + help="refetch the Kitaru cohort instead of reusing the snapshot", + ) run_parser.add_argument("--manifest", required=True, help="run manifest JSON") run_parser.add_argument("--run-id", default=None) run_parser.add_argument("--session-id", default=None) @@ -421,7 +589,7 @@ def build_parser() -> argparse.ArgumentParser: apply_parser.add_argument( "--force", action="store_true", - help="revert even though the template changed after it was applied", + help="override the verification gate, or revert even though the template changed", ) apply_parser.set_defaults(handler=command_apply) @@ -431,6 +599,19 @@ def build_parser() -> argparse.ArgumentParser: status_parser.add_argument("--manifest", default=None) status_parser.set_defaults(handler=command_status) + verify_parser = subparsers.add_parser( + "verify", + parents=[common], + help="replay-verify a candidate against a frozen Kitaru cohort", + ) + verify_parser.add_argument( + "--backend", + default=None, + help="verification backend (kitaru). Required; without one, verify exits non-zero", + ) + verify_parser.add_argument("--run", dest="run_id", default=None, help="tracegrad run id") + verify_parser.set_defaults(handler=command_verify) + return parser @@ -446,10 +627,12 @@ def main(argv: Sequence[str] | None = None, out: TextIO | None = None) -> int: DistillError, IngestError, InventoryError, + KitaruError, LLMError, PipelineError, StateError, SynthesisError, + VerifyError, ) as exc: print(f"tracegrad: {exc}", file=sys.stderr) return 1 diff --git a/src/tracegrad/config.py b/src/tracegrad/config.py index bd40b16..bbd524b 100644 --- a/src/tracegrad/config.py +++ b/src/tracegrad/config.py @@ -44,6 +44,15 @@ class HarnessPreset(BaseModel): enabled: StrictBool = True +class KitaruSettings(BaseModel): + """Non-secret Kitaru selection. Credentials stay in ``kitaru login``.""" + + model_config = ConfigDict(extra="forbid", populate_by_name=True) + + cohort: StrictStr | None = None + evaluation: StrictStr | None = None + + class TracegradConfig(BaseModel): """Validated project configuration loaded from `.tracegradrc`.""" @@ -58,6 +67,7 @@ class TracegradConfig(BaseModel): "synthesis": HarnessPreset(provider="claude"), } ) + kitaru: KitaruSettings = Field(default_factory=KitaruSettings) def _resolve_rc_path(path: str | Path) -> Path: diff --git a/src/tracegrad/integrations/__init__.py b/src/tracegrad/integrations/__init__.py new file mode 100644 index 0000000..f0bd02d --- /dev/null +++ b/src/tracegrad/integrations/__init__.py @@ -0,0 +1,6 @@ +"""Optional third-party integrations. + +Nothing in this package is imported by ``import tracegrad``. Kitaru SDK +imports live only under ``tracegrad.integrations.kitaru`` and are loaded +lazily when a Kitaru CLI path is used. +""" diff --git a/src/tracegrad/integrations/kitaru/__init__.py b/src/tracegrad/integrations/kitaru/__init__.py new file mode 100644 index 0000000..85954e7 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/__init__.py @@ -0,0 +1,43 @@ +"""Optional Kitaru integration. + +Importing this package does not import the Kitaru SDK. Mapping, scoring, and +graph helpers are usable in a core-only install; client/source/verify modules +call :func:`require_kitaru` before touching the SDK. +""" + +from __future__ import annotations + +from .errors import ( + INSTALL_MESSAGE, + KITARU_EXTRA, + KITARU_PIN, + NO_BACKEND_MESSAGE, + KitaruError, + KitaruNotInstalled, + KitaruSourceError, + KitaruVerifyError, +) +from .mapping import ( + MAPPING_VERSION, + MappedTrace, + SourceDrop, + map_session, +) +from .require import kitaru_available, require_kitaru + +__all__ = [ + "INSTALL_MESSAGE", + "KITARU_EXTRA", + "KITARU_PIN", + "MAPPING_VERSION", + "NO_BACKEND_MESSAGE", + "KitaruError", + "KitaruNotInstalled", + "KitaruSourceError", + "KitaruVerifyError", + "MappedTrace", + "SourceDrop", + "kitaru_available", + "map_session", + "require_kitaru", +] diff --git a/src/tracegrad/integrations/kitaru/accounting.py b/src/tracegrad/integrations/kitaru/accounting.py new file mode 100644 index 0000000..23e91ec --- /dev/null +++ b/src/tracegrad/integrations/kitaru/accounting.py @@ -0,0 +1,32 @@ +"""Source-drop vs batch-drop accounting. The two tables are never merged.""" + +from __future__ import annotations + +from collections import Counter +from typing import Iterable, Mapping + +from .mapping import SourceDrop + + +def format_source_table( + *, + sessions_selected: int, + traces_mapped: int, + dropped: Iterable[SourceDrop], + in_batch: int | None = None, + batch_drops: Mapping[str, int] | None = None, +) -> str: + """Render the two-table drop report from issue #8.""" + + reasons = Counter(drop.reason for drop in dropped) + lines = [ + f"Sessions selected: {sessions_selected:>6}", + f"Traces mapped: {traces_mapped:>6}", + ] + for reason, count in sorted(reasons.items()): + lines.append(f" {reason:<32} {count:>4}") + if in_batch is not None: + lines.append(f"In batch: {in_batch:>6}") + for reason, count in sorted((batch_drops or {}).items()): + lines.append(f" {reason:<32} {count:>4}") + return "\n".join(lines) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py new file mode 100644 index 0000000..637dd4c --- /dev/null +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -0,0 +1,563 @@ +"""Kitaru ``VerificationBackend`` (issue #9, ADR 0001 / 0006 / 0007). + +SDK imports stay in this module. The orchestrator in ``tracegrad.verify`` +never imports Kitaru. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Sequence +from typing import Any + +from tracegrad.verify import ( + DIVERGENCE_EVALUATOR_VERSION, + DIVERGENCE_HISTORY, + DIVERGENCE_SCOPE, + DIVERGENCE_SCORE, + DIVERGENCE_SELECT, + Divergence, + ReplayFailure, + SubmittedVerification, + VerificationRequest, + VerificationResult, + verification_fingerprint_for, +) + +from .client import FETCH_JOBS, KitaruGateway, run_async, worker_covers_agent_version +from .errors import KitaruVerifyError +from .graph import index_nodes, is_root_llm_node, llm_nodes, node_index, root_llm_nodes +from .mapping import extract_system_prompt +from .policy import asserts_no_passthrough, recorded_history_policy +from .require import require_kitaru +from .scores import evaluator_version_of, map_score, select_evaluation + + +def _history_tool_policy() -> Any: + """The only tool policy the tracegrad path can construct.""" + + require_kitaru() + spec = recorded_history_policy() + asserts_no_passthrough(spec) + from kitaru.api_models.v1.replay_config import ( + HistoryConfig, + HistoryScope, + ToolPolicy, + ToolPolicyOnMiss, + ) + + policy = ToolPolicy( + default=HistoryConfig( + scope=HistoryScope.COHORT_VERSION, + on_miss=ToolPolicyOnMiss.FAIL, + ) + ) + # Belt: never let a default or per-tool override sneak passthrough in. + default = policy.default + on_miss = getattr(default, "on_miss", None) + if str(getattr(on_miss, "value", on_miss)).lower() == "passthrough": + raise KitaruVerifyError("passthrough tool policy is not reachable through tracegrad") + if getattr(policy, "tools", None): + raise KitaruVerifyError("per-tool policy overrides are not reachable through tracegrad") + return policy + + +def _override(system_prompt: str) -> Any: + from kitaru.api_models.v1.replay_config import ReplayOverride + + return ReplayOverride(system_prompt=system_prompt) + + +def _evaluator_config(request: VerificationRequest) -> Any: + from kitaru.api_models.v1.replay_config import EvaluatorConfig + + return EvaluatorConfig( + evaluator=request.evaluator_name, + version=int(request.evaluator_version), + ) + + +def is_tool_history_miss(error: object | None) -> bool: + """Whether a replay error is a recorded-history miss (issue #9). + + Kitaru adapters raise ``ToolPolicyMissError`` with + ``No history result for tool '…'``. Invented needles never match those + messages, so the miss would otherwise land in ``replay_failures`` instead + of ``TOOL_HISTORY_MISS`` / incomparable. + """ + + if error is None: + return False + names = {type(error).__name__, type(error).__qualname__.rsplit(".", 1)[-1]} + if "ToolPolicyMissError" in names: + return True + text = str(error).strip() + if not text: + return False + lowered = text.lower() + return "no history result for tool" in lowered or "toolpolicymisserror" in lowered + + +def mixed_agent_version_message(counts: dict[str, int]) -> str: + parts = [f"{version}: {count} session(s)" for version, count in sorted(counts.items())] + return ( + "mixed-agent-version-cohort refused: every session in the cohort " + "version must share one agent_version_id " + f"({'; '.join(parts)}). A report that says the prompt caused this " + "when the agent code also moved is worse. See ADR 0007." + ) + + +def assert_override_scope( + baseline_nodes: list[Any] | tuple[Any, ...], + result_nodes: list[Any] | tuple[Any, ...], + candidate_prompt: str, +) -> str | None: + """Return a detail string when the override did not land on root nodes only. + + ADR 0006 requires every root LLM node to carry the candidate. An empty + result graph, or one with no root ``llm_call``, is + ``OVERRIDE_SCOPE_DIVERGENCE`` — not a vacuous pass. + """ + + roots = root_llm_nodes(result_nodes) + if not roots: + return "no root llm node carrying the candidate prompt" + base_by = index_nodes(baseline_nodes) + result_by = index_nodes(result_nodes) + for node in llm_nodes(result_nodes): + recorded = extract_system_prompt(node) + if is_root_llm_node(node, result_by): + if recorded != candidate_prompt: + return ( + f"root llm node {node_index(node)} did not carry the candidate prompt" + ) + continue + counterpart = base_by.get(node_index(node)) + expected = extract_system_prompt(counterpart) if counterpart is not None else None + if recorded != expected: + return ( + f"non-root llm node {node_index(node)} did not keep its baseline prompt" + ) + return None + + +def classify_replay_session( + *, + session_id: str, + number: int | None, + baseline_eval: Any, + candidate_eval: Any, + requested_evaluator_version: int, +) -> str | Divergence: + """Bucket one replay: improved/regressed/unchanged, or a typed divergence. + + ``select_evaluation`` drop reasons, evaluator-version mismatch, and an + unclassified score are incomparable — they go in ``diverged``, never + ``replay_failures``, and never vanish from the per-session buckets. + """ + + if isinstance(baseline_eval, str) or isinstance(candidate_eval, str): + parts: list[str] = [] + if isinstance(baseline_eval, str): + parts.append(f"baseline: {baseline_eval}") + if isinstance(candidate_eval, str): + parts.append(f"candidate: {candidate_eval}") + return Divergence( + session_id=session_id, + kind=DIVERGENCE_SELECT, + detail="; ".join(parts), + number=number, + ) + baseline_version = evaluator_version_of(baseline_eval) + candidate_version = evaluator_version_of(candidate_eval) + if ( + baseline_version != requested_evaluator_version + or candidate_version != requested_evaluator_version + ): + return Divergence( + session_id=session_id, + kind=DIVERGENCE_EVALUATOR_VERSION, + detail=( + f"requested evaluator_version {requested_evaluator_version}; " + f"baseline={baseline_version}; candidate={candidate_version}" + ), + number=number, + ) + verdict = classify_scores(baseline_eval, candidate_eval) + if verdict is None: + return Divergence( + session_id=session_id, + kind=DIVERGENCE_SCORE, + detail="scores could not be classified as improved, regressed, or unchanged", + number=number, + ) + return verdict + + +def classify_scores( + baseline: Any | None, candidate: Any | None +) -> str | None: + """Return improved / regressed / unchanged, or None if incomparable.""" + + if baseline is None or candidate is None: + return None + base_passed = getattr(baseline, "passed", None) + cand_passed = getattr(candidate, "passed", None) + if isinstance(base_passed, bool) and isinstance(cand_passed, bool): + if (not base_passed) and cand_passed: + return "improved" + if base_passed and (not cand_passed): + return "regressed" + if base_passed == cand_passed: + # Fall through to score for a finer signal when both exist. + pass + else: + return "unchanged" + base_score = map_score(baseline) + cand_score = map_score(candidate) + if isinstance(base_score, str) or isinstance(cand_score, str): + return None + if cand_score > base_score: + return "improved" + if cand_score < base_score: + return "regressed" + return "unchanged" + + +async def _await_verify(awaitable: Any, *, what: str) -> Any: + """Run one SDK awaitable; wrap raw failures as ``KitaruVerifyError``. + + Fail-closed: do not isolate as ReplayFailure, persist a result, or ungate + apply. Resume redoes collect. + """ + + try: + return await awaitable + except KitaruVerifyError: + raise + except Exception as exc: + raise KitaruVerifyError( + f"{what} failed ({type(exc).__name__}: {exc}). " + "Apply stays gated. Resume redoes collect. Pass --force on apply " + "if you must write anyway." + ) from exc + + +class KitaruVerificationBackend: + """Submit and collect a Kitaru experiment run for one candidate prompt.""" + + name = "kitaru" + + def __init__(self, gateway: KitaruGateway | None = None) -> None: + self._gateway = gateway + self._owns = gateway is None + if self._owns: + require_kitaru() + + def _gw(self) -> KitaruGateway: + if self._gateway is None: + self._gateway = KitaruGateway() + return self._gateway + + def _run_async(self, coro: Any) -> Any: + """Run one backend coroutine on a fresh event loop. + + ``asyncio.run`` closes the loop when it returns. Caching an + ``httpx.AsyncClient`` across preflight/submit/collect therefore + binds the next call to a dead loop. Owned gateways are closed and + dropped in the same coroutine that used them. + """ + + async def _isolated() -> Any: + try: + return await coro + finally: + if self._owns: + gateway = self._gateway + self._gateway = None + if gateway is not None: + await gateway.close() + + return run_async(_isolated()) + + def preflight(self, request: VerificationRequest) -> None: + self._run_async(self._preflight(request)) + + async def _preflight(self, request: VerificationRequest) -> None: + gateway = self._gw() + try: + await gateway.server_info() + except Exception as exc: + raise KitaruVerifyError( + "kitaru server is not reachable. Start the server and run " + "`kitaru login` first; tracegrad does not host verification " + "(ADR 0001)." + ) from exc + try: + await gateway.get_cohort_version(request.cohort_version_id) + except Exception as exc: + raise KitaruVerifyError( + f"cohort version {request.cohort_version_id} did not resolve" + ) from exc + try: + await gateway.get_agent_version(request.agent_version_id) + except Exception as exc: + raise KitaruVerifyError( + f"agent version {request.agent_version_id} did not resolve" + ) from exc + counts = request.agent_version_counts + distinct = {key: count for key, count in counts.items() if key != "unspecified"} + if "unspecified" in counts or len(distinct) != 1: + raise KitaruVerifyError(mixed_agent_version_message(counts or {"unspecified": 0})) + only = next(iter(distinct)) + if only != request.agent_version_id: + raise KitaruVerifyError(mixed_agent_version_message(counts)) + try: + workers = await gateway.list_live_workers() + except Exception as exc: + raise KitaruVerifyError( + "could not list workers from the kitaru server. Start the " + "server and run `kitaru login` first; tracegrad does not host " + "workers (ADR 0001)." + ) from exc + if not any(worker_covers_agent_version(worker, request.agent_version_id) for worker in workers): + raise KitaruVerifyError( + "no live worker is polling for agent version " + f"{request.agent_version_id}. Start a worker in the virtualenv " + "where the agent runs; tracegrad does not host workers (ADR 0001)." + ) + + def submit(self, request: VerificationRequest) -> SubmittedVerification: + return self._run_async(self._submit(request)) + + async def _submit(self, request: VerificationRequest) -> SubmittedVerification: + from kitaru.api_models.v1.experiment import ExperimentCreateRequest + from kitaru.api_models.v1.experiment_run import ExperimentRunCreateRequest + + gateway = self._gw() + digest = request.candidate_prompt_hash.removeprefix("sha256:")[:8] + experiment = await _await_verify( + gateway.create_experiment( + ExperimentCreateRequest( + name=f"tracegrad-{request.run_id}-{digest}", + description=f"tracegrad verification of {request.run_id}", + agent_id=uuid.UUID(request.agent_id), + override=_override(request.candidate_prompt), + tool_policy=_history_tool_policy(), + evaluators=[_evaluator_config(request)], + ) + ), + what="submit aborted: creating the experiment", + ) + run = await _await_verify( + gateway.start_run( + str(experiment.id), + ExperimentRunCreateRequest( + cohort_version_id=uuid.UUID(request.cohort_version_id), + agent_version_id=uuid.UUID(request.agent_version_id), + evaluate_baselines=True, + ), + ), + what="submit aborted: starting the experiment run", + ) + return SubmittedVerification( + experiment_id=str(experiment.id), + experiment_run_id=str(run.id), + ) + + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + return self._run_async(self._collect(request, submitted)) + + async def _collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + gateway = self._gw() + run = await _await_verify( + gateway.wait_for_experiment_run(submitted.experiment_run_id), + what="collect aborted: waiting for the experiment run", + ) + status = str(getattr(getattr(run, "status", None), "value", getattr(run, "status", "completed"))) + replays = await _await_verify( + gateway.list_replays(submitted.experiment_run_id), + what="collect aborted: listing replays", + ) + improved: list[str] = [] + regressed: list[str] = [] + unchanged: list[str] = [] + diverged: list[Divergence] = [] + failures: list[ReplayFailure] = [] + fetch_targets: list[Any] = [] + + for replay in replays: + session_id = str(replay.baseline_session_id) + number = request.session_numbers.get(session_id) + replay_status = str( + getattr(getattr(replay, "status", None), "value", getattr(replay, "status", "")) + ) + error = getattr(replay, "error", None) + if replay_status == "failed" or (error and not getattr(replay, "result_session_id", None)): + if is_tool_history_miss(error): + diverged.append( + Divergence( + session_id=session_id, + kind=DIVERGENCE_HISTORY, + detail=str(error or "tool history miss"), + number=number, + ) + ) + else: + failures.append( + ReplayFailure( + session_id=session_id, + error=str(error or replay_status or "replay failed"), + number=number, + ) + ) + continue + result_id = getattr(replay, "result_session_id", None) + if result_id is None: + failures.append( + ReplayFailure(session_id=session_id, error="missing result session", number=number) + ) + continue + fetch_targets.append(replay) + + payloads = await _fetch_replay_payloads(gateway, fetch_targets) + for replay, payload in zip(fetch_targets, payloads, strict=True): + session_id = str(replay.baseline_session_id) + number = request.session_numbers.get(session_id) + baseline_nodes, result_nodes, baseline_evals, candidate_evals = payload + scope = assert_override_scope( + baseline_nodes, result_nodes, request.candidate_prompt + ) + if scope: + diverged.append( + Divergence( + session_id=session_id, + kind=DIVERGENCE_SCOPE, + detail=scope, + number=number, + ) + ) + continue + outcome = classify_replay_session( + session_id=session_id, + number=number, + baseline_eval=select_evaluation(baseline_evals, request.evaluation_name), + candidate_eval=select_evaluation(candidate_evals, request.evaluation_name), + requested_evaluator_version=int(request.evaluator_version), + ) + if isinstance(outcome, Divergence): + diverged.append(outcome) + continue + if outcome == "improved": + improved.append(session_id) + elif outcome == "regressed": + regressed.append(session_id) + elif outcome == "unchanged": + unchanged.append(session_id) + + # Headline numbers from /api/v1/ui/experiment-runs/{id}/evaluation-aggregates + # so they match the Kitaru UI. That namespace is UI-support, not an obvious + # third-party contract; the <0.23 pin contains it. + aggregates = await _await_verify( + gateway.evaluation_aggregates(submitted.experiment_run_id), + what="collect aborted: fetching evaluation aggregates", + ) + baseline_stats, candidate_stats = _pick_aggregate(aggregates, request.evaluation_name) + run_status = "failed" if status == "failed" else ( + "partial" if failures or status != "completed" else "completed" + ) + result = VerificationResult( + status=run_status, # type: ignore[arg-type] + baseline_count=int((baseline_stats or {}).get("count") or 0), + candidate_count=int((candidate_stats or {}).get("count") or 0), + baseline_mean_score=_maybe_float((baseline_stats or {}).get("mean")), + candidate_mean_score=_maybe_float((candidate_stats or {}).get("mean")), + baseline_pass_rate=_maybe_float((baseline_stats or {}).get("pass_rate")), + candidate_pass_rate=_maybe_float((candidate_stats or {}).get("pass_rate")), + improved_sessions=improved, + regressed_sessions=regressed, + unchanged_sessions=unchanged, + diverged_sessions=diverged, + replay_failures=failures, + cohort_version_id=request.cohort_version_id, + agent_version_id=request.agent_version_id, + evaluator_version=str(request.evaluator_version), + baseline_prompt_hash=request.baseline_prompt_hash, + candidate_prompt_hash=request.candidate_prompt_hash, + verification_fingerprint=verification_fingerprint_for(request), + experiment_run_id=submitted.experiment_run_id, + ) + return result + + +_ReplayPayload = tuple[tuple[Any, ...], tuple[Any, ...], list[Any], list[Any]] + + +async def _fetch_replay_payloads( + gateway: KitaruGateway, + replays: Sequence[Any], + *, + jobs: int = FETCH_JOBS, +) -> list[_ReplayPayload]: + """Fetch baseline/result nodes and evaluations with bounded concurrency. + + One replay occupies one semaphore slot, matching ``KitaruGateway.fetch_records``. + The four HTTP calls for that replay run concurrently inside the slot. + + Fail-closed: one 404/timeout aborts the whole collect. Do not isolate as + ReplayFailure, persist a result, or ungate apply. Resume redoes collect. + """ + + if not replays: + return [] + semaphore = asyncio.Semaphore(max(1, jobs)) + + async def one(replay: Any) -> _ReplayPayload: + baseline_id = str(replay.baseline_session_id) + result_id = str(replay.result_session_id) + async with semaphore: + baseline_nodes, result_nodes, baseline_evals, candidate_evals = await asyncio.gather( + gateway.session_nodes(baseline_id), + gateway.session_nodes(result_id), + gateway.evaluations_for(baseline_id), + gateway.evaluations_for(result_id), + ) + return baseline_nodes, result_nodes, baseline_evals, candidate_evals + + # Fail-closed: gather raises on the first fetch error (no return_exceptions). + return await _await_verify( + asyncio.gather(*(one(replay) for replay in replays)), + what="collect aborted: fetching a replay payload", + ) + + +def _maybe_float(value: Any) -> float | None: + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return None + + +def _pick_aggregate( + payloads: list[Any], evaluation_name: str +) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + for item in payloads: + raw = item if isinstance(item, dict) else ( + item.model_dump(mode="json") if hasattr(item, "model_dump") else None + ) + if not isinstance(raw, dict): + continue + if raw.get("name") != evaluation_name: + continue + baseline = raw.get("baseline") + result = raw.get("result") + return ( + baseline if isinstance(baseline, dict) else None, + result if isinstance(result, dict) else None, + ) + return None, None diff --git a/src/tracegrad/integrations/kitaru/client.py b/src/tracegrad/integrations/kitaru/client.py new file mode 100644 index 0000000..d7efe8b --- /dev/null +++ b/src/tracegrad/integrations/kitaru/client.py @@ -0,0 +1,356 @@ +"""Kitaru SDK gateway. The only module that constructs ``KitaruAPIClient``. + +Imported only after :func:`require_kitaru`. Credentials and server URL come +from Kitaru's own config; tracegrad stores no secrets. +""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from .errors import KitaruSourceError, KitaruVerifyError +from .require import require_kitaru + +FETCH_JOBS = 8 +_EXPERIMENT_RUN_POLL_INTERVAL = 2.0 +_TERMINAL_EXPERIMENT_RUN_STATUSES = frozenset({"completed", "failed", "canceled"}) + + +@dataclass(frozen=True) +class CohortResolution: + """Immutable cohort version chosen once for a run (ADR 0004).""" + + cohort_id: str + cohort_name: str + cohort_version_id: str + display_version: str | None + version_number: int + agent_id: str + session_count: int + + +def configured_server_url() -> str: + """Kitaru's own configured server URL; empty when unset.""" + + try: + require_kitaru() + from kitaru.client.config import get_server_url + except Exception: + return "" + return get_server_url() or "" + + +def _uuid(value: str | uuid.UUID) -> uuid.UUID: + return value if isinstance(value, uuid.UUID) else uuid.UUID(str(value)) + + +class KitaruGateway: + """Thin async wrapper around ``KitaruAPIClient``.""" + + def __init__(self, client: Any | None = None) -> None: + require_kitaru() + if client is not None: + self._client = client + self._owns = False + return + from kitaru.client import KitaruAPIClient + + try: + self._client = KitaruAPIClient() + except RuntimeError as exc: + raise KitaruSourceError( + "no Kitaru server URL is configured. Run `kitaru login` " + "against your server; tracegrad stores no Kitaru secrets." + ) from exc + self._owns = True + + @property + def client(self) -> Any: + return self._client + + @property + def base_url(self) -> str: + url = getattr(self._client, "base_url", None) or getattr( + self._client, "_base_url", None + ) + return str(url) if url else "" + + async def close(self) -> None: + if self._owns: + await self._client.close() + + async def server_info(self) -> Any: + return await self._client.info.get() + + async def resolve_cohort( + self, name: str, version_ref: str | None = None + ) -> CohortResolution: + """Resolve a cohort name to one immutable version. Once, for the run.""" + + from kitaru.api_models.v1.cohort import CohortListParams + from kitaru.api_models.v1.filter import FilterCondition, FilterOp + + page = await self._client.cohorts.list( + CohortListParams( + filter=FilterCondition(field="name", op=FilterOp.EQ, value=name) + ) + ) + if not page.items: + raise KitaruSourceError(f"kitaru cohort {name!r} was not found") + cohort = page.items[0] + versions = await self._list_versions(cohort.id) + chosen = self._pick_version(versions, cohort.latest_version, version_ref) + if chosen is None: + raise KitaruSourceError( + f"kitaru cohort {name!r} has no version matching {version_ref!r}" + if version_ref + else f"kitaru cohort {name!r} has no versions" + ) + return CohortResolution( + cohort_id=str(cohort.id), + cohort_name=cohort.name, + cohort_version_id=str(chosen.id), + display_version=chosen.display_version, + version_number=int(chosen.version), + agent_id=str(cohort.agent_id), + session_count=int(chosen.session_count), + ) + + def _pick_version( + self, versions: list[Any], latest: int, version_ref: str | None + ) -> Any | None: + if not versions: + return None + if version_ref is None: + return max( + (item for item in versions if item.version == latest), + default=max(versions, key=lambda item: item.version), + key=lambda item: item.version, + ) + for item in versions: + if str(item.id) == version_ref: + return item + if item.display_version == version_ref: + return item + if str(item.version) == version_ref: + return item + return None + + async def _list_versions(self, cohort_id: uuid.UUID) -> list[Any]: + from kitaru.api_models.v1.cohort_version import CohortVersionListParams + + versions: list[Any] = [] + params = CohortVersionListParams() + while True: + page = await self._client.cohorts.list_versions(cohort_id, params) + versions.extend(page.items) + cursor = getattr(page, "next_cursor", None) + if not cursor: + return versions + params = CohortVersionListParams(cursor=cursor) + + async def list_sessions(self, cohort_version_id: str) -> list[Any]: + from kitaru.api_models.v1.filter import FilterCondition, FilterOp + from kitaru.api_models.v1.session import SessionListParams + + params = SessionListParams( + filter=FilterCondition( + field="cohort_version_id", + op=FilterOp.EQ, + value=cohort_version_id, + ), + size=100, + ) + sessions: list[Any] = [] + async for session in self._client.sessions.iter(params): + sessions.append(session) + return sessions + + async def session_bundle(self, session_id: str) -> tuple[Any, tuple[Any, ...], tuple[Any, ...]]: + session_uuid = _uuid(session_id) + full = await self._client.sessions.get_with_nodes(session_uuid) + evaluations = await self._evaluations_for(session_uuid) + return full.session, tuple(full.nodes), tuple(evaluations) + + async def _evaluations_for(self, session_id: uuid.UUID) -> list[Any]: + from kitaru.api_models.v1.evaluation import EvaluationListParams + from kitaru.api_models.v1.filter import FilterCondition, FilterOp + + params = EvaluationListParams( + filter=FilterCondition( + field="session_id", op=FilterOp.EQ, value=str(session_id) + ) + ) + return [item async for item in self._client.evaluations.iter(params)] + + async def fetch_records( + self, sessions: Sequence[Any], *, jobs: int = FETCH_JOBS + ) -> list[tuple[Any, tuple[Any, ...], tuple[Any, ...]]]: + semaphore = asyncio.Semaphore(max(1, jobs)) + + async def one(session: Any) -> tuple[Any, tuple[Any, ...], tuple[Any, ...]]: + async with semaphore: + return await self.session_bundle(str(session.id)) + + if not sessions: + return [] + # Fail-closed: one 404/timeout aborts the batch. Do not map a partial + # cohort. Wrap so the CLI prints KitaruSourceError instead of a traceback. + try: + return list(await asyncio.gather(*(one(session) for session in sessions))) + except KitaruSourceError: + raise + except Exception as exc: + raise KitaruSourceError( + "source fetch aborted: fetching a session payload failed " + f"({type(exc).__name__}: {exc}). The batch is not mapped." + ) from exc + + async def evaluator_id(self, name: str) -> str: + from kitaru.api_models.v1.evaluator import EvaluatorListParams + from kitaru.api_models.v1.filter import FilterCondition, FilterOp + + page = await self._client.evaluators.list( + EvaluatorListParams( + filter=FilterCondition(field="name", op=FilterOp.EQ, value=name) + ) + ) + if not page.items: + raise KitaruSourceError(f"kitaru evaluator {name!r} was not found") + return str(page.items[0].id) + + async def get_agent_version(self, agent_version_id: str) -> Any: + return await self._client.agent_versions.get(_uuid(agent_version_id)) + + async def get_cohort_version(self, cohort_version_id: str) -> Any: + return await self._client.cohort_versions.get(_uuid(cohort_version_id)) + + async def list_live_workers(self) -> list[Any]: + from kitaru.api_models.v1.worker import WorkerListParams + + # kitaru 0.22 WorkerListParams is FilterableListParams (extra='forbid'). + # include_stale is not a field; worker_covers_agent_version skips live=False. + return [item async for item in self._client.workers.iter(WorkerListParams())] + + async def create_experiment(self, request: Any) -> Any: + return await self._client.experiments.create(request) + + async def start_run(self, experiment_id: str, request: Any) -> Any: + return await self._client.experiments.start_run(_uuid(experiment_id), request) + + async def get_experiment_run(self, run_id: str) -> Any: + return await self._client.experiment_runs.get(_uuid(run_id)) + + async def wait_for_experiment_run(self, run_id: str, timeout: float | None = None) -> Any: + """Poll ``experiment_runs.get`` until completed, failed, or canceled. + + Same loop as kitaru CLI ``poll_run``. Do not wrap ``KitaruClient``: + its ``close()`` shuts the shared API client, and wait is not a method + on the 0.22 API client (close / context-manager only). + """ + + run_uuid = _uuid(run_id) + deadline = None if timeout is None else time.monotonic() + timeout + while True: + if deadline is None: + run = await self._client.experiment_runs.get(run_uuid) + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for experiment run {run_id}") + try: + run = await asyncio.wait_for( + self._client.experiment_runs.get(run_uuid), + timeout=remaining, + ) + except TimeoutError as exc: + raise TimeoutError( + f"Timed out waiting for experiment run {run_id}" + ) from exc + + status = getattr(run, "status", None) + if str(getattr(status, "value", status) or "") in _TERMINAL_EXPERIMENT_RUN_STATUSES: + return run + + if deadline is None: + await asyncio.sleep(_EXPERIMENT_RUN_POLL_INTERVAL) + continue + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for experiment run {run_id}") + try: + await asyncio.wait_for( + asyncio.sleep(min(_EXPERIMENT_RUN_POLL_INTERVAL, remaining)), + timeout=remaining, + ) + except TimeoutError as exc: + raise TimeoutError( + f"Timed out waiting for experiment run {run_id}" + ) from exc + + async def list_replays(self, experiment_run_id: str) -> list[Any]: + from kitaru.api_models.v1.filter import FilterCondition, FilterOp + from kitaru.api_models.v1.replay import ReplayListParams + + params = ReplayListParams( + filter=FilterCondition( + field="experiment_run_id", + op=FilterOp.EQ, + value=experiment_run_id, + ) + ) + return [item async for item in self._client.replays.iter(params)] + + async def evaluation_aggregates(self, experiment_run_id: str) -> list[Any]: + """Headline stats from the UI-support namespace. + + ``/api/v1/ui/`` is a UI-support namespace, not an obvious third-party + contract. This is the most likely thing to move under us; the + ``<0.23`` pin contains it. + """ + + response = await self._client.request( + "GET", + f"/api/v1/ui/experiment-runs/{experiment_run_id}/evaluation-aggregates", + ) + return response.json() + + async def session_nodes(self, session_id: str) -> tuple[Any, ...]: + full = await self._client.sessions.get_with_nodes(_uuid(session_id)) + return tuple(full.nodes) + + async def evaluations_for(self, session_id: str) -> list[Any]: + return await self._evaluations_for(_uuid(session_id)) + + +def worker_covers_agent_version(worker: Any, agent_version_id: str) -> bool: + """Whether a live worker claims this agent version.""" + + if getattr(worker, "live", False) is False: + return False + scope = getattr(worker, "scope", None) + claims = getattr(scope, "claims", None) or () + target = str(agent_version_id) + for claim in claims: + kind = getattr(getattr(claim, "kind", None), "value", getattr(claim, "kind", None)) + if str(kind) != "agent": + continue + claimed = getattr(claim, "agent_version_id", None) + if claimed is None or str(claimed) == target: + return True + return False + + +def run_async(coro: Any) -> Any: + """Run one coroutine from the sync CLI.""" + + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coro) + raise KitaruVerifyError("kitaru client cannot nest inside a running event loop") diff --git a/src/tracegrad/integrations/kitaru/errors.py b/src/tracegrad/integrations/kitaru/errors.py new file mode 100644 index 0000000..b11b4d1 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/errors.py @@ -0,0 +1,47 @@ +"""Named errors for the optional Kitaru integration. + +This module does not import the Kitaru SDK. Importing it must be safe in a +core-only install. +""" + +from __future__ import annotations + +KITARU_EXTRA = "tracegrad[kitaru]" +KITARU_PIN = "kitaru>=0.22,<0.23" + +INSTALL_MESSAGE = ( + "Kitaru support is an optional extra, not part of core tracegrad.\n" + f' install it with: uv tool install "{KITARU_EXTRA}"\n' + f" the extra pins {KITARU_PIN}\n" + "Then run `kitaru login` against your server. " + "tracegrad stores no Kitaru secrets.\n" + "Core commands still work without it: " + "tracegrad run --traces … / apply / trends." +) + +NO_BACKEND_MESSAGE = ( + "tracegrad verify needs a backend; without one nothing was verified.\n" + " install the extra: uv tool install \"tracegrad[kitaru]\"\n" + " then: tracegrad verify --backend kitaru --run \n" + "A verify that exits 0 having done nothing would read as verified in CI.\n" + "run, apply, and trends still work without a backend." +) + + +class KitaruError(ValueError): + """A Kitaru integration failure with an actionable message.""" + + +class KitaruNotInstalled(KitaruError): + """The extra is not installed.""" + + def __init__(self, message: str = INSTALL_MESSAGE) -> None: + super().__init__(message) + + +class KitaruSourceError(KitaruError): + """Fetching or mapping a Kitaru cohort failed.""" + + +class KitaruVerifyError(KitaruError): + """Replay verification could not start or complete.""" diff --git a/src/tracegrad/integrations/kitaru/graph.py b/src/tracegrad/integrations/kitaru/graph.py new file mode 100644 index 0000000..df7793a --- /dev/null +++ b/src/tracegrad/integrations/kitaru/graph.py @@ -0,0 +1,96 @@ +"""Root LLM node classification over Kitaru's session DAG. + +A root LLM node is an ``llm_call`` with no ``subagent_call`` anywhere in its +ancestry, following ``parent_index`` **and** ``secondary_parent_indexes``. A +node reachable from a subagent is not root even when one of its parents is. +See ADR 0005. + +This module does not import the Kitaru SDK; nodes are duck-typed. +""" + +from __future__ import annotations + +from typing import Any, Iterable, Sequence + +LLM_CALL = "llm_call" +SUBAGENT_CALL = "subagent_call" +TOOL_CALL = "tool_call" + + +def node_type_value(node: Any) -> str: + """Return the node type as a plain string, enum or not.""" + + raw = getattr(node, "node_type", "") + value = getattr(raw, "value", raw) + return "" if value is None else str(value) + + +def node_index(node: Any) -> int: + return int(node.index) + + +def parent_indexes(node: Any) -> tuple[int, ...]: + """Every parent of ``node``, primary first, then secondary, de-duplicated.""" + + parents: list[int] = [] + primary = getattr(node, "parent_index", None) + if primary is not None: + parents.append(int(primary)) + secondary = getattr(node, "secondary_parent_indexes", None) or () + for item in secondary: + index = int(item) + if index not in parents: + parents.append(index) + return tuple(parents) + + +def index_nodes(nodes: Sequence[Any]) -> dict[int, Any]: + return {node_index(node): node for node in nodes} + + +def has_subagent_ancestor(node: Any, by_index: dict[int, Any]) -> bool: + """Whether any ancestor of ``node`` is a ``subagent_call``.""" + + seen: set[int] = set() + stack = list(parent_indexes(node)) + while stack: + current = stack.pop() + if current in seen: + continue + seen.add(current) + ancestor = by_index.get(current) + if ancestor is None: + continue + if node_type_value(ancestor) == SUBAGENT_CALL: + return True + stack.extend(parent_indexes(ancestor)) + return False + + +def is_root_llm_node(node: Any, by_index: dict[int, Any] | None = None) -> bool: + """Whether ``node`` is an ``llm_call`` with no subagent in its ancestry.""" + + if node_type_value(node) != LLM_CALL: + return False + table = by_index if by_index is not None else index_nodes((node,)) + return not has_subagent_ancestor(node, table) + + +def root_llm_nodes(nodes: Iterable[Any]) -> tuple[Any, ...]: + """Root LLM nodes in session order (ascending index).""" + + materialised = tuple(nodes) + by_index = index_nodes(materialised) + roots = [node for node in materialised if is_root_llm_node(node, by_index)] + return tuple(sorted(roots, key=node_index)) + + +def llm_nodes(nodes: Iterable[Any]) -> tuple[Any, ...]: + """Every LLM node, root or not, in session order.""" + + return tuple( + sorted( + (node for node in nodes if node_type_value(node) == LLM_CALL), + key=node_index, + ) + ) diff --git a/src/tracegrad/integrations/kitaru/mapping.py b/src/tracegrad/integrations/kitaru/mapping.py new file mode 100644 index 0000000..0f1a3eb --- /dev/null +++ b/src/tracegrad/integrations/kitaru/mapping.py @@ -0,0 +1,270 @@ +"""Session → ``Trace`` mapping (ADR 0002, 0003, 0005). + +Kitaru sessions are duck-typed. This module does not import the Kitaru SDK, so +core-only tests can exercise every mapping rule with plain objects. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Sequence + +from tracegrad.canonical import text_hash +from tracegrad.schema import Trace, TraceMeta + +from .graph import llm_nodes, node_index, root_llm_nodes +from .pointer import resolve_text_selector +from .scores import ( + REASON_AMBIGUOUS_EVALUATION, + evaluator_name_of, + evaluator_version_id_of, + evaluator_version_of, + map_judge, + select_evaluation, +) + +MAPPING_VERSION = 1 + +REASON_SYSTEM_PROMPT_UNAVAILABLE = "system-prompt-unavailable" +REASON_MULTIPLE_SYSTEM_PROMPTS = "multiple-system-prompts" +REASON_OUTPUT_UNAVAILABLE = "output-unavailable" +REASON_INPUT_UNAVAILABLE = "input-unavailable" +REASON_ROOT_LLM_UNAVAILABLE = "root-llm-unavailable" +REASON_FORMAT_ENGINE_REFUSED = "format-engine-refused" +REASON_JUDGE_FINGERPRINT_CONFLICT = "judge-fingerprint-conflict" + +SOURCE_DROP_REASONS = ( + REASON_SYSTEM_PROMPT_UNAVAILABLE, + REASON_MULTIPLE_SYSTEM_PROMPTS, + REASON_OUTPUT_UNAVAILABLE, + REASON_INPUT_UNAVAILABLE, + REASON_ROOT_LLM_UNAVAILABLE, + REASON_AMBIGUOUS_EVALUATION, + "judge-rationale-missing", + "judge-score-out-of-range", + "judge-score-unsupported", + "judge-score-unavailable", +) + + +@dataclass(frozen=True) +class SourceDrop: + """A Session that could not become a Trace, with the kebab-case reason.""" + + session_id: str + reason: str + detail: str = "" + number: int | None = None + + +@dataclass(frozen=True) +class MappedTrace: + """A Trace plus the Kitaru metadata Phase 2 reuses.""" + + trace: Trace + session_number: int | None + evaluator_name: str + evaluator_version: int + evaluator_version_id: str | None + multi_turn: bool + system_prompt: str + + +def _session_id(session: Any) -> str: + value = getattr(session, "id", None) + if value is None: + raise TypeError("session is missing id") + return str(value) + + +def _session_number(session: Any) -> int | None: + value = getattr(session, "number", None) + return int(value) if isinstance(value, int) else None + + +def _extract_system_prompts(roots: Sequence[Any]) -> list[str] | str: + """Unique system prompts from root LLM nodes, or a drop reason.""" + + values: list[str] = [] + for node in roots: + prompt = resolve_text_selector( + getattr(node, "inputs", None), + getattr(node, "system_prompt_selector", None), + ) + if prompt is None or prompt == "": + continue + if prompt not in values: + values.append(prompt) + if not values: + return REASON_SYSTEM_PROMPT_UNAVAILABLE + if len(values) > 1: + return REASON_MULTIPLE_SYSTEM_PROMPTS + return values + + +def _root_model(roots: Sequence[Any]) -> str | None: + for node in roots: + model = getattr(node, "model", None) + if isinstance(model, str) and model: + return model + return None + + +def map_session( + session: Any, + nodes: Sequence[Any], + evaluations: Sequence[Any], + evaluation_name: str, +) -> MappedTrace | SourceDrop: + """Map one Kitaru session onto a ``Trace``, or drop it by name. + + Never guesses a system prompt. Never reads session-level inputs/outputs. + Never lets a tool output or a subagent prompt become the artifact. + """ + + session_id = _session_id(session) + number = _session_number(session) + roots = root_llm_nodes(nodes) + if not roots: + return SourceDrop( + session_id, + REASON_ROOT_LLM_UNAVAILABLE, + "no llm_call node is free of subagent ancestry", + number, + ) + + prompts = _extract_system_prompts(roots) + if isinstance(prompts, str): + detail = ( + "system_prompt_selector did not resolve to a string on any root LLM node" + if prompts == REASON_SYSTEM_PROMPT_UNAVAILABLE + else "root LLM nodes recorded more than one distinct system prompt" + ) + return SourceDrop(session_id, prompts, detail, number) + system_prompt = prompts[0] + + first, last = roots[0], roots[-1] + mapped_input = resolve_text_selector( + getattr(first, "inputs", None), + getattr(first, "input_text_selector", None), + ) + if mapped_input is None: + return SourceDrop( + session_id, + REASON_INPUT_UNAVAILABLE, + "input_text_selector did not resolve to a string on the first root LLM node", + number, + ) + mapped_output = resolve_text_selector( + getattr(last, "outputs", None), + getattr(last, "output_text_selector", None), + ) + if mapped_output is None: + return SourceDrop( + session_id, + REASON_OUTPUT_UNAVAILABLE, + "output_text_selector did not resolve to a string on the last root LLM node", + number, + ) + + selected = select_evaluation(list(evaluations), evaluation_name) + if isinstance(selected, str): + return SourceDrop(session_id, selected, f"evaluation {evaluation_name!r}", number) + + judge = map_judge(selected) + if isinstance(judge, str): + return SourceDrop(session_id, judge, f"evaluation {evaluation_name!r}", number) + + name = evaluator_name_of(selected) or evaluation_name + version = evaluator_version_of(selected) + if version is None: + return SourceDrop( + session_id, + REASON_AMBIGUOUS_EVALUATION, + "evaluation has no evaluator_version", + number, + ) + + model = _root_model(roots) + multi_turn = len(roots) > 1 or len(llm_nodes(nodes)) > 1 + trace = Trace( + trace_id=session_id, + input=mapped_input, + output=mapped_output, + judge=judge, + prompt_hash=text_hash(system_prompt), + meta=TraceMeta(model=model) if model is not None else None, + ) + return MappedTrace( + trace=trace, + session_number=number, + evaluator_name=name, + evaluator_version=version, + evaluator_version_id=evaluator_version_id_of(selected), + multi_turn=multi_turn, + system_prompt=system_prompt, + ) + + +def extract_system_prompt(node: Any) -> str | None: + """The recorded system prompt on one node, or ``None`` if unresolved.""" + + return resolve_text_selector( + getattr(node, "inputs", None), + getattr(node, "system_prompt_selector", None), + ) + + +def counterpart_prompt(nodes: Sequence[Any], index: int) -> str | None: + """System prompt of the LLM node at ``index``, if any.""" + + for node in nodes: + if node_index(node) == index: + return extract_system_prompt(node) + return None + + +@dataclass(frozen=True) +class BatchMapping: + """The mapped traces of one cohort, plus source drops, never mixed.""" + + mapped: tuple[MappedTrace, ...] + dropped: tuple[SourceDrop, ...] + evaluator_name: str | None + evaluator_version: int | None + evaluator_version_id: str | None + multi_turn_count: int + + +def map_batch( + records: Sequence[tuple[Any, Sequence[Any], Sequence[Any]]], + evaluation_name: str, +) -> BatchMapping | str: + """Map every session. Mixed evaluator versions refuse the batch. + + Returns a :class:`BatchMapping`, or ``ambiguous-evaluation`` when the + successfully mapped traces do not share one evaluator version. + """ + + mapped: list[MappedTrace] = [] + dropped: list[SourceDrop] = [] + for session, nodes, evaluations in records: + result = map_session(session, nodes, evaluations, evaluation_name) + if isinstance(result, SourceDrop): + dropped.append(result) + else: + mapped.append(result) + + versions = {(item.evaluator_name, item.evaluator_version) for item in mapped} + if len(versions) > 1: + return REASON_AMBIGUOUS_EVALUATION + + first = mapped[0] if mapped else None + return BatchMapping( + mapped=tuple(mapped), + dropped=tuple(dropped), + evaluator_name=first.evaluator_name if first else None, + evaluator_version=first.evaluator_version if first else None, + evaluator_version_id=first.evaluator_version_id if first else None, + multi_turn_count=sum(1 for item in mapped if item.multi_turn), + ) diff --git a/src/tracegrad/integrations/kitaru/pointer.py b/src/tracegrad/integrations/kitaru/pointer.py new file mode 100644 index 0000000..72d9d3b --- /dev/null +++ b/src/tracegrad/integrations/kitaru/pointer.py @@ -0,0 +1,66 @@ +"""RFC 6901 JSON Pointer resolution for Kitaru node selectors. + +Selectors are never guessed and never stringified as a fallback: a pointer that +does not land on a string is a resolution failure, and a missing selector is +the same. This module does not import the Kitaru SDK. +""" + +from __future__ import annotations + +from typing import Any + + +def unescape_token(token: str) -> str: + """Decode one JSON Pointer token (``~1`` → ``/``, ``~0`` → ``~``).""" + + return token.replace("~1", "/").replace("~0", "~") + + +def resolve_pointer(document: Any, pointer: str) -> Any: + """Resolve an RFC 6901 JSON Pointer. + + Raises: + ValueError: The pointer is malformed or does not exist. + """ + + if pointer == "": + return document + if not pointer.startswith("/"): + raise ValueError(f"JSON Pointer must be empty or start with '/': {pointer!r}") + + current = document + for raw_token in pointer[1:].split("/"): + token = unescape_token(raw_token) + if isinstance(current, list): + if token == "-": + raise ValueError("JSON Pointer '-' is not resolvable on a list") + if not token.isdigit() or (len(token) > 1 and token.startswith("0")): + raise ValueError(f"invalid array index {token!r}") + index = int(token) + if index >= len(current): + raise ValueError(f"array index {index} out of range") + current = current[index] + continue + if isinstance(current, dict): + if token not in current: + raise ValueError(f"key {token!r} not found") + current = current[token] + continue + raise ValueError(f"cannot traverse {type(current).__name__} with {token!r}") + return current + + +def resolve_text_selector(document: Any, selector: str | None) -> str | None: + """Resolve a selector onto a string, or ``None`` on any failure. + + Non-string landings are failures: tool payloads and nested objects must + not become ``Trace.input`` / ``Trace.output`` / a system prompt. + """ + + if selector is None: + return None + try: + value = resolve_pointer(document, selector) + except (TypeError, ValueError): + return None + return value if isinstance(value, str) else None diff --git a/src/tracegrad/integrations/kitaru/policy.py b/src/tracegrad/integrations/kitaru/policy.py new file mode 100644 index 0000000..358c241 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/policy.py @@ -0,0 +1,43 @@ +"""Tool-policy spec the Kitaru backend is allowed to emit. + +The values are the Kitaru wire strings. The only policy the tracegrad path +constructs is recorded history with ``on_miss=fail``. There is no parameter +that can select passthrough (ADR / issue #9 hard invariant). +""" + +from __future__ import annotations + +from typing import Final, Literal + +HistoryScopeName = Literal["cohort_version"] +OnMissName = Literal["fail"] + +HISTORY_SCOPE: Final[HistoryScopeName] = "cohort_version" +ON_MISS: Final[OnMissName] = "fail" +POLICY_TYPE: Final[str] = "history" + +# Wire form of HistoryConfig(scope=COHORT_VERSION, on_miss=FAIL). +RECORDED_HISTORY_POLICY: Final[dict[str, str]] = { + "type": POLICY_TYPE, + "scope": HISTORY_SCOPE, + "on_miss": ON_MISS, +} + + +def recorded_history_policy() -> dict[str, str]: + """The only tool policy tracegrad will attach to a replay.""" + + return dict(RECORDED_HISTORY_POLICY) + + +def asserts_no_passthrough(policy: dict[str, str]) -> None: + """Fail closed if a policy other than recorded-history/fail is supplied.""" + + if policy.get("type") != POLICY_TYPE: + raise ValueError("tracegrad refusals: only history tool policy is allowed") + if policy.get("scope") != HISTORY_SCOPE: + raise ValueError("tracegrad refusals: history scope must be cohort_version") + if policy.get("on_miss") != ON_MISS: + raise ValueError("tracegrad refusals: history on_miss must be fail") + if str(policy.get("on_miss", "")).lower() == "passthrough": + raise ValueError("passthrough tool policy is not reachable through tracegrad") diff --git a/src/tracegrad/integrations/kitaru/require.py b/src/tracegrad/integrations/kitaru/require.py new file mode 100644 index 0000000..474af19 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/require.py @@ -0,0 +1,24 @@ +"""Lazy availability check for the optional Kitaru extra.""" + +from __future__ import annotations + +from importlib import import_module + +from .errors import KitaruNotInstalled + + +def kitaru_available() -> bool: + """Whether the Kitaru SDK can be imported in this environment.""" + + try: + import_module("kitaru") + except ImportError: + return False + return True + + +def require_kitaru() -> None: + """Raise an actionable install message when the extra is missing.""" + + if not kitaru_available(): + raise KitaruNotInstalled() diff --git a/src/tracegrad/integrations/kitaru/scores.py b/src/tracegrad/integrations/kitaru/scores.py new file mode 100644 index 0000000..ed8a5b5 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/scores.py @@ -0,0 +1,134 @@ +"""Kitaru evaluation → tracegrad ``Judge`` mapping (ADR 0003). + +Scores are accepted, never rescaled. ``schema.Score`` is ``[0, 1]``; anything +outside that, or non-numeric, drops by name. This module does not import the +Kitaru SDK; evaluations are duck-typed. +""" + +from __future__ import annotations + +from typing import Any + +from tracegrad.schema import Judge + +REASON_RATIONALE_MISSING = "judge-rationale-missing" +REASON_SCORE_OUT_OF_RANGE = "judge-score-out-of-range" +REASON_SCORE_UNSUPPORTED = "judge-score-unsupported" +REASON_SCORE_UNAVAILABLE = "judge-score-unavailable" +REASON_AMBIGUOUS_EVALUATION = "ambiguous-evaluation" + +CATEGORICAL = "categorical" +STRING = "str" +BOOL = "bool" +FLOAT = "float" + + +def _data_type(evaluation: Any) -> str: + raw = getattr(evaluation, "data_type", None) + value = getattr(raw, "value", raw) + if value: + return str(value) + score = getattr(evaluation, "score", None) + label = getattr(evaluation, "value", None) + if score is not None and label is not None: + return CATEGORICAL + if score is None and label is not None: + return STRING + if isinstance(score, bool): + return BOOL + if score is not None: + return FLOAT + return "" + + +def _rationale(evaluation: Any) -> str | None: + explanation = getattr(evaluation, "explanation", None) + if not isinstance(explanation, str): + return None + stripped = explanation.strip() + return stripped or None + + +def map_score(evaluation: Any) -> float | str: + """Return a ``[0, 1]`` score, or the kebab-case drop reason.""" + + data_type = _data_type(evaluation) + if data_type in {STRING, CATEGORICAL}: + return REASON_SCORE_UNSUPPORTED + + score = getattr(evaluation, "score", None) + passed = getattr(evaluation, "passed", None) + + if isinstance(score, bool): + return 1.0 if score else 0.0 + if score is None and isinstance(passed, bool): + return 1.0 if passed else 0.0 + if isinstance(score, (int, float)): + value = float(score) + if 0.0 <= value <= 1.0: + return value + return REASON_SCORE_OUT_OF_RANGE + return REASON_SCORE_UNAVAILABLE + + +def map_judge(evaluation: Any) -> Judge | str: + """Map one evaluation onto a ``Judge``, or return a drop reason.""" + + rationale = _rationale(evaluation) + if rationale is None: + return REASON_RATIONALE_MISSING + mapped = map_score(evaluation) + if isinstance(mapped, str): + return mapped + return Judge(score=mapped, rationale=rationale) + + +def evaluator_name_of(evaluation: Any) -> str | None: + name = getattr(evaluation, "evaluator_name", None) + if isinstance(name, str) and name: + return name + fallback = getattr(evaluation, "name", None) + return fallback if isinstance(fallback, str) and fallback else None + + +def evaluator_version_of(evaluation: Any) -> int | None: + version = getattr(evaluation, "evaluator_version", None) + return int(version) if isinstance(version, int) else None + + +def evaluator_version_id_of(evaluation: Any) -> str | None: + value = getattr(evaluation, "evaluator_version_id", None) + return str(value) if value is not None else None + + +def select_evaluation( + evaluations: list[Any] | tuple[Any, ...], + evaluation_name: str, +) -> Any | str: + """Pick the evaluation named ``evaluation_name``, or a drop reason. + + Several rows with the same evaluator version are collapsed + deterministically (sorted by id). Differing versions on one session are + ``ambiguous-evaluation``. + """ + + matching = [ + item + for item in evaluations + if getattr(item, "name", None) == evaluation_name + ] + if not matching: + return REASON_SCORE_UNAVAILABLE + + versions = {evaluator_version_of(item) for item in matching} + if len(versions) > 1: + return REASON_AMBIGUOUS_EVALUATION + + matching.sort(key=lambda item: str(getattr(item, "id", ""))) + return matching[0] + + +def judge_fingerprint_for(evaluator_name: str, evaluator_version: int) -> str: + """The fingerprint derived from the evaluator that actually scored the batch.""" + + return f"{evaluator_name}@{evaluator_version}" diff --git a/src/tracegrad/integrations/kitaru/snapshot.py b/src/tracegrad/integrations/kitaru/snapshot.py new file mode 100644 index 0000000..4a615b3 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -0,0 +1,421 @@ +"""On-disk snapshot of a mapped Kitaru cohort (ADR 0004). + +The deterministic core never sees this module. Re-runs read the JSONL ingest +already knows; ``--refresh`` refetches. No Kitaru SDK import. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Mapping, Sequence + +from pydantic import BaseModel, ConfigDict, Field +from pydantic.types import StrictInt, StrictStr + +from tracegrad.state import ( + PathContainmentError, + StateLayout, + atomic_write, + atomic_write_json, + contained_path, + validate_run_id, +) + +from .mapping import MAPPING_VERSION, MappedTrace, SourceDrop + +SOURCE_KIND = "kitaru" +BATCH_FILENAME = "batch.jsonl" +FINGERPRINT_FILENAME = "fingerprint.json" +META_FILENAME = "meta.json" +DROPS_FILENAME = "source-drops.jsonl" +LATEST_POINTER = "latest.json" +RUN_SOURCE_FILENAME = "kitaru-source.json" + + +class SourceFingerprint(BaseModel): + """The fingerprint persisted next to the mapped batch.""" + + model_config = ConfigDict(extra="forbid") + + source: StrictStr = SOURCE_KIND + cohort_id: StrictStr + cohort_version_id: StrictStr + evaluation_name: StrictStr + evaluator_id: StrictStr + evaluator_version: StrictInt + agent_id: StrictStr + mapping_version: StrictInt = MAPPING_VERSION + + +class SourceMeta(BaseModel): + """Sidecar Phase 2 reuses; not part of the fingerprint contract.""" + + model_config = ConfigDict(extra="forbid") + + cohort_name: StrictStr + display_version: StrictStr | None = None + version_number: StrictInt | None = None + agent_version_id: StrictStr | None = None + agent_version_counts: dict[StrictStr, StrictInt] = Field(default_factory=dict) + session_numbers: dict[StrictStr, StrictInt] = Field(default_factory=dict) + system_prompts: dict[StrictStr, StrictStr] = Field(default_factory=dict) + multi_turn_count: StrictInt = 0 + sessions_selected: StrictInt = 0 + traces_mapped: StrictInt = 0 + evaluator_name: StrictStr | None = None + + +def kitaru_root(layout: StateLayout) -> Path: + return layout.sources / SOURCE_KIND + + +def _path_safe_component(value: str) -> str: + """A single directory name: no separators, no ``.`` / ``..``.""" + + text = str(value).strip() + if not text: + raise ValueError("snapshot path component must be non-empty") + collapsed = text.replace("\\", "/") + parts = [part.replace("..", "__") for part in collapsed.split("/") if part not in {"", ".", ".."}] + cleaned = "--".join(parts) + if not cleaned or Path(cleaned).name != cleaned: + raise ValueError(f"unsafe snapshot path component: {value!r}") + return cleaned + + +def snapshot_key(cohort_version_id: str, evaluation_name: str) -> str: + """On-disk identity: one frozen cohort version × one evaluation. + + Fingerprint identity includes ``evaluation_name``; keying only by + ``cohort_version_id`` would let a later ``--kitaru-evaluation`` clobber + the earlier JSONL. + """ + + return f"{_path_safe_component(cohort_version_id)}/{_path_safe_component(evaluation_name)}" + + +def snapshot_identity_key(fingerprint: SourceFingerprint) -> str: + return snapshot_key(fingerprint.cohort_version_id, fingerprint.evaluation_name) + + +def snapshot_dir(layout: StateLayout, snapshot_id: str) -> Path: + return contained_path(kitaru_root(layout), snapshot_id) + + +def _is_snapshot_dir(target: Path) -> bool: + return (target / BATCH_FILENAME).is_file() and (target / FINGERPRINT_FILENAME).is_file() + + +def snapshot_exists(layout: StateLayout, snapshot_id: str) -> bool: + try: + target = snapshot_dir(layout, snapshot_id) + except (ValueError, PathContainmentError, OSError): + return False + return _is_snapshot_dir(target) + + +def write_snapshot( + layout: StateLayout, + *, + fingerprint: SourceFingerprint, + meta: SourceMeta, + mapped: Sequence[MappedTrace], + dropped: Sequence[SourceDrop], +) -> Path: + """Write JSONL + fingerprint + meta under ``.tracegrad/sources/kitaru/``.""" + + snapshot_id = snapshot_identity_key(fingerprint) + target = snapshot_dir(layout, snapshot_id) + target.mkdir(parents=True, exist_ok=True) + lines = [ + json.dumps(item.trace.model_dump(mode="json"), ensure_ascii=False, separators=(",", ":")) + for item in mapped + ] + atomic_write(target / BATCH_FILENAME, "\n".join(lines) + ("\n" if lines else "")) + atomic_write_json(target / FINGERPRINT_FILENAME, fingerprint.model_dump(mode="json")) + atomic_write_json(target / META_FILENAME, meta.model_dump(mode="json")) + drop_lines = [ + json.dumps( + { + "session_id": drop.session_id, + "reason": drop.reason, + "detail": drop.detail, + "number": drop.number, + }, + ensure_ascii=False, + separators=(",", ":"), + ) + for drop in dropped + ] + atomic_write(target / DROPS_FILENAME, "\n".join(drop_lines) + ("\n" if drop_lines else "")) + _upsert_latest_pointer( + layout, + { + "cohort_name": meta.cohort_name, + "evaluation_name": fingerprint.evaluation_name, + "cohort_version_id": fingerprint.cohort_version_id, + "snapshot_id": snapshot_id, + "display_version": meta.display_version, + "version_number": meta.version_number, + "batch": str(target / BATCH_FILENAME), + }, + ) + return target + + +def load_fingerprint(layout: StateLayout, snapshot_id: str) -> SourceFingerprint: + path = snapshot_dir(layout, snapshot_id) / FINGERPRINT_FILENAME + return SourceFingerprint.model_validate_json(path.read_text(encoding="utf-8")) + + +def load_meta(layout: StateLayout, snapshot_id: str) -> SourceMeta: + path = snapshot_dir(layout, snapshot_id) / META_FILENAME + return SourceMeta.model_validate_json(path.read_text(encoding="utf-8")) + + +def load_source_drops(layout: StateLayout, snapshot_id: str) -> tuple[SourceDrop, ...]: + path = snapshot_dir(layout, snapshot_id) / DROPS_FILENAME + if not path.exists(): + return () + drops: list[SourceDrop] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + raw = json.loads(line) + drops.append( + SourceDrop( + session_id=str(raw["session_id"]), + reason=str(raw["reason"]), + detail=str(raw.get("detail") or ""), + number=raw.get("number") if isinstance(raw.get("number"), int) else None, + ) + ) + return tuple(drops) + + +def batch_path(layout: StateLayout, snapshot_id: str) -> Path: + return snapshot_dir(layout, snapshot_id) / BATCH_FILENAME + + +def persist_run_source( + layout: StateLayout, + run_id: str, + fingerprint: SourceFingerprint, + meta: SourceMeta, +) -> Path: + """Copy source identity into the run directory so verify can find it.""" + + target = layout.runs / validate_run_id(run_id) / RUN_SOURCE_FILENAME + atomic_write_json( + target, + { + "fingerprint": fingerprint.model_dump(mode="json"), + "meta": meta.model_dump(mode="json"), + }, + ) + return target + + +def fingerprints_compatible(stored: SourceFingerprint, requested: Mapping[str, Any]) -> bool: + """Whether a snapshot can be reused for this request without refetching.""" + + if stored.source != SOURCE_KIND: + return False + if stored.mapping_version != MAPPING_VERSION: + return False + if stored.evaluation_name != requested.get("evaluation_name"): + return False + requested_id = requested.get("cohort_version_id") + if requested_id is not None and str(stored.cohort_version_id) != str(requested_id): + return False + return True + + +def _read_pointer_file(layout: StateLayout) -> dict[str, Any] | None: + path = kitaru_root(layout) / LATEST_POINTER + if not path.is_file(): + return None + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if not isinstance(raw, dict): + return None + return raw + + +def _pointer_entries(raw: dict[str, Any]) -> list[dict[str, Any]]: + entries = raw.get("entries") + if isinstance(entries, list): + return [item for item in entries if isinstance(item, dict)] + if isinstance(raw.get("cohort_version_id"), str): + return [raw] + return [] + + +def _upsert_latest_pointer(layout: StateLayout, entry: Mapping[str, Any]) -> None: + raw = _read_pointer_file(layout) or {} + kept: list[dict[str, Any]] = [] + key = (entry.get("cohort_name"), entry.get("evaluation_name")) + for existing in _pointer_entries(raw): + existing_key = (existing.get("cohort_name"), existing.get("evaluation_name")) + if existing_key == key: + continue + kept.append(existing) + kept.append(dict(entry)) + atomic_write_json(kitaru_root(layout) / LATEST_POINTER, {"entries": kept}) + + +def load_latest_pointer( + layout: StateLayout, + *, + cohort_name: str, + evaluation_name: str, +) -> dict[str, Any] | None: + """The last-fetched snapshot for this cohort name + evaluation, if any.""" + + raw = _read_pointer_file(layout) + if raw is None: + return None + for entry in _pointer_entries(raw): + if entry.get("cohort_name") == cohort_name and entry.get("evaluation_name") == evaluation_name: + return entry + return None + + +def list_snapshot_ids(layout: StateLayout) -> list[str]: + """Relative ids under the Kitaru source root, including nested eval dirs. + + Legacy snapshots keyed only by ``cohort_version_id`` are still listed so + an older tree remains readable. + """ + + root = kitaru_root(layout) + if not root.is_dir(): + return [] + found: list[str] = [] + for child in sorted(root.iterdir()): + if not child.is_dir(): + continue + for grandchild in sorted(child.iterdir()): + rel = f"{child.name}/{grandchild.name}" + if grandchild.is_dir() and snapshot_exists(layout, rel): + found.append(rel) + if snapshot_exists(layout, child.name): + found.append(child.name) + return found + + +def snapshot_matches_request( + layout: StateLayout, + snapshot_id: str, + *, + cohort_name: str, + evaluation_name: str, + cohort_version: str | None = None, +) -> bool: + """Whether this on-disk snapshot is the one the request asked for.""" + + if not snapshot_exists(layout, snapshot_id): + return False + try: + fingerprint = load_fingerprint(layout, snapshot_id) + if not fingerprints_compatible(fingerprint, {"evaluation_name": evaluation_name}): + return False + meta = load_meta(layout, snapshot_id) + # Malformed JSON or a missing session_id/reason must not traceback + # in prepare_kitaru_source; treat the snapshot as corrupt and skip. + load_source_drops(layout, snapshot_id) + except (OSError, ValueError, KeyError, TypeError): + return False + if meta.cohort_name != cohort_name: + return False + if cohort_version is None: + return True + return ( + fingerprint.cohort_version_id == cohort_version + or meta.display_version == cohort_version + or snapshot_id == cohort_version + or ( + meta.version_number is not None + and str(meta.version_number) == cohort_version + ) + ) + + +def _snapshot_mtime(layout: StateLayout, snapshot_id: str) -> float: + path = snapshot_dir(layout, snapshot_id) / FINGERPRINT_FILENAME + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def _newest_matching(layout: StateLayout, snapshot_ids: Sequence[str]) -> str | None: + if not snapshot_ids: + return None + return max(snapshot_ids, key=lambda item: (_snapshot_mtime(layout, item), item)) + + +def _pointer_snapshot_id(entry: Mapping[str, Any], evaluation_name: str) -> str | None: + snapshot_id = entry.get("snapshot_id") + if isinstance(snapshot_id, str) and snapshot_id: + return snapshot_id + cohort_version_id = entry.get("cohort_version_id") + if not isinstance(cohort_version_id, str) or not cohort_version_id: + return None + return snapshot_key(cohort_version_id, evaluation_name) + + +def find_local_snapshot( + layout: StateLayout, + *, + cohort_name: str, + evaluation_name: str, + cohort_version: str | None = None, +) -> str | None: + """Return a local snapshot id that satisfies this request. + + Prefers the ``latest.json`` entry for ``(cohort_name, evaluation_name)`` + so a version-less re-run keeps the last fetched version of *this* pair + even after another cohort was sourced. When that pointer is missing, + picks the newest matching snapshot rather than the lexicographically + first directory. Does not contact the server. ``--refresh`` is the + caller's decision not to call this. + """ + + def matches(snapshot_id: str) -> bool: + return snapshot_matches_request( + layout, + snapshot_id, + cohort_name=cohort_name, + evaluation_name=evaluation_name, + cohort_version=cohort_version, + ) + + if cohort_version is not None: + candidates = [ + snapshot_key(cohort_version, evaluation_name), + _path_safe_component(cohort_version), + ] + seen: set[str] = set() + for candidate in candidates: + if candidate in seen: + continue + seen.add(candidate) + if matches(candidate): + return candidate + return _newest_matching(layout, [item for item in list_snapshot_ids(layout) if matches(item)]) + + pointer = load_latest_pointer( + layout, cohort_name=cohort_name, evaluation_name=evaluation_name + ) + if pointer is not None: + pointed = _pointer_snapshot_id(pointer, evaluation_name) + if pointed and matches(pointed): + return pointed + legacy = pointer.get("cohort_version_id") + if isinstance(legacy, str) and legacy and matches(legacy): + return legacy + return _newest_matching(layout, [item for item in list_snapshot_ids(layout) if matches(item)]) diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py new file mode 100644 index 0000000..3e7a46c --- /dev/null +++ b/src/tracegrad/integrations/kitaru/source.py @@ -0,0 +1,321 @@ +"""Fetch-and-map a frozen Kitaru cohort onto JSONL (ADR 0004, 0010). + +The deterministic core never imports this module. Callers write a snapshot +and pass the JSONL path to ``ingest_traces``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +from tracegrad.schema import Manifest, TemplateEngine +from tracegrad.state import StateLayout, initialize + +from .accounting import format_source_table +from .errors import KitaruSourceError +from .mapping import ( + REASON_AMBIGUOUS_EVALUATION, + REASON_FORMAT_ENGINE_REFUSED, + REASON_JUDGE_FINGERPRINT_CONFLICT, + SourceDrop, + map_batch, +) +from .require import require_kitaru +from .scores import judge_fingerprint_for +from .snapshot import ( + SourceFingerprint, + SourceMeta, + batch_path, + find_local_snapshot, + load_fingerprint, + load_meta, + load_source_drops, + persist_run_source, + snapshot_identity_key, + write_snapshot, +) + + +@dataclass(frozen=True) +class PreparedSource: + """What the CLI hands the existing pipeline.""" + + traces_path: Path + fingerprint: SourceFingerprint + meta: SourceMeta + source_table: str + refreshed: bool + + +def refuse_format_engine(manifest: Manifest) -> None: + if manifest.engine is TemplateEngine.FORMAT: + raise KitaruSourceError( + f"{REASON_FORMAT_ENGINE_REFUSED}: --source kitaru supports " + 'engine="none" only. A format template renders per request, so ' + "hashing recorded prompts would collapse the batch. See ADR 0002." + ) + + +def check_judge_fingerprint(manifest: Manifest, derived: str) -> None: + if manifest.judge_fingerprint != derived: + raise KitaruSourceError( + f"{REASON_JUDGE_FINGERPRINT_CONFLICT}: manifest judge_fingerprint " + f"{manifest.judge_fingerprint!r} does not match the evaluator that " + f"scored this cohort ({derived}). Drift detection only works if it " + "reads the thing that actually drifts. See ADR 0003." + ) + + +def _check_mapped_fingerprint( + manifest: Manifest, + *, + fingerprint: SourceFingerprint, + meta: SourceMeta, + evaluation_name: str, +) -> None: + if not meta.traces_mapped: + return + derived = judge_fingerprint_for( + meta.evaluator_name or evaluation_name, + fingerprint.evaluator_version, + ) + check_judge_fingerprint(manifest, derived) + + +def _agent_version_counts(sessions: list[Any]) -> dict[str, int]: + counts: dict[str, int] = {} + for session in sessions: + version = getattr(session, "agent_version_id", None) + key = str(version) if version is not None else "unspecified" + counts[key] = counts.get(key, 0) + 1 + return counts + + +def _single_agent_version(counts: dict[str, int]) -> str | None: + real = {key: count for key, count in counts.items() if key != "unspecified"} + if len(real) == 1: + return next(iter(real)) + return None + + +async def _await_source(awaitable: Any, *, what: str) -> Any: + """Run one SDK awaitable; wrap raw failures as ``KitaruSourceError``. + + Fail-closed: the batch is not mapped. Do not persist a partial snapshot. + """ + + try: + return await awaitable + except KitaruSourceError: + raise + except Exception as exc: + raise KitaruSourceError( + "source fetch aborted: " + f"{what} failed ({type(exc).__name__}: {exc}). The batch is not mapped." + ) from exc + + +async def _fetch_and_map( + *, + gateway: Any, + resolution: Any, + evaluation_name: str, +) -> tuple[Any, Any, Any, Any]: + sessions = await _await_source( + gateway.list_sessions(resolution.cohort_version_id), + what="listing sessions", + ) + records = await gateway.fetch_records(sessions) + mapping = map_batch(records, evaluation_name) + if isinstance(mapping, str) and mapping == REASON_AMBIGUOUS_EVALUATION: + versions: dict[str, int] = {} + for session, _nodes, evaluations in records: + from .scores import evaluator_version_of, select_evaluation + + selected = select_evaluation(list(evaluations), evaluation_name) + if isinstance(selected, str): + continue + version = evaluator_version_of(selected) + key = "unknown" if version is None else str(version) + versions[key] = versions.get(key, 0) + 1 + breakdown = ", ".join(f"v{ver}: {count}" for ver, count in sorted(versions.items())) + raise KitaruSourceError( + f"{REASON_AMBIGUOUS_EVALUATION}: evaluation {evaluation_name!r} " + f"resolves to more than one evaluator_version across the cohort " + f"({breakdown}). Refusing to mix. See ADR 0003." + ) + evaluator_id = await _await_source( + gateway.evaluator_id(mapping.evaluator_name or evaluation_name), + what="looking up the evaluator", + ) + fingerprint = SourceFingerprint( + source="kitaru", + cohort_id=resolution.cohort_id, + cohort_version_id=resolution.cohort_version_id, + evaluation_name=evaluation_name, + evaluator_id=evaluator_id, + evaluator_version=int(mapping.evaluator_version or 0), + agent_id=resolution.agent_id, + ) + if mapping.mapped and mapping.evaluator_version is None: + raise KitaruSourceError( + f"{REASON_AMBIGUOUS_EVALUATION}: mapped traces have no evaluator_version" + ) + if mapping.mapped: + fingerprint = fingerprint.model_copy( + update={"evaluator_version": int(mapping.evaluator_version)} + ) + counts = _agent_version_counts(sessions) + meta = SourceMeta( + cohort_name=resolution.cohort_name, + display_version=resolution.display_version, + version_number=int(resolution.version_number), + agent_version_id=_single_agent_version(counts), + agent_version_counts=counts, + session_numbers={ + item.trace.trace_id: item.session_number + for item in mapping.mapped + if item.session_number is not None + }, + system_prompts={item.trace.trace_id: item.system_prompt for item in mapping.mapped}, + multi_turn_count=mapping.multi_turn_count, + sessions_selected=len(sessions), + traces_mapped=len(mapping.mapped), + evaluator_name=mapping.evaluator_name or evaluation_name, + ) + return fingerprint, meta, mapping.mapped, mapping.dropped + + +def _assembled_source( + layout: StateLayout, + *, + snapshot_id: str, + fingerprint: SourceFingerprint, + meta: SourceMeta, + dropped: Sequence[SourceDrop], + refreshed: bool, + manifest: Manifest, + evaluation_name: str, + run_id: str | None, +) -> PreparedSource: + _check_mapped_fingerprint( + manifest, + fingerprint=fingerprint, + meta=meta, + evaluation_name=evaluation_name, + ) + if run_id is not None: + persist_run_source(layout, run_id, fingerprint, meta) + table = format_source_table( + sessions_selected=meta.sessions_selected, + traces_mapped=meta.traces_mapped, + dropped=dropped, + ) + if meta.multi_turn_count: + table += ( + f"\n{meta.multi_turn_count} multi-turn session(s) collapsed to " + "first-root input and last-root output" + ) + return PreparedSource( + traces_path=batch_path(layout, snapshot_id), + fingerprint=fingerprint, + meta=meta, + source_table=table, + refreshed=refreshed, + ) + + +def prepare_kitaru_source( + *, + project_root: str | Path, + manifest: Manifest, + cohort_name: str, + evaluation_name: str, + cohort_version: str | None = None, + refresh: bool = False, + gateway: Any | None = None, + run_id: str | None = None, +) -> PreparedSource: + """Resolve, snapshot, and return the JSONL path ingest already reads. + + Re-runs read the local snapshot (including ``latest.json``) unless + ``--refresh``. The gateway is constructed only when a fetch is needed so + a pinned snapshot stays reproducible with the server unreachable + (ADR 0004 / issue #8). + """ + + refuse_format_engine(manifest) + layout = initialize(project_root) + + if not refresh: + local_id = find_local_snapshot( + layout, + cohort_name=cohort_name, + evaluation_name=evaluation_name, + cohort_version=cohort_version, + ) + if local_id is not None: + fingerprint = load_fingerprint(layout, local_id) + meta = load_meta(layout, local_id) + dropped = load_source_drops(layout, local_id) + return _assembled_source( + layout, + snapshot_id=local_id, + fingerprint=fingerprint, + meta=meta, + dropped=dropped, + refreshed=False, + manifest=manifest, + evaluation_name=evaluation_name, + run_id=run_id, + ) + + from .client import run_async + + owns = gateway is None + if owns: + require_kitaru() + from .client import KitaruGateway + + gateway = KitaruGateway() + + async def _run() -> PreparedSource: + try: + resolution = await _await_source( + gateway.resolve_cohort(cohort_name, cohort_version), + what="resolving the cohort", + ) + fingerprint, meta, mapped, dropped = await _fetch_and_map( + gateway=gateway, + resolution=resolution, + evaluation_name=evaluation_name, + ) + # Check before write_snapshot so a conflicting --refresh cannot + # clobber the last-good snapshot or latest.json pointer. + _check_mapped_fingerprint( + manifest, + fingerprint=fingerprint, + meta=meta, + evaluation_name=evaluation_name, + ) + write_snapshot( + layout, fingerprint=fingerprint, meta=meta, mapped=mapped, dropped=dropped + ) + return _assembled_source( + layout, + snapshot_id=snapshot_identity_key(fingerprint), + fingerprint=fingerprint, + meta=meta, + dropped=dropped, + refreshed=True, + manifest=manifest, + evaluation_name=evaluation_name, + run_id=run_id, + ) + finally: + if owns: + await gateway.close() + + return run_async(_run()) diff --git a/src/tracegrad/ports.py b/src/tracegrad/ports.py index 6087bf6..79bf733 100644 --- a/src/tracegrad/ports.py +++ b/src/tracegrad/ports.py @@ -6,6 +6,10 @@ deterministic side, and the implementations live in ``llm``. Nothing in this module can talk to anything. That is the point. + +``VerificationBackend`` is the same idea for Phase 2 (ADR 0010): the +orchestrator holds a backend without becoming backend-aware. The Kitaru +implementation lives under ``integrations/kitaru/``. """ from __future__ import annotations @@ -28,3 +32,16 @@ def complete( schema: Mapping[str, Any] | None = None, timeout: float | None = None, ) -> Any: ... + + +@runtime_checkable +class VerificationBackend(Protocol): + """Replay-verify a candidate without the orchestrator knowing the vendor.""" + + name: str + + def preflight(self, request: Any) -> None: ... + + def submit(self, request: Any) -> Any: ... + + def collect(self, request: Any, submitted: Any) -> Any: ... diff --git a/src/tracegrad/state.py b/src/tracegrad/state.py index d149238..2f7f76a 100644 --- a/src/tracegrad/state.py +++ b/src/tracegrad/state.py @@ -8,6 +8,8 @@ * ``reports/`` — persisted reports; * ``runs//resume.json`` — per-run interruption checkpoints; * ``snapshots/`` — pre-apply prompt snapshots; +* ``sources/`` — optional mapped-trace snapshots (JSONL the pipeline already reads); +* ``verification/`` — persisted replay-verification state; * ``.gitignore`` — keeps generated state out of version control. """ @@ -45,6 +47,8 @@ class StateLayout: reports: Path snapshots: Path runs: Path + sources: Path + verification: Path gitignore: Path @property @@ -53,7 +57,14 @@ def resume_directory(self) -> Path: @property def data_directories(self) -> tuple[Path, ...]: - return (self.distilled, self.ledgers, self.reports, self.snapshots) + return ( + self.distilled, + self.ledgers, + self.reports, + self.snapshots, + self.sources, + self.verification, + ) def resume_path(self, run_id: str) -> Path: return self.runs / validate_run_id(run_id) / "resume.json" @@ -80,6 +91,8 @@ def _layout(project_root: str | Path | StateLayout) -> StateLayout: reports=state_root / "reports", snapshots=state_root / "snapshots", runs=state_root / "runs", + sources=state_root / "sources", + verification=state_root / "verification", gitignore=state_root / ".gitignore", ) diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py new file mode 100644 index 0000000..b10eb7b --- /dev/null +++ b/src/tracegrad/verify.py @@ -0,0 +1,549 @@ +"""Replay verification orchestrator. + +Holds a :class:`~tracegrad.ports.VerificationBackend` without becoming +backend-aware (ADR 0010). Persist / resume lives here so an interrupted +verify never duplicates the experiment. Apply-gating on +``candidate_prompt_hash`` plus the originating cohort lives here so the core +of ``apply`` stays a writer, not a Kitaru client. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field +from pydantic.types import StrictFloat, StrictInt, StrictStr + +from .apply import Proposal, StaleProposalError, candidate_prompt, is_stale, load_proposal +from .canonical import content_hash, text_hash +from .ports import VerificationBackend +from .state import ( + StateLayout, + atomic_write_json, + contained_path, + initialize, + validate_run_id, +) + +RUN_SOURCE_FILENAME = "kitaru-source.json" +_FINGERPRINT_REQUIRED = ( + "cohort_id", + "cohort_version_id", + "evaluation_name", + "evaluator_id", + "evaluator_version", + "agent_id", +) + +DIVERGENCE_HISTORY = "TOOL_HISTORY_MISS" +DIVERGENCE_SCOPE = "OVERRIDE_SCOPE_DIVERGENCE" +DIVERGENCE_EVALUATOR_VERSION = "EVALUATOR_VERSION_MISMATCH" +DIVERGENCE_SELECT = "SELECT_EVALUATION_FAILED" +DIVERGENCE_SCORE = "SCORE_UNCLASSIFIED" +VERIFICATION_FILENAME = "state.json" + +DivergenceKind = Literal[ + "TOOL_HISTORY_MISS", + "OVERRIDE_SCOPE_DIVERGENCE", + "EVALUATOR_VERSION_MISMATCH", + "SELECT_EVALUATION_FAILED", + "SCORE_UNCLASSIFIED", +] + + +class VerifyError(ValueError): + """Verification cannot start or cannot gate apply.""" + + +class Divergence(BaseModel): + model_config = ConfigDict(extra="forbid") + + session_id: StrictStr + kind: DivergenceKind + detail: StrictStr = "" + number: StrictInt | None = None + + +class ReplayFailure(BaseModel): + model_config = ConfigDict(extra="forbid") + + session_id: StrictStr + error: StrictStr + number: StrictInt | None = None + + +class VerificationResult(BaseModel): + model_config = ConfigDict(extra="forbid") + + status: Literal["completed", "partial", "failed"] + baseline_count: StrictInt = 0 + candidate_count: StrictInt = 0 + baseline_mean_score: StrictFloat | None = None + candidate_mean_score: StrictFloat | None = None + baseline_pass_rate: StrictFloat | None = None + candidate_pass_rate: StrictFloat | None = None + improved_sessions: list[StrictStr] = Field(default_factory=list) + regressed_sessions: list[StrictStr] = Field(default_factory=list) + unchanged_sessions: list[StrictStr] = Field(default_factory=list) + diverged_sessions: list[Divergence] = Field(default_factory=list) + replay_failures: list[ReplayFailure] = Field(default_factory=list) + cohort_version_id: StrictStr + agent_version_id: StrictStr + evaluator_version: StrictStr + baseline_prompt_hash: StrictStr + candidate_prompt_hash: StrictStr + verification_fingerprint: StrictStr + experiment_run_id: StrictStr + + +class SubmittedVerification(BaseModel): + model_config = ConfigDict(extra="forbid") + + experiment_id: StrictStr + experiment_run_id: StrictStr + + +class VerificationRequest(BaseModel): + """What a backend needs, all sourced from the originating run's snapshot.""" + + model_config = ConfigDict(extra="forbid") + + run_id: StrictStr + proposal_id: StrictStr + candidate_prompt: StrictStr + candidate_prompt_hash: StrictStr + baseline_prompt_hash: StrictStr + cohort_id: StrictStr + cohort_version_id: StrictStr + cohort_name: StrictStr + display_version: StrictStr | None = None + evaluation_name: StrictStr + evaluator_id: StrictStr + evaluator_version: StrictInt + evaluator_name: StrictStr + agent_id: StrictStr + agent_version_id: StrictStr + agent_version_counts: dict[StrictStr, StrictInt] = Field(default_factory=dict) + system_prompts: dict[StrictStr, StrictStr] = Field(default_factory=dict) + session_numbers: dict[StrictStr, StrictInt] = Field(default_factory=dict) + + +class VerificationState(BaseModel): + """Persisted verification record under ``.tracegrad/verification/``.""" + + model_config = ConfigDict(extra="forbid") + + verification_id: StrictStr + run_id: StrictStr + proposal_id: StrictStr + experiment_id: StrictStr | None = None + experiment_run_id: StrictStr | None = None + cohort_version_id: StrictStr + evaluator_version: StrictStr + agent_version_id: StrictStr + baseline_prompt_hash: StrictStr + candidate_prompt_hash: StrictStr + tool_policy: dict[str, str] = Field(default_factory=dict) + per_session: dict[str, str] = Field(default_factory=dict) + result: VerificationResult | None = None + verification_fingerprint: StrictStr | None = None + + +def verification_id_for(run_id: str, candidate_prompt_hash: str) -> str: + digest = candidate_prompt_hash.removeprefix("sha256:")[:12] + return f"verify-{validate_run_id(run_id)}-{digest}" + + +def verification_dir(layout: StateLayout, verification_id: str) -> Path: + return layout.verification / verification_id + + +def verification_path(layout: StateLayout, verification_id: str) -> Path: + return verification_dir(layout, verification_id) / VERIFICATION_FILENAME + + +def save_verification_state(layout: StateLayout, state: VerificationState) -> Path: + target = verification_path(layout, state.verification_id) + atomic_write_json(target, state.model_dump(mode="json")) + return target + + +def load_verification_state(layout: StateLayout, verification_id: str) -> VerificationState | None: + """Read one verification record, or None if it is missing or unreadable. + + Corrupt ``state.json`` (including pydantic ``ValidationError``) is treated + as missing so ``run_verification`` can resubmit. Apply stays gated: + ``list_verification_states`` already skips those files. + """ + + target = verification_path(layout, verification_id) + if not target.exists(): + return None + try: + return VerificationState.model_validate_json(target.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def list_verification_states(layout: StateLayout) -> list[VerificationState]: + records: list[VerificationState] = [] + for path in sorted(layout.verification.glob(f"*/{VERIFICATION_FILENAME}")): + try: + records.append(VerificationState.model_validate_json(path.read_text(encoding="utf-8"))) + except (OSError, ValueError): + continue + return records + + +def matching_verification( + project_root: str | Path, + candidate_prompt_hash: str, + *, + cohort_version_id: str, +) -> VerificationState | None: + """A persisted verification of this candidate on this cohort, if any. + + Requires a stored ``result``. Hash-only is not enough: a verify against + a previous cohort must not ungate apply on a new one. Apply's ``run_id`` + is not the match key. A finished REVIEW/FAILED report still matches; + status is not required to be ``completed``. + """ + + if not candidate_prompt_hash or not cohort_version_id: + return None + layout = initialize(project_root) + for state in list_verification_states(layout): + if ( + state.candidate_prompt_hash == candidate_prompt_hash + and state.cohort_version_id == cohort_version_id + and state.experiment_run_id + and state.result is not None + ): + return state + return None + + +def _source_cohort_version_id(project_root: str | Path, run_id: str) -> str | None: + """Cohort version from the originating run's sidecar, if usable.""" + + payload = load_run_source_payload(project_root, run_id) + if payload is None: + return None + fingerprint = payload.get("fingerprint") + if not isinstance(fingerprint, dict): + return None + value = fingerprint.get("cohort_version_id") + if value is None or str(value) == "": + return None + return str(value) + + +def load_run_source_payload(project_root: str | Path, run_id: str) -> dict[str, Any] | None: + """Read the originating-run source sidecar, if the run had a backend.""" + + layout = initialize(project_root) + target = layout.runs / validate_run_id(run_id) / RUN_SOURCE_FILENAME + if not target.exists(): + return None + import json + + try: + payload = json.loads(target.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return _usable_run_source_payload(payload) + + +def _usable_run_source_payload(payload: object) -> dict[str, Any] | None: + """A sidecar verify can consume without KeyError/TypeError in build_request.""" + + if not isinstance(payload, dict): + return None + fingerprint = payload.get("fingerprint") + meta = payload.get("meta") + if not isinstance(fingerprint, dict) or not isinstance(meta, dict): + return None + for key in _FINGERPRINT_REQUIRED: + if key not in fingerprint or fingerprint[key] is None: + return None + return payload + + +def backend_is_configured(project_root: str | Path, run_id: str) -> bool: + """Whether this run originated from a Kitaru source (ADR 0009). + + Presence of the sidecar is enough. A corrupt file must not ungate apply. + """ + + layout = initialize(project_root) + target = layout.runs / validate_run_id(run_id) / RUN_SOURCE_FILENAME + return target.is_file() + + +def refuse_ungated_apply( + project_root: str | Path, + *, + run_id: str, + candidate_prompt_hash: str, + force: bool, +) -> None: + """Refuse apply when a backend is configured and no matching verify exists. + + Match is ``candidate_prompt_hash`` plus the sidecar's ``cohort_version_id``. + ``run_id`` loads that sidecar; it is not the match key. A new cohort + needs re-verify even when the candidate text is unchanged. A selected + subset cannot be ungated by re-running verify; the operator must + ``apply --all`` or ``--force`` (ADR 0009). + """ + + if force or not backend_is_configured(project_root, run_id): + return + cohort_version_id = _source_cohort_version_id(project_root, run_id) + if ( + cohort_version_id is None + or matching_verification( + project_root, + candidate_prompt_hash, + cohort_version_id=cohort_version_id, + ) + is None + ): + raise VerifyError( + "apply is gated on a hash-matching verification for this candidate " + "on this cohort. Verify always hashes the full proposal. A subset " + "of a verified proposal (interactive or --accept) needs --force or " + "`tracegrad apply --all`; re-running verify cannot ungate it. " + "A new cohort needs `tracegrad verify` even when the candidate " + "text is unchanged. See ADR 0009." + ) + + +def build_request( + *, + project_root: str | Path, + run_id: str, + proposal: Proposal, + base_directory: str | Path = ".", + source: dict[str, Any], +) -> VerificationRequest: + if is_stale(proposal, base_directory=base_directory): + raise StaleProposalError( + f"{proposal.template_file} changed since run {run_id}; " + "the proposal is stale — re-run tracegrad" + ) + template = contained_path(base_directory, proposal.template_file) + try: + current = template.read_text(encoding="utf-8") + except OSError as exc: + raise VerifyError(f"could not read template {template}: {exc}") from exc + candidate = candidate_prompt( + current, proposal, range(len(proposal.edits)) + ) + fingerprint = source.get("fingerprint") if isinstance(source, dict) else None + meta = source.get("meta") if isinstance(source, dict) else None + if not isinstance(fingerprint, dict) or not isinstance(meta, dict): + raise VerifyError( + f"run {run_id} source metadata is not a usable Kitaru sidecar" + ) + try: + agent_version_id = meta.get("agent_version_id") + if not agent_version_id: + raise VerifyError( + "verification requires a single agent_version_id on the originating " + "cohort; this run's source metadata does not have one. See ADR 0007." + ) + return VerificationRequest( + run_id=run_id, + proposal_id=proposal.run_id, + candidate_prompt=candidate, + candidate_prompt_hash=text_hash(candidate), + baseline_prompt_hash=proposal.base_prompt_hash, + cohort_id=str(fingerprint["cohort_id"]), + cohort_version_id=str(fingerprint["cohort_version_id"]), + cohort_name=str(meta.get("cohort_name") or ""), + display_version=meta.get("display_version"), + evaluation_name=str(fingerprint["evaluation_name"]), + evaluator_id=str(fingerprint["evaluator_id"]), + evaluator_version=int(fingerprint["evaluator_version"]), + evaluator_name=str(meta.get("evaluator_name") or fingerprint["evaluation_name"]), + agent_id=str(fingerprint["agent_id"]), + agent_version_id=str(agent_version_id), + agent_version_counts=dict(meta.get("agent_version_counts") or {}), + system_prompts=dict(meta.get("system_prompts") or {}), + session_numbers={ + key: int(value) for key, value in (meta.get("session_numbers") or {}).items() + }, + ) + except VerifyError: + raise + except (KeyError, TypeError, ValueError) as exc: + raise VerifyError( + f"run {run_id} source metadata is not a usable Kitaru sidecar" + ) from exc + + +def verification_fingerprint_for(request: VerificationRequest) -> str: + return content_hash( + { + "cohort_version_id": request.cohort_version_id, + "agent_version_id": request.agent_version_id, + "evaluator_version": request.evaluator_version, + "evaluation_name": request.evaluation_name, + "baseline_prompt_hash": request.baseline_prompt_hash, + "candidate_prompt_hash": request.candidate_prompt_hash, + "tool_policy": { + "type": "history", + "scope": "cohort_version", + "on_miss": "fail", + }, + } + ) + + +def run_verification( + project_root: str | Path, + request: VerificationRequest, + backend: VerificationBackend, +) -> VerificationResult: + """Submit (or resume) verification and persist the result.""" + + layout = initialize(project_root) + vid = verification_id_for(request.run_id, request.candidate_prompt_hash) + state = load_verification_state(layout, vid) or VerificationState( + verification_id=vid, + run_id=request.run_id, + proposal_id=request.proposal_id, + cohort_version_id=request.cohort_version_id, + evaluator_version=str(request.evaluator_version), + agent_version_id=request.agent_version_id, + baseline_prompt_hash=request.baseline_prompt_hash, + candidate_prompt_hash=request.candidate_prompt_hash, + tool_policy={ + "type": "history", + "scope": "cohort_version", + "on_miss": "fail", + }, + verification_fingerprint=verification_fingerprint_for(request), + ) + save_verification_state(layout, state) + + if state.experiment_run_id is None: + backend.preflight(request) + submitted = backend.submit(request) + if not isinstance(submitted, SubmittedVerification): + submitted = SubmittedVerification.model_validate( + submitted if isinstance(submitted, dict) else submitted.__dict__ + ) + state = state.model_copy( + update={ + "experiment_id": submitted.experiment_id, + "experiment_run_id": submitted.experiment_run_id, + } + ) + save_verification_state(layout, state) + + submitted = SubmittedVerification( + experiment_id=state.experiment_id or "", + experiment_run_id=state.experiment_run_id or "", + ) + result = backend.collect(request, submitted) + if not isinstance(result, VerificationResult): + result = VerificationResult.model_validate(result) + state = state.model_copy( + update={ + "result": result, + "verification_fingerprint": result.verification_fingerprint, + "per_session": { + **{sid: "improved" for sid in result.improved_sessions}, + **{sid: "regressed" for sid in result.regressed_sessions}, + **{sid: "unchanged" for sid in result.unchanged_sessions}, + **{item.session_id: item.kind for item in result.diverged_sessions}, + }, + } + ) + save_verification_state(layout, state) + return result + + +def format_verification_report( + result: VerificationResult, + request: VerificationRequest, + *, + server_url: str = "", +) -> str: + """Human-readable summary. Never prints SHIP (issue #9).""" + + label = request.cohort_name + if request.display_version: + label = f"{request.cohort_name}/{request.display_version}" + sessions = ( + result.baseline_count + or len(result.improved_sessions) + + len(result.regressed_sessions) + + len(result.unchanged_sessions) + + len(result.diverged_sessions) + + len(result.replay_failures) + ) + def _pct(value: float | None) -> str: + return " n/a" if value is None else f"{value * 100:8.1f}%" + + def _num(value: float | None) -> str: + return " n/a" if value is None else f"{value:8.3f}" + + lines = [ + "tracegrad verification", + "────────────────────────────────", + f"Cohort: {label}", + f"Sessions: {sessions}", + "", + " Baseline Candidate", + f"Mean score {_num(result.baseline_mean_score)} {_num(result.candidate_mean_score)}", + f"Pass rate {_pct(result.baseline_pass_rate)} {_pct(result.candidate_pass_rate)}", + "", + f"Improved {len(result.improved_sessions):8d}", + f"Regressed {len(result.regressed_sessions):8d}", + f"Unchanged {len(result.unchanged_sessions):8d}", + f"Diverged {len(result.diverged_sessions):8d}", + f"Replay failures {len(result.replay_failures):8d}", + ] + if result.regressed_sessions: + lines.append("") + lines.append("Regressions") + for session_id in result.regressed_sessions: + number = request.session_numbers.get(session_id) + prefix = f"#{number}" if number is not None else session_id[:8] + short = _short_id(session_id) + lines.append(f"{prefix} session {short}") + if result.diverged_sessions: + lines.append("") + lines.append("Divergence") + for item in result.diverged_sessions: + number = item.number or request.session_numbers.get(item.session_id) + prefix = f"#{number}" if number is not None else item.session_id[:8] + lines.append(f"{prefix} {item.kind} {item.detail}".rstrip()) + if result.replay_failures: + lines.append("") + lines.append("Replay failures") + for item in result.replay_failures: + number = item.number or request.session_numbers.get(item.session_id) + prefix = f"#{number}" if number is not None else item.session_id[:8] + lines.append(f"{prefix} {item.error}".rstrip()) + lines.append("") + lines.append(f"Experiment run: {_short_id(result.experiment_run_id)}") + if server_url: + lines.append(f"Kitaru: {server_url}") + verdict = "FAILED" if result.status == "failed" else "REVIEW" + lines.append(f"Verdict: {verdict}") + return "\n".join(lines) + + +def _short_id(value: str) -> str: + text = value.removeprefix("sha256:") + if len(text) <= 12: + return text + return f"{text[:4]}…{text[-4:]}" + + +def load_proposal_for_verify(project_root: str | Path, run_id: str) -> Proposal: + return load_proposal(project_root, run_id) diff --git a/tests/test_apply.py b/tests/test_apply.py index 2bf088b..e47a94c 100644 --- a/tests/test_apply.py +++ b/tests/test_apply.py @@ -6,10 +6,13 @@ from tracegrad.apply import ( ApplyError, + Proposal, + ProposedEdit, StaleProposalError, applied_history, apply_proposal, build_proposal, + candidate_prompt, current_baseline, is_stale, latest_run_id, @@ -20,7 +23,7 @@ ) from tracegrad.canonical import text_hash from tracegrad.distill import DistillConfig, distill_trace -from tracegrad.edits import resolve_edits +from tracegrad.edits import REASON_UNKNOWN_ANCHOR, resolve_edits from tracegrad.gates import GateFlag, GateOutcome from tracegrad.inventory import build_inventory from tracegrad.schema import AttributionResult, Edit, Trace @@ -151,6 +154,47 @@ def test_loading_an_unknown_run_is_an_apply_error(tmp_path: Path) -> None: load_proposal(tmp_path, "run-9999") +def test_apply_proposal_writes_the_same_text_as_candidate_prompt(tmp_path: Path) -> None: + template = _template(tmp_path) + outcome, _ = _outcome() + proposal = build_proposal( + run_id="run-0001", template_file="prompt.md", prompt=PROMPT, outcome=outcome + ) + expected = candidate_prompt(PROMPT, proposal, [0]) + + result = apply_proposal(tmp_path, proposal, [0], base_directory=tmp_path) + + assert template.read_text(encoding="utf-8") == expected + assert result.applied_prompt_hash == text_hash(expected) + assert result.unchanged is False + + +def test_apply_proposal_records_resolution_rejections(tmp_path: Path) -> None: + template = _template(tmp_path) + inventory = build_inventory(PROMPT) + good = _edit(inventory.instructions[-1].instruction_id) + bad = _edit("i-missingxxx-99", "Ghost.") + proposal = Proposal( + run_id="run-0001", + template_file="prompt.md", + base_prompt_hash=text_hash(PROMPT), + edits=[ + ProposedEdit(edit=good, before="Cite the doc.", after="Always cite the doc."), + ProposedEdit(edit=bad, before="", after="Ghost."), + ], + ) + + written = apply_proposal(tmp_path, proposal, [0, 1], base_directory=tmp_path) + assert written.unchanged is False + assert REASON_UNKNOWN_ANCHOR in written.resolution_rejections + assert template.read_text(encoding="utf-8") == candidate_prompt(PROMPT, proposal, [0, 1]) + + template.write_text(PROMPT, encoding="utf-8") + unchanged = apply_proposal(tmp_path, proposal, [1], base_directory=tmp_path) + assert unchanged.unchanged is True + assert unchanged.resolution_rejections == (REASON_UNKNOWN_ANCHOR,) + + def test_apply_writes_the_template_and_records_the_baseline(tmp_path: Path) -> None: template = _template(tmp_path) outcome, _ = _outcome() diff --git a/tests/test_config.py b/tests/test_config.py index e326cfb..343e225 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -47,6 +47,21 @@ def test_each_rc_field_parses(tmp_path: Path) -> None: assert config.harness_presets["remote"].model == "gpt-test" +def test_kitaru_selection_parses_from_rc(tmp_path: Path) -> None: + (tmp_path / ".tracegradrc").write_text( + """ +[kitaru] +cohort = "support-production" +evaluation = "quality" +""".lstrip(), + encoding="utf-8", + ) + + config = load_config(tmp_path) + assert config.kitaru.cohort == "support-production" + assert config.kitaru.evaluation == "quality" + + def test_malformed_toml_is_rejected_with_filename(tmp_path: Path) -> None: rc = tmp_path / ".tracegradrc" rc.write_text("minEffect = [", encoding="utf-8") diff --git a/tests/test_kitaru_mapping.py b/tests/test_kitaru_mapping.py new file mode 100644 index 0000000..f3905f0 --- /dev/null +++ b/tests/test_kitaru_mapping.py @@ -0,0 +1,348 @@ +"""Kitaru session → Trace mapping (issue #8 definition of done). + +These tests do not import the Kitaru SDK. Nodes and evaluations are +duck-typed so core-only CI stays green without the extra. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from tracegrad.canonical import text_hash +from tracegrad.integrations.kitaru.graph import is_root_llm_node, root_llm_nodes +from tracegrad.integrations.kitaru.mapping import ( + REASON_FORMAT_ENGINE_REFUSED, + REASON_INPUT_UNAVAILABLE, + REASON_MULTIPLE_SYSTEM_PROMPTS, + REASON_OUTPUT_UNAVAILABLE, + REASON_ROOT_LLM_UNAVAILABLE, + REASON_SYSTEM_PROMPT_UNAVAILABLE, + MappedTrace, + SourceDrop, + map_batch, + map_session, +) +from tracegrad.integrations.kitaru.pointer import resolve_pointer, resolve_text_selector +from tracegrad.integrations.kitaru.scores import ( + REASON_AMBIGUOUS_EVALUATION, + REASON_RATIONALE_MISSING, + REASON_SCORE_OUT_OF_RANGE, + REASON_SCORE_UNAVAILABLE, + REASON_SCORE_UNSUPPORTED, + judge_fingerprint_for, + map_judge, + select_evaluation, +) +from tracegrad.integrations.kitaru.source import check_judge_fingerprint, refuse_format_engine +from tracegrad.schema import Manifest, TemplateEngine + +SYSTEM = "You are the root agent." +RATIONALE = "The agent skipped the citation the customer asked for." + + +def _node( + index: int, + *, + node_type: str = "llm_call", + parent: int | None = None, + secondary: list[int] | None = None, + system: str | None = SYSTEM, + input_text: str = "user question", + output_text: str = "model answer", + model: str = "gpt-4.1", + in_sel: str | None = "/input", + out_sel: str | None = "/output", + sys_sel: str | None = "/system", +) -> SimpleNamespace: + inputs: dict[str, object] = {"input": input_text} + if system is not None: + inputs["system"] = system + return SimpleNamespace( + index=index, + parent_index=parent, + secondary_parent_indexes=list(secondary or []), + node_type=node_type, + inputs=inputs, + outputs={"output": output_text}, + input_text_selector=in_sel, + output_text_selector=out_sel, + system_prompt_selector=sys_sel, + model=model, + ) + + +def _eval( + *, + name: str = "quality", + score: object = 0.25, + explanation: str | None = RATIONALE, + passed: bool | None = False, + data_type: str = "float", + version: int | None = 3, + value: str | None = None, + eval_id: str = "e1", +) -> SimpleNamespace: + return SimpleNamespace( + id=eval_id, + name=name, + score=score, + explanation=explanation, + passed=passed, + value=value, + data_type=data_type, + evaluator_name="quality", + evaluator_version=version, + evaluator_version_id="ev-3", + ) + + +def _session(session_id: str = "0f3a0000-0000-4000-8000-00000000c19d", number: int = 4811): + return SimpleNamespace(id=session_id, number=number) + + +def test_json_pointer_resolves_and_unescapes() -> None: + document = {"a": [{"b~c": "ok"}], "x/y": "slash"} + assert resolve_pointer(document, "") is document + assert resolve_pointer(document, "/a/0/b~0c") == "ok" + + +def test_selector_failure_is_none_not_a_guess() -> None: + assert resolve_text_selector({"n": 1}, "/n") is None + assert resolve_text_selector({"n": "s"}, None) is None + assert resolve_text_selector({"n": "s"}, "/missing") is None + + +def test_session_maps_to_a_trace() -> None: + result = map_session( + _session(), + [_node(0)], + [_eval()], + "quality", + ) + assert isinstance(result, MappedTrace) + assert result.trace.trace_id == "0f3a0000-0000-4000-8000-00000000c19d" + assert result.trace.input == "user question" + assert result.trace.output == "model answer" + assert result.trace.prompt_hash == text_hash(SYSTEM) + assert result.trace.judge.score == 0.25 + assert result.trace.meta is not None + assert result.trace.meta.model == "gpt-4.1" + assert result.session_number == 4811 + assert result.evaluator_version == 3 + + +def test_single_system_prompt_is_accepted() -> None: + result = map_session(_session(), [_node(0), _node(1, parent=0, system=SYSTEM)], [_eval()], "quality") + assert isinstance(result, MappedTrace) + assert result.system_prompt == SYSTEM + + +def test_missing_system_prompt_drops() -> None: + node = _node(0, sys_sel="/missing") + result = map_session(_session(), [node], [_eval()], "quality") + assert isinstance(result, SourceDrop) + assert result.reason == REASON_SYSTEM_PROMPT_UNAVAILABLE + + +def test_empty_system_prompt_is_not_guessed() -> None: + node = _node(0, system="") + result = map_session(_session(), [node], [_eval()], "quality") + assert isinstance(result, SourceDrop) + assert result.reason == REASON_SYSTEM_PROMPT_UNAVAILABLE + + +def test_multiple_system_prompts_drop() -> None: + nodes = [_node(0, system="prompt A"), _node(1, parent=0, system="prompt B")] + result = map_session(_session(), nodes, [_eval()], "quality") + assert isinstance(result, SourceDrop) + assert result.reason == REASON_MULTIPLE_SYSTEM_PROMPTS + + +def test_selector_resolution_failures_drop_by_name() -> None: + missing_input = map_session( + _session(), [_node(0, in_sel="/nope")], [_eval()], "quality" + ) + missing_output = map_session( + _session(), [_node(0, out_sel="/nope")], [_eval()], "quality" + ) + assert isinstance(missing_input, SourceDrop) + assert missing_input.reason == REASON_INPUT_UNAVAILABLE + assert isinstance(missing_output, SourceDrop) + assert missing_output.reason == REASON_OUTPUT_UNAVAILABLE + + +def test_root_vs_subagent_including_secondary_parent() -> None: + # 0 root llm, 1 subagent, 2 llm under subagent, 3 root llm, + # 4 llm with parents 3 (root) and 2 (subagent descendant) — not root. + nodes = [ + _node(0), + _node(1, node_type="subagent_call", parent=0, system="subagent prompt"), + _node(2, parent=1, system="subagent prompt", output_text="subagent said"), + _node(3, parent=0, input_text="second turn", output_text="last answer"), + _node(4, parent=3, secondary=[2], system="should not be used"), + _node(5, node_type="tool_call", parent=3, output_text="TOOL OUTPUT"), + ] + by_index = {node.index: node for node in nodes} + assert is_root_llm_node(nodes[0], by_index) + assert not is_root_llm_node(nodes[1], by_index) + assert not is_root_llm_node(nodes[2], by_index) + assert is_root_llm_node(nodes[3], by_index) + assert not is_root_llm_node(nodes[4], by_index) + roots = root_llm_nodes(nodes) + assert [n.index for n in roots] == [0, 3] + + result = map_session(_session(), nodes, [_eval()], "quality") + assert isinstance(result, MappedTrace) + assert result.trace.input == "user question" + assert result.trace.output == "last answer" + assert result.trace.output != "TOOL OUTPUT" + assert result.trace.output != "subagent said" + assert result.system_prompt == SYSTEM + assert result.multi_turn is True + + +def test_tool_output_cannot_become_trace_output() -> None: + nodes = [ + _node(0, output_text="llm answer"), + _node(1, node_type="tool_call", parent=0, output_text="secret tool payload"), + ] + result = map_session(_session(), nodes, [_eval()], "quality") + assert isinstance(result, MappedTrace) + assert result.trace.output == "llm answer" + + +def test_no_root_llm_drops() -> None: + nodes = [_node(0, node_type="subagent_call"), _node(1, parent=0)] + result = map_session(_session(), nodes, [_eval()], "quality") + assert isinstance(result, SourceDrop) + assert result.reason == REASON_ROOT_LLM_UNAVAILABLE + + +def test_float_bool_passed_out_of_range_and_categorical_mapping() -> None: + assert map_judge(_eval(score=0.4, data_type="float")).score == 0.4 # type: ignore[union-attr] + bool_judge = map_judge(_eval(score=True, data_type="bool", passed=None)) + assert bool_judge.score == 1.0 # type: ignore[union-attr] + passed_only = map_judge(_eval(score=None, passed=False, data_type="bool")) + assert passed_only.score == 0.0 # type: ignore[union-attr] + assert map_judge(_eval(score=1.5, data_type="float")) == REASON_SCORE_OUT_OF_RANGE + assert map_judge(_eval(score=-0.1, data_type="float")) == REASON_SCORE_OUT_OF_RANGE + assert ( + map_judge(_eval(score=None, value="bad", data_type="str", passed=None)) + == REASON_SCORE_UNSUPPORTED + ) + assert ( + map_judge(_eval(score=0.2, value="cat", data_type="categorical")) + == REASON_SCORE_UNSUPPORTED + ) + + +def test_missing_rationale_drops() -> None: + assert map_judge(_eval(explanation=None)) == REASON_RATIONALE_MISSING + assert map_judge(_eval(explanation=" ")) == REASON_RATIONALE_MISSING + + +def test_missing_score_drops() -> None: + assert ( + map_judge(_eval(score=None, passed=None, data_type="float")) + == REASON_SCORE_UNAVAILABLE + ) + + +def test_ambiguous_evaluator_version_on_one_session() -> None: + selected = select_evaluation( + [_eval(version=2, eval_id="a"), _eval(version=3, eval_id="b")], + "quality", + ) + assert selected == REASON_AMBIGUOUS_EVALUATION + + +def test_ambiguous_evaluator_version_across_the_cohort_refuses() -> None: + records = [ + (_session("s1"), [_node(0)], [_eval(version=2)]), + (_session("s2", number=2), [_node(0)], [_eval(version=3)]), + ] + result = map_batch(records, "quality") + assert result == REASON_AMBIGUOUS_EVALUATION + + +def test_format_engine_is_refused_with_a_named_error(tmp_path) -> None: + manifest = Manifest( + template_file=tmp_path / "prompt.md", + engine=TemplateEngine.FORMAT, + judge_fingerprint="quality@3", + ) + with pytest.raises(Exception, match=REASON_FORMAT_ENGINE_REFUSED): + refuse_format_engine(manifest) + + +def test_conflicting_judge_fingerprint_is_an_error(tmp_path) -> None: + manifest = Manifest( + template_file=tmp_path / "prompt.md", + engine=TemplateEngine.NONE, + judge_fingerprint="other-judge", + ) + with pytest.raises(Exception, match="judge-fingerprint-conflict"): + check_judge_fingerprint(manifest, judge_fingerprint_for("quality", 3)) + + +def test_matching_derived_fingerprint_is_accepted(tmp_path) -> None: + derived = judge_fingerprint_for("quality", 3) + manifest = Manifest( + template_file=tmp_path / "prompt.md", + engine="none", + judge_fingerprint=derived, + ) + check_judge_fingerprint(manifest, derived) + + +def test_source_and_batch_reasons_stay_kebab_case() -> None: + drop = map_session(_session(), [_node(0, sys_sel=None)], [_eval()], "quality") + assert isinstance(drop, SourceDrop) + assert "-" in drop.reason + assert drop.reason == drop.reason.lower() + + +def test_evaluator_id_looks_up_mapped_evaluator_name_not_the_cli_flag() -> None: + from tracegrad.integrations.kitaru.client import CohortResolution + from tracegrad.integrations.kitaru.source import _fetch_and_map + + looked_up: list[str] = [] + evaluation = _eval() + evaluation.evaluator_name = "quality-judge" + + class Gateway: + async def list_sessions(self, cohort_version_id: str) -> list[object]: + return [_session()] + + async def fetch_records(self, sessions: list[object]) -> list[object]: + return [(sessions[0], [_node(0)], [evaluation])] + + async def evaluator_id(self, name: str) -> str: + looked_up.append(name) + if name != "quality-judge": + raise LookupError(f"kitaru evaluator {name!r} was not found") + return "eid-judge" + + fingerprint, meta, mapped, _dropped = asyncio.run( + _fetch_and_map( + gateway=Gateway(), + resolution=CohortResolution( + cohort_id="c", + cohort_name="support-production", + cohort_version_id="cv", + display_version="week-34", + version_number=1, + agent_id="a", + session_count=1, + ), + evaluation_name="quality", + ) + ) + assert looked_up == ["quality-judge"] + assert fingerprint.evaluator_id == "eid-judge" + assert fingerprint.evaluation_name == "quality" + assert meta.evaluator_name == "quality-judge" + assert mapped diff --git a/tests/test_kitaru_optional.py b/tests/test_kitaru_optional.py new file mode 100644 index 0000000..40cd0bf --- /dev/null +++ b/tests/test_kitaru_optional.py @@ -0,0 +1,143 @@ +"""Core-only regression: Kitaru is optional (issue #8 / #9).""" + +from __future__ import annotations + +import ast +import io +import sys +from pathlib import Path + +import pytest + +from tracegrad import cli +from tracegrad.integrations.kitaru.errors import ( + INSTALL_MESSAGE, + KITARU_PIN, + NO_BACKEND_MESSAGE, + KitaruNotInstalled, +) +from tracegrad.integrations.kitaru.require import kitaru_available, require_kitaru + +PACKAGE = Path(__file__).resolve().parents[1] / "src" / "tracegrad" + + +def test_import_tracegrad_does_not_import_kitaru() -> None: + sys.modules.pop("kitaru", None) + import tracegrad as package + + assert package.__version__ + assert "kitaru" not in sys.modules + + +def test_only_the_integration_package_imports_the_kitaru_sdk() -> None: + offenders: list[str] = [] + for source in PACKAGE.rglob("*.py"): + rel = source.relative_to(PACKAGE) + if rel.parts[:2] == ("integrations", "kitaru"): + continue + tree = ast.parse(source.read_text(encoding="utf-8"), filename=str(source)) + for node in ast.walk(tree): + names: list[str] = [] + if isinstance(node, ast.Import): + names.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + names.append(node.module) + if any(name == "kitaru" or name.startswith("kitaru.") for name in names): + offenders.append(str(rel)) + assert offenders == [], offenders + + +def test_kitaru_pin_matches_the_issue() -> None: + assert KITARU_PIN == "kitaru>=0.22,<0.23" + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + text = pyproject.read_text(encoding="utf-8") + assert 'kitaru = ["kitaru>=0.22,<0.23"]' in text + assert "kitaru>=" not in text.split("[project]")[1].split("[project.optional-dependencies]")[0] + + +def test_require_kitaru_is_actionable_when_missing() -> None: + if kitaru_available(): + return + try: + require_kitaru() + except KitaruNotInstalled as exc: + assert "tracegrad[kitaru]" in str(exc) + assert "uv tool install" in str(exc) + assert INSTALL_MESSAGE.splitlines()[0] in str(exc) + else: + raise AssertionError("expected KitaruNotInstalled") + + +def test_source_kitaru_without_the_extra_is_actionable( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + if kitaru_available(): + return + stream = io.StringIO() + code = cli.main( + [ + "run", + "--source", + "kitaru", + "--kitaru-cohort", + "support-production", + "--kitaru-evaluation", + "quality", + "--manifest", + str(tmp_path / "missing.json"), + "--project-root", + str(tmp_path), + ], + out=stream, + ) + assert code == 1 + err = capsys.readouterr().err + assert "tracegrad[kitaru]" in err + assert "ImportError" not in err + + +def test_verify_without_a_backend_exits_nonzero(tmp_path: Path) -> None: + stream = io.StringIO() + code = cli.main(["verify", "--project-root", str(tmp_path)], out=stream) + assert code == 1 + output = stream.getvalue() + assert "needs a backend" in output + assert NO_BACKEND_MESSAGE.splitlines()[0] in output + + +def test_verify_does_not_break_run_apply_or_trends(tmp_path: Path) -> None: + """A missing backend is a verify failure, not a lock on the rest of the CLI.""" + + init_code, _ = _run("init", "--project-root", str(tmp_path)) + trends_code, trends_out = _run("trends", "--project-root", str(tmp_path)) + apply_code, apply_out = _run("apply", "--all", "--project-root", str(tmp_path)) + assert init_code == 0 + assert trends_code == 0 + assert "at least two runs" in trends_out + assert apply_code == 1 + assert "no proposal" in apply_out + + +def _run(*argv: str) -> tuple[int, str]: + stream = io.StringIO() + code = cli.main(list(argv), out=stream) + return code, stream.getvalue() + + +def test_traces_and_source_are_mutually_exclusive(tmp_path: Path) -> None: + stream = io.StringIO() + code = cli.main( + [ + "run", + "--traces", + str(tmp_path / "batch.jsonl"), + "--source", + "kitaru", + "--manifest", + str(tmp_path / "manifest.json"), + "--project-root", + str(tmp_path), + ], + out=stream, + ) + assert code == 1 diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py new file mode 100644 index 0000000..e0c653b --- /dev/null +++ b/tests/test_kitaru_snapshot.py @@ -0,0 +1,834 @@ +"""Snapshot reuse for the Kitaru source (ADR 0004).""" + +from __future__ import annotations + +import asyncio +import inspect +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tracegrad.canonical import text_hash +from tracegrad.integrations.kitaru.accounting import format_source_table +from tracegrad.integrations.kitaru.client import CohortResolution, KitaruGateway +from tracegrad.integrations.kitaru.errors import KitaruSourceError +from tracegrad.integrations.kitaru.mapping import MappedTrace, SourceDrop +from tracegrad.integrations.kitaru.snapshot import ( + DROPS_FILENAME, + LATEST_POINTER, + SourceFingerprint, + SourceMeta, + find_local_snapshot, + fingerprints_compatible, + kitaru_root, + load_fingerprint, + load_latest_pointer, + load_meta, + load_source_drops, + snapshot_exists, + snapshot_key, + write_snapshot, +) +from tracegrad.integrations.kitaru.source import prepare_kitaru_source +from tracegrad.schema import Manifest, Trace +from tracegrad.state import initialize + + +def _sample_mapped( + *, + trace_id: str = "s1", + evaluator_name: str = "quality", + evaluator_version: int = 3, +) -> list[MappedTrace]: + trace = Trace( + trace_id=trace_id, + input="q", + output="a", + judge={"score": 0.2, "rationale": "needs a citation in the answer now"}, + prompt_hash=text_hash("prompt"), + ) + return [ + MappedTrace( + trace=trace, + session_number=12, + evaluator_name=evaluator_name, + evaluator_version=evaluator_version, + evaluator_version_id="ev", + multi_turn=False, + system_prompt="prompt", + ) + ] + + +def _sample_fingerprint(**overrides: object) -> SourceFingerprint: + payload = dict( + cohort_id="c", + cohort_version_id="cv", + evaluation_name="quality", + evaluator_id="eid", + evaluator_version=3, + agent_id="a", + ) + payload.update(overrides) + return SourceFingerprint.model_validate(payload) + + +def _sample_meta(**overrides: object) -> SourceMeta: + payload = dict( + cohort_name="support-production", + display_version="week-34", + version_number=1, + traces_mapped=1, + sessions_selected=2, + evaluator_name="quality", + ) + payload.update(overrides) + return SourceMeta.model_validate(payload) + + +def _manifest(judge_fingerprint: str = "quality@3") -> Manifest: + return Manifest(template_file=Path("prompt.md"), judge_fingerprint=judge_fingerprint) + + +def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: + layout = initialize(tmp_path) + mapped = _sample_mapped() + dropped = [SourceDrop("s2", "system-prompt-unavailable", number=13)] + fingerprint = _sample_fingerprint() + meta = _sample_meta() + write_snapshot(layout, fingerprint=fingerprint, meta=meta, mapped=mapped, dropped=dropped) + + assert snapshot_exists(layout, snapshot_key("cv", "quality")) + stored = load_fingerprint(layout, snapshot_key("cv", "quality")) + assert stored.mapping_version == 1 + assert stored.source == "kitaru" + assert fingerprints_compatible( + stored, {"evaluation_name": "quality", "cohort_version_id": "cv"} + ) + assert not fingerprints_compatible( + stored, {"evaluation_name": "other", "cohort_version_id": "cv"} + ) + assert not fingerprints_compatible( + stored, {"evaluation_name": "quality", "cohort_version_id": "other"} + ) + drops = load_source_drops(layout, snapshot_key("cv", "quality")) + assert drops[0].reason == "system-prompt-unavailable" + batch = ( + layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl" + ).read_text(encoding="utf-8") + assert "s1" in batch + pointer = load_latest_pointer( + layout, cohort_name="support-production", evaluation_name="quality" + ) + assert pointer is not None + assert pointer["cohort_version_id"] == "cv" + assert pointer["snapshot_id"] == snapshot_key("cv", "quality") + assert pointer["cohort_name"] == "support-production" + assert pointer["evaluation_name"] == "quality" + assert pointer["display_version"] == "week-34" + stored_meta = load_meta(layout, snapshot_key("cv", "quality")) + assert stored_meta.version_number == 1 + assert pointer["version_number"] == 1 + assert find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) == snapshot_key("cv", "quality") + assert find_local_snapshot( + layout, + cohort_name="support-production", + evaluation_name="quality", + cohort_version="1", + ) == snapshot_key("cv", "quality") + assert ( + find_local_snapshot( + layout, + cohort_name="support-production", + evaluation_name="quality", + cohort_version="2", + ) + is None + ) + + +def test_source_and_batch_tables_are_not_merged() -> None: + table = format_source_table( + sessions_selected=10, + traces_mapped=8, + dropped=[ + SourceDrop("a", "system-prompt-unavailable"), + SourceDrop("b", "judge-rationale-missing"), + ], + in_batch=6, + batch_drops={"prompt-hash-partition": 2}, + ) + assert "Sessions selected" in table + assert "Traces mapped" in table + assert "In batch" in table + assert "system-prompt-unavailable" in table + assert "prompt-hash-partition" in table + # The two tables stay labeled separately rather than summed. + assert table.index("Traces mapped") < table.index("In batch") + + +def test_prepare_reuses_latest_pointer_without_the_server(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + prepared = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + ) + assert prepared.refreshed is False + assert prepared.fingerprint.cohort_version_id == "cv" + assert prepared.traces_path == layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl" + + +def test_prepare_does_not_follow_moved_remote_latest(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + + class MovingLatest: + resolved = 0 + + async def resolve_cohort(self, name: str, version: str | None = None) -> CohortResolution: + type(self).resolved += 1 + return CohortResolution( + cohort_id="c", + cohort_name=name, + cohort_version_id="cv-new", + display_version="week-35", + version_number=2, + agent_id="a", + session_count=1, + ) + + async def close(self) -> None: + return None + + prepared = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + gateway=MovingLatest(), + ) + assert MovingLatest.resolved == 0 + assert prepared.fingerprint.cohort_version_id == "cv" + assert prepared.refreshed is False + + +def test_prepare_reuses_explicit_version_and_display_version(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + by_id = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + cohort_version="cv", + ) + by_display = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + cohort_version="week-34", + ) + by_number = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + cohort_version="1", + ) + assert by_id.fingerprint.cohort_version_id == "cv" + assert by_display.fingerprint.cohort_version_id == "cv" + assert by_number.fingerprint.cohort_version_id == "cv" + assert by_id.refreshed is False + assert by_display.refreshed is False + assert by_number.refreshed is False + + +def test_prepare_refresh_fetches_even_when_latest_exists(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + + class FetchGateway: + def __init__(self) -> None: + self.resolved = 0 + + async def resolve_cohort(self, name: str, version: str | None = None) -> CohortResolution: + self.resolved += 1 + return CohortResolution( + cohort_id="c", + cohort_name=name, + cohort_version_id="cv-new", + display_version="week-35", + version_number=2, + agent_id="a", + session_count=0, + ) + + async def list_sessions(self, cohort_version_id: str) -> list[object]: + return [] + + async def fetch_records(self, sessions: list[object]) -> list[object]: + return [] + + async def evaluator_id(self, name: str) -> str: + return "eid" + + async def close(self) -> None: + return None + + gateway = FetchGateway() + prepared = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + refresh=True, + gateway=gateway, + ) + assert gateway.resolved == 1 + assert prepared.refreshed is True + assert prepared.fingerprint.cohort_version_id == "cv-new" + pointer = load_latest_pointer( + layout, cohort_name="support-production", evaluation_name="quality" + ) + assert pointer is not None + assert pointer["cohort_version_id"] == "cv-new" + assert pointer["snapshot_id"] == snapshot_key("cv-new", "quality") + + +def test_refresh_fingerprint_conflict_does_not_clobber_last_good_snapshot(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + previous_batch = ( + layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl" + ).read_text(encoding="utf-8") + + session = SimpleNamespace(id="0f3a0000-0000-4000-8000-00000000c19d", number=12) + node = SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"input": "q", "system": "prompt"}, + outputs={"output": "a"}, + input_text_selector="/input", + output_text_selector="/output", + system_prompt_selector="/system", + model="gpt-4.1", + ) + evaluation = SimpleNamespace( + id="e1", + name="quality", + score=0.2, + explanation="needs a citation in the answer now", + passed=False, + value=None, + data_type="float", + evaluator_name="quality", + evaluator_version=9, + evaluator_version_id="ev-9", + ) + + class ConflictGateway: + async def resolve_cohort(self, name: str, version: str | None = None) -> CohortResolution: + return CohortResolution( + cohort_id="c", + cohort_name=name, + cohort_version_id="cv-new", + display_version="week-35", + version_number=2, + agent_id="a", + session_count=1, + ) + + async def list_sessions(self, cohort_version_id: str) -> list[object]: + return [session] + + async def fetch_records(self, sessions: list[object]) -> list[object]: + return [(sessions[0], [node], [evaluation])] + + async def evaluator_id(self, name: str) -> str: + return "eid-new" + + async def close(self) -> None: + return None + + with pytest.raises(Exception, match="judge-fingerprint-conflict"): + prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest("quality@3"), + cohort_name="support-production", + evaluation_name="quality", + refresh=True, + gateway=ConflictGateway(), + ) + + pointer = load_latest_pointer( + layout, cohort_name="support-production", evaluation_name="quality" + ) + assert pointer is not None + assert pointer["cohort_version_id"] == "cv" + assert pointer["snapshot_id"] == snapshot_key("cv", "quality") + assert snapshot_exists(layout, snapshot_key("cv", "quality")) + assert not snapshot_exists(layout, snapshot_key("cv-new", "quality")) + assert ( + layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl" + ).read_text(encoding="utf-8") == previous_batch + + +def test_old_latest_pointer_without_extra_fields_still_reuses(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + pointer_path = kitaru_root(layout) / LATEST_POINTER + pointer_path.write_text( + '{"cohort_version_id": "cv", "batch": "ignored"}', + encoding="utf-8", + ) + prepared = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + ) + assert prepared.fingerprint.cohort_version_id == "cv" + assert prepared.refreshed is False + + +def test_two_evaluations_of_the_same_cohort_do_not_clobber(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(evaluation_name="quality"), + meta=_sample_meta(evaluator_name="quality"), + mapped=_sample_mapped(trace_id="quality-trace"), + dropped=[], + ) + write_snapshot( + layout, + fingerprint=_sample_fingerprint( + evaluation_name="safety", evaluator_id="eid-s", evaluator_version=1 + ), + meta=_sample_meta(evaluator_name="safety"), + mapped=_sample_mapped( + trace_id="safety-trace", evaluator_name="safety", evaluator_version=1 + ), + dropped=[], + ) + quality_batch = (layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl").read_text( + encoding="utf-8" + ) + safety_batch = (layout.sources / "kitaru" / "cv" / "safety" / "batch.jsonl").read_text( + encoding="utf-8" + ) + assert "quality-trace" in quality_batch + assert "safety-trace" in safety_batch + assert "safety-trace" not in quality_batch + assert find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) == snapshot_key("cv", "quality") + assert find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="safety" + ) == snapshot_key("cv", "safety") + quality = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest("quality@3"), + cohort_name="support-production", + evaluation_name="quality", + ) + safety = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest("safety@1"), + cohort_name="support-production", + evaluation_name="safety", + ) + assert quality.fingerprint.evaluation_name == "quality" + assert "quality-trace" in quality.traces_path.read_text(encoding="utf-8") + assert safety.fingerprint.evaluation_name == "safety" + assert "safety-trace" in safety.traces_path.read_text(encoding="utf-8") + + +def test_versionless_rerun_keeps_last_fetched_after_another_cohort(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(cohort_version_id="aaa-old"), + meta=_sample_meta(display_version="week-33"), + mapped=_sample_mapped(trace_id="old"), + dropped=[], + ) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(cohort_version_id="zzz-new"), + meta=_sample_meta(display_version="week-34"), + mapped=_sample_mapped(trace_id="new"), + dropped=[], + ) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(cohort_id="c2", cohort_version_id="other-cv"), + meta=_sample_meta(cohort_name="billing-production", display_version="week-1"), + mapped=_sample_mapped(trace_id="other"), + dropped=[], + ) + first = load_latest_pointer( + layout, cohort_name="support-production", evaluation_name="quality" + ) + other = load_latest_pointer( + layout, cohort_name="billing-production", evaluation_name="quality" + ) + assert first is not None + assert first["cohort_version_id"] == "zzz-new" + assert other is not None + assert other["cohort_version_id"] == "other-cv" + prepared = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + ) + assert prepared.fingerprint.cohort_version_id == "zzz-new" + assert prepared.refreshed is False + assert "new" in prepared.traces_path.read_text(encoding="utf-8") + + +def test_missing_pointer_picks_newest_matching_snapshot(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(cohort_version_id="aaa-old"), + meta=_sample_meta(display_version="week-33"), + mapped=_sample_mapped(trace_id="old"), + dropped=[], + ) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(cohort_version_id="zzz-new"), + meta=_sample_meta(display_version="week-34"), + mapped=_sample_mapped(trace_id="new"), + dropped=[], + ) + pointer_path = kitaru_root(layout) / LATEST_POINTER + pointer_path.write_text('{"entries": []}', encoding="utf-8") + older = layout.sources / "kitaru" / "aaa-old" / "quality" / "fingerprint.json" + newer = layout.sources / "kitaru" / "zzz-new" / "quality" / "fingerprint.json" + os.utime(newer, (1_000_000_000, 1_000_000_000)) + os.utime(older, (2_000_000_000, 2_000_000_000)) + found = find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) + assert found == snapshot_key("aaa-old", "quality") + + +def test_pick_version_matches_id_display_and_number_then_none() -> None: + versions = [ + SimpleNamespace( + id="aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + display_version="week-34", + version=1, + ), + SimpleNamespace( + id="bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", + display_version="week-35", + version=2, + ), + ] + pick = KitaruGateway._pick_version + sentinel = object() + assert pick(sentinel, versions, 2, None).version == 2 + assert pick(sentinel, versions, 2, "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa").version == 1 + assert pick(sentinel, versions, 2, "week-35").version == 2 + assert pick(sentinel, versions, 2, "1").version == 1 + assert pick(sentinel, versions, 2, "cccccccc-cccc-4ccc-8ccc-cccccccccccc") is None + assert pick(sentinel, versions, 2, "not-a-version") is None + assert pick(sentinel, [], 1, None) is None + + +def test_fetch_records_wraps_gather_failures_as_source_error() -> None: + class Fake: + async def session_bundle(self, session_id: str) -> tuple[object, ...]: + raise TimeoutError(f"404 fetching session {session_id}") + + with pytest.raises(KitaruSourceError, match="source fetch aborted") as caught: + asyncio.run(KitaruGateway.fetch_records(Fake(), [SimpleNamespace(id="s0")])) + message = str(caught.value) + assert "TimeoutError" in message + assert "batch is not mapped" in message + assert caught.value.__cause__ is not None + + +def test_corrupt_snapshot_dir_is_skipped_during_reuse(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[], + ) + bad = layout.sources / "kitaru" / "cv-bad" / "quality" + bad.mkdir(parents=True) + (bad / "batch.jsonl").write_text("{}\n", encoding="utf-8") + (bad / "fingerprint.json").write_text("not-json", encoding="utf-8") + found = find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) + assert found == snapshot_key("cv", "quality") + + meta_path = layout.sources / "kitaru" / "cv" / "quality" / "meta.json" + meta_path.unlink() + assert ( + find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) + is None + ) + + +def test_corrupt_source_drops_are_skipped_during_reuse(tmp_path) -> None: + layout = initialize(tmp_path) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(), + meta=_sample_meta(), + mapped=_sample_mapped(), + dropped=[SourceDrop("s2", "system-prompt-unavailable", number=13)], + ) + write_snapshot( + layout, + fingerprint=_sample_fingerprint(cohort_version_id="cv-bad"), + meta=_sample_meta(display_version="week-36"), + mapped=_sample_mapped(trace_id="bad"), + dropped=[], + ) + bad_drops = layout.sources / "kitaru" / "cv-bad" / "quality" / DROPS_FILENAME + + for payload in ("{not-json\n", "{}\n", '{"reason": "no-session"}\n', "[]\n"): + bad_drops.write_text(payload, encoding="utf-8") + found = find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) + assert found == snapshot_key("cv", "quality") + prepared = prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + ) + assert prepared.refreshed is False + assert prepared.fingerprint.cohort_version_id == "cv" + + good_drops = layout.sources / "kitaru" / "cv" / "quality" / DROPS_FILENAME + good_drops.write_text('{"session_id": "s2"}\n', encoding="utf-8") + assert ( + find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) + is None + ) + + class Boom: + async def resolve_cohort(self, name: str, version: str | None = None) -> object: + raise TimeoutError("server unreachable") + + async def close(self) -> None: + return None + + with pytest.raises(KitaruSourceError, match="resolving the cohort"): + prepare_kitaru_source( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + gateway=Boom(), + ) + + +def test_source_sdk_errors_are_kitaru_source_errors(tmp_path) -> None: + resolution = CohortResolution( + cohort_id="c", + cohort_name="support-production", + cohort_version_id="cv-new", + display_version="week-35", + version_number=2, + agent_id="a", + session_count=0, + ) + + class ResolveBoom: + async def resolve_cohort(self, name: str, version: str | None = None) -> object: + raise TimeoutError("resolve 404") + + async def close(self) -> None: + return None + + class ListBoom: + async def resolve_cohort(self, name: str, version: str | None = None) -> object: + return resolution + + async def list_sessions(self, cohort_version_id: str) -> list[object]: + raise TimeoutError("list sessions 404") + + async def close(self) -> None: + return None + + class EvaluatorBoom: + async def resolve_cohort(self, name: str, version: str | None = None) -> object: + return resolution + + async def list_sessions(self, cohort_version_id: str) -> list[object]: + return [] + + async def fetch_records(self, sessions: list[object]) -> list[object]: + return [] + + async def evaluator_id(self, name: str) -> str: + raise TimeoutError("evaluator 404") + + async def close(self) -> None: + return None + + kwargs = dict( + project_root=tmp_path, + manifest=_manifest(), + cohort_name="support-production", + evaluation_name="quality", + refresh=True, + ) + with pytest.raises(KitaruSourceError, match="resolving the cohort"): + prepare_kitaru_source(**kwargs, gateway=ResolveBoom()) + with pytest.raises(KitaruSourceError, match="listing sessions"): + prepare_kitaru_source(**kwargs, gateway=ListBoom()) + with pytest.raises(KitaruSourceError, match="looking up the evaluator"): + prepare_kitaru_source(**kwargs, gateway=EvaluatorBoom()) + + +def test_wait_for_experiment_run_polls_get_not_a_missing_client_method( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """kitaru 0.22's API client has experiment_runs.get, not wait_for_experiment_run.""" + + pytest.importorskip("kitaru") + from importlib.metadata import version + + from kitaru.client.api_client import KitaruAPIClient + from kitaru.client.client import KitaruClient + from kitaru.client.resources.experiment_runs import ExperimentRunsResource + + assert version("kitaru").startswith("0.22") + assert hasattr(ExperimentRunsResource, "get") + assert not hasattr(KitaruAPIClient, "wait_for_experiment_run") + # Wrapping KitaruClient would close the shared API client on close(). + assert "self._api_client.close" in inspect.getsource(KitaruClient.close) + + slept: list[float] = [] + + async def fake_sleep(delay: float) -> None: + slept.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + class ExperimentRuns: + def __init__(self) -> None: + self.calls = 0 + + async def get(self, experiment_run_id: object) -> SimpleNamespace: + self.calls += 1 + status = "running" if self.calls < 2 else "completed" + return SimpleNamespace(id=experiment_run_id, status=status) + + class ApiClient: + def __init__(self) -> None: + self.experiment_runs = ExperimentRuns() + self.closed = False + + async def close(self) -> None: + self.closed = True + + api = ApiClient() + assert not hasattr(api, "wait_for_experiment_run") + run = asyncio.run( + KitaruGateway(client=api).wait_for_experiment_run( + "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" + ) + ) + assert run.status == "completed" + assert api.experiment_runs.calls == 2 + assert api.closed is False + assert slept == [2.0] + + +@pytest.mark.parametrize("status", ["failed", "canceled"]) +def test_wait_for_experiment_run_returns_terminal_status(status: str) -> None: + pytest.importorskip("kitaru") + + class ExperimentRuns: + async def get(self, experiment_run_id: object) -> SimpleNamespace: + return SimpleNamespace(id=experiment_run_id, status=status) + + api = SimpleNamespace(experiment_runs=ExperimentRuns()) + run = asyncio.run( + KitaruGateway(client=api).wait_for_experiment_run( + "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" + ) + ) + assert run.status == status + + +def test_wait_for_experiment_run_timeout_is_fail_closed() -> None: + pytest.importorskip("kitaru") + + class ExperimentRuns: + async def get(self, experiment_run_id: object) -> SimpleNamespace: + return SimpleNamespace(id=experiment_run_id, status="running") + + api = SimpleNamespace(experiment_runs=ExperimentRuns()) + with pytest.raises(TimeoutError, match="Timed out waiting for experiment run"): + asyncio.run( + KitaruGateway(client=api).wait_for_experiment_run( + "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + timeout=0.01, + ) + ) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py new file mode 100644 index 0000000..f9c9f74 --- /dev/null +++ b/tests/test_kitaru_verify.py @@ -0,0 +1,1495 @@ +"""Phase 2 verification: gate, resume, policy, override scope (issue #9).""" + +from __future__ import annotations + +import asyncio +import io +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tracegrad.apply import ( + Proposal, + ProposedEdit, + StaleProposalError, + applied_history, + apply_proposal, + candidate_prompt, + save_proposal, +) +from tracegrad.canonical import text_hash +from tracegrad.edits import resolve_edits +from tracegrad.integrations.kitaru.backend import ( + KitaruVerificationBackend, + assert_override_scope, + classify_replay_session, + classify_scores, + is_tool_history_miss, + mixed_agent_version_message, +) +from tracegrad.integrations.kitaru.client import FETCH_JOBS, worker_covers_agent_version +from tracegrad.integrations.kitaru.errors import KitaruVerifyError +from tracegrad.integrations.kitaru.policy import ( + RECORDED_HISTORY_POLICY, + asserts_no_passthrough, + recorded_history_policy, +) +from tracegrad.integrations.kitaru.scores import ( + REASON_AMBIGUOUS_EVALUATION, + REASON_SCORE_UNAVAILABLE, +) +from tracegrad.integrations.kitaru.snapshot import RUN_SOURCE_FILENAME +from tracegrad.inventory import build_inventory +from tracegrad.schema import Edit +from tracegrad.state import atomic_write_json, initialize +from tracegrad.verify import ( + DIVERGENCE_EVALUATOR_VERSION, + DIVERGENCE_HISTORY, + DIVERGENCE_SCORE, + DIVERGENCE_SELECT, + Divergence, + ReplayFailure, + SubmittedVerification, + VerificationRequest, + VerificationResult, + VerifyError, + backend_is_configured, + build_request, + format_verification_report, + load_run_source_payload, + load_verification_state, + matching_verification, + refuse_ungated_apply, + run_verification, + verification_id_for, + verification_path, +) + +PROMPT = "Rules:\n- Be concise.\n- Cite the doc.\n" +CANDIDATE = "Rules:\n- Be concise.\n- Always cite the doc.\n" + + +def _proposal(project: Path, prompt: str = PROMPT) -> Proposal: + (project / "prompt.md").write_text(prompt, encoding="utf-8") + inventory = build_inventory(prompt) + resolution = resolve_edits( + inventory, + [ + Edit( + instruction_id=inventory.instructions[-1].instruction_id, + operation="REWRITE", + text="Always cite the doc.", + covers_theme="missing-citation", + watch_metric="missing-citation", + ) + ], + ) + proposal = Proposal( + run_id="run-0001", + template_file="prompt.md", + base_prompt_hash=text_hash(prompt), + edits=[ + ProposedEdit( + edit=item.edit, + before=item.anchor.text if item.anchor else "", + after=item.replacement, + ) + for item in resolution.resolved + ], + ) + save_proposal(project, proposal) + return proposal + + +def _multi_edit_proposal(project: Path, prompt: str = PROMPT) -> Proposal: + (project / "prompt.md").write_text(prompt, encoding="utf-8") + inventory = build_inventory(prompt) + first, second = inventory.instructions[1], inventory.instructions[2] + resolution = resolve_edits( + inventory, + [ + Edit( + instruction_id=first.instruction_id, + operation="REWRITE", + text="Be brief.", + covers_theme="verbosity", + watch_metric="verbosity", + ), + Edit( + instruction_id=second.instruction_id, + operation="REWRITE", + text="Always cite the doc.", + covers_theme="missing-citation", + watch_metric="missing-citation", + ), + ], + ) + proposal = Proposal( + run_id="run-0001", + template_file="prompt.md", + base_prompt_hash=text_hash(prompt), + edits=[ + ProposedEdit( + edit=item.edit, + before=item.anchor.text if item.anchor else "", + after=item.replacement, + ) + for item in resolution.resolved + ], + ) + save_proposal(project, proposal) + return proposal + + +def _source_sidecar( + project: Path, run_id: str = "run-0001", *, cohort_version_id: str = "cv1" +) -> None: + layout = initialize(project) + atomic_write_json( + layout.runs / run_id / RUN_SOURCE_FILENAME, + { + "fingerprint": { + "source": "kitaru", + "cohort_id": "c1", + "cohort_version_id": cohort_version_id, + "evaluation_name": "quality", + "evaluator_id": "ev", + "evaluator_version": 3, + "agent_id": "a1", + "mapping_version": 1, + }, + "meta": { + "cohort_name": "support-production", + "display_version": "week-34", + "agent_version_id": "av1", + "agent_version_counts": {"av1": 2}, + "session_numbers": {}, + "system_prompts": {}, + "multi_turn_count": 0, + "sessions_selected": 2, + "traces_mapped": 2, + "evaluator_name": "quality", + }, + }, + ) + + +def _request(**overrides: object) -> VerificationRequest: + payload = dict( + run_id="run-0001", + proposal_id="run-0001", + candidate_prompt=CANDIDATE, + candidate_prompt_hash=text_hash(CANDIDATE), + baseline_prompt_hash=text_hash(PROMPT), + cohort_id="c1", + cohort_version_id="cv1", + cohort_name="support-production", + display_version="week-34", + evaluation_name="quality", + evaluator_id="ev", + evaluator_version=3, + evaluator_name="quality", + agent_id="a1", + agent_version_id="av1", + ) + payload.update(overrides) + return VerificationRequest.model_validate(payload) + + +class FakeBackend: + name = "fake" + + def __init__(self) -> None: + self.preflighted = 0 + self.submitted = 0 + self.collected = 0 + + def preflight(self, request: VerificationRequest) -> None: + self.preflighted += 1 + + def submit(self, request: VerificationRequest) -> SubmittedVerification: + self.submitted += 1 + return SubmittedVerification(experiment_id="exp-1", experiment_run_id="erun-1") + + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + self.collected += 1 + return VerificationResult( + status="completed", + baseline_count=2, + candidate_count=2, + baseline_mean_score=0.5, + candidate_mean_score=0.7, + baseline_pass_rate=0.5, + candidate_pass_rate=0.7, + improved_sessions=["s1"], + regressed_sessions=[], + unchanged_sessions=["s2"], + diverged_sessions=[], + replay_failures=[], + cohort_version_id=request.cohort_version_id, + agent_version_id=request.agent_version_id, + evaluator_version=str(request.evaluator_version), + baseline_prompt_hash=request.baseline_prompt_hash, + candidate_prompt_hash=request.candidate_prompt_hash, + verification_fingerprint="fp", + experiment_run_id=submitted.experiment_run_id, + ) + + +def test_recorded_history_policy_is_the_only_policy() -> None: + policy = recorded_history_policy() + asserts_no_passthrough(policy) + assert policy["scope"] == "cohort_version" + assert policy["on_miss"] == "fail" + assert policy["type"] == "history" + assert "passthrough" not in json.dumps(RECORDED_HISTORY_POLICY) + + +def test_passthrough_is_rejected() -> None: + with pytest.raises(ValueError, match="fail"): + asserts_no_passthrough({"type": "history", "scope": "cohort_version", "on_miss": "passthrough"}) + + +def test_tool_history_miss_is_typed() -> None: + assert is_tool_history_miss("No history result for tool 'search_account'") + assert is_tool_history_miss("ToolPolicyMissError: No history result for tool 'lookup'") + + class ToolPolicyMissError(Exception): + pass + + assert is_tool_history_miss(ToolPolicyMissError("No history result for tool 'x'")) + assert not is_tool_history_miss("evaluator crashed") + assert not is_tool_history_miss("no recorded call matched under on_miss=fail") + assert not is_tool_history_miss("TOOL_HISTORY_MISS: search_account") + + +def test_override_scope_divergence_on_non_root() -> None: + candidate = "NEW PROMPT" + baseline = [ + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": "ROOT"}, + system_prompt_selector="/system", + ), + SimpleNamespace( + index=1, + parent_index=None, + secondary_parent_indexes=[], + node_type="subagent_call", + inputs={}, + system_prompt_selector=None, + ), + SimpleNamespace( + index=2, + parent_index=1, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": "SUB"}, + system_prompt_selector="/system", + ), + ] + result_ok = [ + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": candidate}, + system_prompt_selector="/system", + ), + SimpleNamespace( + index=1, + parent_index=None, + secondary_parent_indexes=[], + node_type="subagent_call", + inputs={}, + system_prompt_selector=None, + ), + SimpleNamespace( + index=2, + parent_index=1, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": "SUB"}, + system_prompt_selector="/system", + ), + ] + result_bad = [ + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": candidate}, + system_prompt_selector="/system", + ), + SimpleNamespace( + index=1, + parent_index=None, + secondary_parent_indexes=[], + node_type="subagent_call", + inputs={}, + system_prompt_selector=None, + ), + SimpleNamespace( + index=2, + parent_index=1, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": candidate}, + system_prompt_selector="/system", + ), + ] + assert assert_override_scope(baseline, result_ok, candidate) is None + detail = assert_override_scope(baseline, result_bad, candidate) + assert detail is not None + assert "non-root" in detail + + +def test_override_scope_diverges_when_result_has_no_root_llm() -> None: + candidate = "NEW PROMPT" + baseline = [ + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": "ROOT"}, + system_prompt_selector="/system", + ), + ] + empty_detail = assert_override_scope(baseline, (), candidate) + assert empty_detail is not None + assert "no root llm node carrying the candidate" in empty_detail + + tool_only = [ + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="tool_call", + inputs={}, + system_prompt_selector=None, + ), + ] + tool_detail = assert_override_scope(baseline, tool_only, candidate) + assert tool_detail is not None + assert "no root llm node carrying the candidate" in tool_detail + + non_root_only = [ + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="subagent_call", + inputs={}, + system_prompt_selector=None, + ), + SimpleNamespace( + index=1, + parent_index=0, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": candidate}, + system_prompt_selector="/system", + ), + ] + nested_detail = assert_override_scope(baseline, non_root_only, candidate) + assert nested_detail is not None + assert "no root llm node carrying the candidate" in nested_detail + + +def test_mixed_agent_version_message_includes_a_breakdown() -> None: + message = mixed_agent_version_message({"av-1": 12, "av-2": 3}) + assert "av-1: 12 session(s)" in message + assert "av-2: 3 session(s)" in message + + +def test_stale_workers_do_not_cover_an_agent_version() -> None: + claim = SimpleNamespace(kind="agent", agent_version_id="av1") + stale = SimpleNamespace(live=False, scope=SimpleNamespace(claims=[claim])) + live = SimpleNamespace(live=True, scope=SimpleNamespace(claims=[claim])) + assert worker_covers_agent_version(stale, "av1") is False + assert worker_covers_agent_version(live, "av1") is True + + +def test_preflight_wraps_list_live_workers_errors() -> None: + class Gateway: + async def server_info(self) -> str: + return "ok" + + async def get_cohort_version(self, cohort_version_id: str) -> object: + return SimpleNamespace(id=cohort_version_id) + + async def get_agent_version(self, agent_version_id: str) -> object: + return SimpleNamespace(id=agent_version_id) + + async def list_live_workers(self) -> list[object]: + raise RuntimeError("connection reset") + + backend = KitaruVerificationBackend(gateway=Gateway()) + with pytest.raises(KitaruVerifyError, match="could not list workers"): + backend.preflight(_request(agent_version_counts={"av1": 2})) + + +def test_preflight_ignores_stale_workers_when_checking_coverage() -> None: + claim = SimpleNamespace(kind="agent", agent_version_id="av1") + + class Gateway: + async def server_info(self) -> str: + return "ok" + + async def get_cohort_version(self, cohort_version_id: str) -> object: + return SimpleNamespace(id=cohort_version_id) + + async def get_agent_version(self, agent_version_id: str) -> object: + return SimpleNamespace(id=agent_version_id) + + async def list_live_workers(self) -> list[object]: + return [SimpleNamespace(live=False, scope=SimpleNamespace(claims=[claim]))] + + backend = KitaruVerificationBackend(gateway=Gateway()) + with pytest.raises(KitaruVerifyError, match="no live worker is polling"): + backend.preflight(_request(agent_version_counts={"av1": 2})) + + +def test_classify_fail_to_pass_is_improved() -> None: + baseline = SimpleNamespace(score=0.0, passed=False, data_type="float", value=None) + candidate = SimpleNamespace(score=1.0, passed=True, data_type="float", value=None) + assert classify_scores(baseline, candidate) == "improved" + assert classify_scores(candidate, baseline) == "regressed" + assert classify_scores(candidate, candidate) == "unchanged" + + +def _eval(*, version: int = 3, score: float | None = 0.5, passed: bool | None = True) -> SimpleNamespace: + return SimpleNamespace( + evaluator_version=version, + score=score, + passed=passed, + data_type="float", + value=None, + ) + + +def test_select_evaluation_failure_is_select_evaluation_failed() -> None: + outcome = classify_replay_session( + session_id="s-select", + number=4, + baseline_eval=REASON_SCORE_UNAVAILABLE, + candidate_eval=_eval(), + requested_evaluator_version=3, + ) + assert isinstance(outcome, Divergence) + assert outcome.kind == DIVERGENCE_SELECT + assert outcome.kind == "SELECT_EVALUATION_FAILED" + assert "baseline: judge-score-unavailable" in outcome.detail + assert outcome.session_id == "s-select" + assert outcome.number == 4 + + both = classify_replay_session( + session_id="s-both", + number=None, + baseline_eval=REASON_SCORE_UNAVAILABLE, + candidate_eval=REASON_AMBIGUOUS_EVALUATION, + requested_evaluator_version=3, + ) + assert isinstance(both, Divergence) + assert both.kind == "SELECT_EVALUATION_FAILED" + assert "baseline: judge-score-unavailable" in both.detail + assert "candidate: ambiguous-evaluation" in both.detail + + +def test_evaluator_version_mismatch_is_typed() -> None: + outcome = classify_replay_session( + session_id="s-ver", + number=2, + baseline_eval=_eval(version=2), + candidate_eval=_eval(version=3), + requested_evaluator_version=3, + ) + assert isinstance(outcome, Divergence) + assert outcome.kind == DIVERGENCE_EVALUATOR_VERSION + assert outcome.kind == "EVALUATOR_VERSION_MISMATCH" + assert "requested evaluator_version 3" in outcome.detail + assert "baseline=2" in outcome.detail + assert "candidate=3" in outcome.detail + + +def test_unclassified_scores_are_score_unclassified() -> None: + unclassifiable = SimpleNamespace( + evaluator_version=3, + score=None, + passed=None, + data_type="float", + value=None, + ) + assert classify_scores(unclassifiable, unclassifiable) is None + outcome = classify_replay_session( + session_id="s-score", + number=9, + baseline_eval=unclassifiable, + candidate_eval=unclassifiable, + requested_evaluator_version=3, + ) + assert isinstance(outcome, Divergence) + assert outcome.kind == DIVERGENCE_SCORE + assert outcome.kind == "SCORE_UNCLASSIFIED" + + # map_score drop-reason with bool passed flags is incomparable, not unchanged. + out_of_range = _eval(score=1.5, passed=True) + assert classify_scores(out_of_range, out_of_range) is None + mixed = classify_replay_session( + session_id="s-range", + number=8, + baseline_eval=out_of_range, + candidate_eval=_eval(score=0.9, passed=True), + requested_evaluator_version=3, + ) + assert isinstance(mixed, Divergence) + assert mixed.kind == "SCORE_UNCLASSIFIED" + # Pass-flip still classifies before the unreadable score. + assert classify_scores(_eval(score=1.5, passed=False), _eval(score=1.5, passed=True)) == "improved" + + +def test_replay_session_still_classifies_comparable_scores() -> None: + outcome = classify_replay_session( + session_id="s-ok", + number=1, + baseline_eval=_eval(score=0.2, passed=False), + candidate_eval=_eval(score=0.9, passed=True), + requested_evaluator_version=3, + ) + assert outcome == "improved" + + +def test_new_divergence_kinds_stay_in_per_session_buckets(tmp_path: Path) -> None: + request = _request() + result = VerificationResult( + status="completed", + baseline_count=4, + candidate_count=4, + baseline_mean_score=0.4, + candidate_mean_score=0.4, + improved_sessions=["s-improved"], + diverged_sessions=[ + Divergence( + session_id="s-select", + kind="SELECT_EVALUATION_FAILED", + detail="baseline: judge-score-unavailable", + number=1, + ), + Divergence( + session_id="s-ver", + kind="EVALUATOR_VERSION_MISMATCH", + detail="requested evaluator_version 3; baseline=2; candidate=3", + number=2, + ), + Divergence( + session_id="s-score", + kind="SCORE_UNCLASSIFIED", + detail="scores could not be classified as improved, regressed, or unchanged", + number=3, + ), + ], + replay_failures=[], + cohort_version_id=request.cohort_version_id, + agent_version_id=request.agent_version_id, + evaluator_version=str(request.evaluator_version), + baseline_prompt_hash=request.baseline_prompt_hash, + candidate_prompt_hash=request.candidate_prompt_hash, + verification_fingerprint="fp", + experiment_run_id="erun-1", + ) + + class DivergedBackend(FakeBackend): + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + self.collected += 1 + return result + + stored = run_verification(tmp_path, request, DivergedBackend()) + assert stored.replay_failures == [] + assert {item.kind for item in stored.diverged_sessions} == { + "SELECT_EVALUATION_FAILED", + "EVALUATOR_VERSION_MISMATCH", + "SCORE_UNCLASSIFIED", + } + matched = matching_verification( + tmp_path, + request.candidate_prompt_hash, + cohort_version_id=request.cohort_version_id, + ) + assert matched is not None + assert matched.result is not None + assert matched.per_session["s-select"] == "SELECT_EVALUATION_FAILED" + assert matched.per_session["s-ver"] == "EVALUATOR_VERSION_MISMATCH" + assert matched.per_session["s-score"] == "SCORE_UNCLASSIFIED" + assert matched.per_session["s-improved"] == "improved" + report = format_verification_report(stored, request) + assert "SELECT_EVALUATION_FAILED" in report + assert "EVALUATOR_VERSION_MISMATCH" in report + assert "SCORE_UNCLASSIFIED" in report + assert "judge-score-unavailable" in report + assert stored.baseline_count == 4 + assert stored.candidate_mean_score == 0.4 + + +def test_interrupted_verify_does_not_duplicate_the_experiment(tmp_path: Path) -> None: + backend = FakeBackend() + request = _request() + first = run_verification(tmp_path, request, backend) + assert backend.submitted == 1 + second = run_verification(tmp_path, request, backend) + assert backend.submitted == 1 + assert backend.collected == 2 + assert first.experiment_run_id == second.experiment_run_id == "erun-1" + assert ( + matching_verification( + tmp_path, + request.candidate_prompt_hash, + cohort_version_id=request.cohort_version_id, + ) + is not None + ) + + +def test_corrupt_verification_state_is_treated_as_missing_and_stays_gated( + tmp_path: Path, +) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + request = _request(candidate_prompt=written, candidate_prompt_hash=digest) + run_verification(tmp_path, request, FakeBackend()) + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + layout = initialize(tmp_path) + vid = verification_id_for("run-0001", digest) + path = verification_path(layout, vid) + + for payload in ("not-json", "{}"): + path.write_text(payload, encoding="utf-8") + assert load_verification_state(layout, vid) is None + assert matching_verification(tmp_path, digest, cohort_version_id="cv1") is None + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + backend = FakeBackend() + run_verification(tmp_path, request, backend) + assert backend.submitted == 1 + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + +def test_apply_is_gated_on_a_matching_hash(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + about_to_write = candidate_prompt( + PROMPT, proposal, range(len(proposal.edits)) + ) + digest = text_hash(about_to_write) + with pytest.raises(VerifyError, match="hash-matching verification"): + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=True + ) + + +def test_partial_apply_after_full_verify_needs_force_or_all(tmp_path: Path) -> None: + proposal = _multi_edit_proposal(tmp_path) + _source_sidecar(tmp_path) + full = candidate_prompt(PROMPT, proposal, range(len(proposal.edits))) + subset = candidate_prompt(PROMPT, proposal, [1]) + assert text_hash(full) != text_hash(subset) + run_verification( + tmp_path, + _request(candidate_prompt=full, candidate_prompt_hash=text_hash(full)), + FakeBackend(), + ) + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=text_hash(full), force=False + ) + with pytest.raises(VerifyError, match="apply --all") as caught: + refuse_ungated_apply( + tmp_path, + run_id="run-0001", + candidate_prompt_hash=text_hash(subset), + force=False, + ) + message = str(caught.value) + assert "--force" in message + assert "re-running verify cannot ungate" in message + assert "Run `tracegrad verify --backend kitaru` first" not in message + + +def test_apply_gate_does_not_affect_core_only_runs(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + about_to_write = candidate_prompt(PROMPT, proposal, [0]) + refuse_ungated_apply( + tmp_path, + run_id="run-0001", + candidate_prompt_hash=text_hash(about_to_write), + force=False, + ) + result = apply_proposal(tmp_path, proposal, [0], base_directory=tmp_path) + assert "Always cite the doc." in (tmp_path / "prompt.md").read_text(encoding="utf-8") + assert result.accepted + + +def test_hand_edit_after_verify_misses_the_gate(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + backend = FakeBackend() + written = candidate_prompt(PROMPT, proposal, [0]) + request = _request( + candidate_prompt=written, + candidate_prompt_hash=text_hash(written), + ) + run_verification(tmp_path, request, backend) + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=text_hash(written), force=False + ) + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, + run_id="run-0001", + candidate_prompt_hash=text_hash(written + "\n# hand edit\n"), + force=False, + ) + + +def test_apply_gate_requires_a_stored_result(tmp_path: Path) -> None: + """A submit that never stored result must not ungate apply.""" + + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + + class SubmitOnlyBackend(FakeBackend): + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + raise RuntimeError("collect exploded") + + request = _request(candidate_prompt=written, candidate_prompt_hash=digest) + with pytest.raises(RuntimeError, match="collect exploded"): + run_verification(tmp_path, request, SubmitOnlyBackend()) + assert matching_verification(tmp_path, digest, cohort_version_id="cv1") is None + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + +def test_collect_fetch_error_does_not_persist_result_or_ungate_apply(tmp_path: Path) -> None: + """A session_nodes/evaluations 404 or timeout aborts collect fail-closed.""" + + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + request = _request(candidate_prompt=written, candidate_prompt_hash=digest) + + def _root(system: str) -> SimpleNamespace: + return SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": system}, + system_prompt_selector="/system", + ) + + class FetchErrorGateway: + async def wait_for_experiment_run( + self, run_id: str, timeout: float | None = None + ) -> object: + return SimpleNamespace(status="completed") + + async def list_replays(self, experiment_run_id: str) -> list[object]: + return [ + SimpleNamespace( + baseline_session_id="s0", + result_session_id="r0", + status="completed", + error=None, + ), + SimpleNamespace( + baseline_session_id="s1", + result_session_id="r1", + status="completed", + error=None, + ), + ] + + async def session_nodes(self, session_id: str) -> tuple[object, ...]: + system = CANDIDATE if session_id.startswith("r") else PROMPT + return (_root(system),) + + async def evaluations_for(self, session_id: str) -> list[object]: + if session_id == "r1": + raise TimeoutError("404/timeout fetching evaluations for r1") + passed = session_id.startswith("r") + return [ + SimpleNamespace( + name="quality", + evaluator_version=3, + passed=passed, + score=0.9 if passed else 0.2, + id=session_id, + ) + ] + + async def evaluation_aggregates(self, experiment_run_id: str) -> list[object]: + return [ + { + "name": "quality", + "baseline": {"count": 2, "mean": 0.2, "pass_rate": 0.0}, + "result": {"count": 2, "mean": 0.9, "pass_rate": 1.0}, + } + ] + + class FetchErrorBackend: + name = "kitaru" + + def preflight(self, request: VerificationRequest) -> None: + return None + + def submit(self, request: VerificationRequest) -> SubmittedVerification: + return SubmittedVerification(experiment_id="exp-1", experiment_run_id="erun-1") + + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + return KitaruVerificationBackend(gateway=FetchErrorGateway()).collect( + request, submitted + ) + + with pytest.raises(KitaruVerifyError, match="collect aborted") as caught: + run_verification(tmp_path, request, FetchErrorBackend()) + message = str(caught.value) + assert "Apply stays gated" in message + assert "Resume redoes collect" in message + assert "--force" in message + assert "TimeoutError" in message + assert caught.value.__cause__ is not None + + assert matching_verification(tmp_path, digest, cohort_version_id="cv1") is None + state = load_verification_state( + initialize(tmp_path), verification_id_for("run-0001", digest) + ) + assert state is not None + assert state.experiment_run_id == "erun-1" + assert state.result is None + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + +def _collect_sdk_error_gateway(broken: str) -> object: + def _root(system: str) -> SimpleNamespace: + return SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": system}, + system_prompt_selector="/system", + ) + + class Gateway: + async def wait_for_experiment_run( + self, run_id: str, timeout: float | None = None + ) -> object: + if broken == "wait": + raise TimeoutError("wait timed out") + return SimpleNamespace(status="completed") + + async def list_replays(self, experiment_run_id: str) -> list[object]: + if broken == "replays": + raise RuntimeError("APIError listing replays") + return [ + SimpleNamespace( + baseline_session_id="s0", + result_session_id="r0", + status="completed", + error=None, + ) + ] + + async def session_nodes(self, session_id: str) -> tuple[object, ...]: + system = CANDIDATE if session_id.startswith("r") else PROMPT + return (_root(system),) + + async def evaluations_for(self, session_id: str) -> list[object]: + passed = session_id.startswith("r") + return [ + SimpleNamespace( + name="quality", + evaluator_version=3, + passed=passed, + score=0.9 if passed else 0.2, + id=session_id, + ) + ] + + async def evaluation_aggregates(self, experiment_run_id: str) -> list[object]: + if broken == "aggregates": + raise TimeoutError("aggregates 404") + return [ + { + "name": "quality", + "baseline": {"count": 1, "mean": 0.2, "pass_rate": 0.0}, + "result": {"count": 1, "mean": 0.9, "pass_rate": 1.0}, + } + ] + + return Gateway() + + +@pytest.mark.parametrize("broken", ["wait", "replays", "aggregates"]) +def test_collect_sdk_errors_do_not_persist_result_or_ungate_apply( + tmp_path: Path, broken: str +) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + request = _request(candidate_prompt=written, candidate_prompt_hash=digest) + + class Backend: + name = "kitaru" + + def preflight(self, request: VerificationRequest) -> None: + return None + + def submit(self, request: VerificationRequest) -> SubmittedVerification: + return SubmittedVerification(experiment_id="exp-1", experiment_run_id="erun-1") + + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + return KitaruVerificationBackend( + gateway=_collect_sdk_error_gateway(broken) + ).collect(request, submitted) + + with pytest.raises(KitaruVerifyError, match="collect aborted") as caught: + run_verification(tmp_path, request, Backend()) + assert "Apply stays gated" in str(caught.value) + assert matching_verification(tmp_path, digest, cohort_version_id="cv1") is None + state = load_verification_state( + initialize(tmp_path), verification_id_for("run-0001", digest) + ) + assert state is not None + assert state.experiment_run_id == "erun-1" + assert state.result is None + + +def test_submit_sdk_errors_are_kitaru_verify_errors() -> None: + pytest.importorskip("kitaru") + + class Gateway: + async def create_experiment(self, request: object) -> object: + raise RuntimeError("APIError creating experiment") + + async def start_run(self, experiment_id: str, request: object) -> object: + raise AssertionError("start_run must not run after create_experiment fails") + + backend = KitaruVerificationBackend(gateway=Gateway()) + request = _request( + agent_id="aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", + agent_version_id="bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", + cohort_version_id="cccccccc-cccc-4ccc-8ccc-cccccccccccc", + ) + with pytest.raises(KitaruVerifyError, match="submit aborted") as caught: + backend.submit(request) + assert "creating the experiment" in str(caught.value) + assert "Apply stays gated" in str(caught.value) + + +def test_apply_gate_accepts_a_finished_failed_report(tmp_path: Path) -> None: + """A finished REVIEW/FAILED report still allows apply; do not require completed.""" + + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + + class FailedBackend(FakeBackend): + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + result = super().collect(request, submitted) + return result.model_copy(update={"status": "failed"}) + + request = _request(candidate_prompt=written, candidate_prompt_hash=digest) + run_verification(tmp_path, request, FailedBackend()) + matched = matching_verification(tmp_path, digest, cohort_version_id="cv1") + assert matched is not None + assert matched.result is not None + assert matched.result.status == "failed" + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + +def test_same_hash_and_same_cohort_ungates_without_matching_run_id(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path, run_id="run-0001") + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + run_verification( + tmp_path, + _request(candidate_prompt=written, candidate_prompt_hash=digest), + FakeBackend(), + ) + _source_sidecar(tmp_path, run_id="run-0002") + refuse_ungated_apply( + tmp_path, run_id="run-0002", candidate_prompt_hash=digest, force=False + ) + matched = matching_verification(tmp_path, digest, cohort_version_id="cv1") + assert matched is not None + assert matched.run_id == "run-0001" + + +def test_same_hash_and_different_cohort_stays_gated(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path, run_id="run-0001") + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + run_verification( + tmp_path, + _request(candidate_prompt=written, candidate_prompt_hash=digest), + FakeBackend(), + ) + _source_sidecar(tmp_path, run_id="run-0002", cohort_version_id="cv-new") + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, run_id="run-0002", candidate_prompt_hash=digest, force=False + ) + assert matching_verification(tmp_path, digest, cohort_version_id="cv-new") is None + assert matching_verification(tmp_path, digest, cohort_version_id="cv1") is not None + + +def test_verification_id_is_path_safe() -> None: + vid = verification_id_for("run-0001", "sha256:abcdef1234567890") + assert vid.startswith("verify-run-0001-") + assert "/" not in vid + + +def test_build_request_refuses_a_stale_proposal_without_submitting(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + source = load_run_source_payload(tmp_path, "run-0001") + assert source is not None + (tmp_path / "prompt.md").write_text(PROMPT + "- Edited by hand.\n", encoding="utf-8") + backend = FakeBackend() + with pytest.raises(StaleProposalError, match="stale"): + build_request( + project_root=tmp_path, + run_id="run-0001", + proposal=proposal, + base_directory=tmp_path, + source=source, + ) + assert backend.submitted == 0 + assert backend.preflighted == 0 + layout = initialize(tmp_path) + assert list(layout.verification.glob("*")) == [] + + +def test_build_request_wraps_unreadable_template_as_verify_error( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + source = load_run_source_payload(tmp_path, "run-0001") + assert source is not None + (tmp_path / "prompt.md").unlink() + monkeypatch.setattr("tracegrad.verify.is_stale", lambda *args, **kwargs: False) + with pytest.raises(VerifyError, match="could not read template"): + build_request( + project_root=tmp_path, + run_id="run-0001", + proposal=proposal, + base_directory=tmp_path, + source=source, + ) + + +def test_corrupt_run_source_sidecar_does_not_traceback_or_ungate(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + sidecar = tmp_path / ".tracegrad" / "runs" / "run-0001" / "kitaru-source.json" + sidecar.write_text("not-json", encoding="utf-8") + + assert load_run_source_payload(tmp_path, "run-0001") is None + assert backend_is_configured(tmp_path, "run-0001") is True + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, + run_id="run-0001", + candidate_prompt_hash=text_hash(candidate_prompt(PROMPT, proposal, [0])), + force=False, + ) + + +def test_incomplete_run_source_sidecar_is_rejected_and_stays_gated(tmp_path: Path) -> None: + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + sidecar = tmp_path / ".tracegrad" / "runs" / "run-0001" / "kitaru-source.json" + digest = text_hash(candidate_prompt(PROMPT, proposal, [0])) + + for payload in ({}, {"fingerprint": {}, "meta": {}}, {"fingerprint": {"cohort_id": "c1"}, "meta": {}}): + sidecar.write_text(json.dumps(payload), encoding="utf-8") + assert load_run_source_payload(tmp_path, "run-0001") is None + assert backend_is_configured(tmp_path, "run-0001") is True + with pytest.raises(VerifyError, match="hash-matching"): + refuse_ungated_apply( + tmp_path, run_id="run-0001", candidate_prompt_hash=digest, force=False + ) + + sidecar.write_text(json.dumps({"fingerprint": "nope", "meta": {}}), encoding="utf-8") + assert load_run_source_payload(tmp_path, "run-0001") is None + with pytest.raises(VerifyError, match="usable Kitaru sidecar"): + build_request( + project_root=tmp_path, + run_id="run-0001", + proposal=proposal, + base_directory=tmp_path, + source={"fingerprint": "nope", "meta": {}}, + ) + + +def test_verify_cli_refuses_a_stale_proposal_before_kitaru(tmp_path: Path) -> None: + from tracegrad import cli + + _proposal(tmp_path) + _source_sidecar(tmp_path) + (tmp_path / "prompt.md").write_text(PROMPT + "- Edited by hand.\n", encoding="utf-8") + stream = io.StringIO() + code = cli.main( + [ + "verify", + "--backend", + "kitaru", + "--run", + "run-0001", + "--project-root", + str(tmp_path), + "--base-directory", + str(tmp_path), + ], + out=stream, + ) + assert code == 1 + output = stream.getvalue() + assert "stale" in output + assert "re-run tracegrad" in output + assert "prompt.md changed since run run-0001" in output + layout = initialize(tmp_path) + assert list(layout.verification.glob("*")) == [] + assert any(record.get("event") == "stale" for record in applied_history(tmp_path)) + + +@pytest.mark.parametrize( + ("status", "code"), + [ + ("completed", 0), + ("partial", 1), + ("canceled", 1), + ("failed", 1), + ("incomplete", 1), + ], +) +def test_verify_exit_code_is_zero_only_for_completed(status: str, code: int) -> None: + from tracegrad.cli import verify_exit_code + + assert verify_exit_code(status) == code + + +@pytest.mark.parametrize( + ("status", "code"), + [ + ("completed", 0), + ("partial", 1), + ("failed", 1), + ], +) +def test_verify_cli_exits_zero_only_when_status_is_completed( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, status: str, code: int +) -> None: + from tracegrad import cli + + proposal = _proposal(tmp_path) + _source_sidecar(tmp_path) + written = candidate_prompt(PROMPT, proposal, [0]) + digest = text_hash(written) + + def fake_run(project_root: object, request: VerificationRequest, backend: object) -> VerificationResult: + return VerificationResult( + status=status, # type: ignore[arg-type] + baseline_count=1, + candidate_count=1, + improved_sessions=[], + regressed_sessions=[], + unchanged_sessions=[], + diverged_sessions=[], + replay_failures=[], + cohort_version_id=request.cohort_version_id, + agent_version_id=request.agent_version_id, + evaluator_version=str(request.evaluator_version), + baseline_prompt_hash=request.baseline_prompt_hash, + candidate_prompt_hash=digest, + verification_fingerprint="fp", + experiment_run_id="erun-1", + ) + + monkeypatch.setattr("tracegrad.verify.run_verification", fake_run) + monkeypatch.setattr( + "tracegrad.integrations.kitaru.backend.KitaruVerificationBackend", + lambda *args, **kwargs: object(), + ) + stream = io.StringIO() + assert ( + cli.main( + [ + "verify", + "--backend", + "kitaru", + "--run", + "run-0001", + "--project-root", + str(tmp_path), + "--base-directory", + str(tmp_path), + ], + out=stream, + ) + == code + ) + assert "tracegrad verification" in stream.getvalue() + + +def test_report_lists_replay_failures_next_to_divergence() -> None: + request = _request(session_numbers={"crash-session": 7}) + result = VerificationResult( + status="partial", + baseline_count=2, + candidate_count=1, + improved_sessions=[], + regressed_sessions=[], + unchanged_sessions=[], + diverged_sessions=[ + Divergence( + session_id="miss-session", + kind="TOOL_HISTORY_MISS", + detail="No history result for tool 'search'", + number=3, + ) + ], + replay_failures=[ + ReplayFailure( + session_id="crash-session", + error="worker process exited 1", + number=7, + ) + ], + cohort_version_id=request.cohort_version_id, + agent_version_id=request.agent_version_id, + evaluator_version=str(request.evaluator_version), + baseline_prompt_hash=request.baseline_prompt_hash, + candidate_prompt_hash=request.candidate_prompt_hash, + verification_fingerprint="fp", + experiment_run_id="erun-1", + ) + report = format_verification_report(result, request) + assert "Replay failures 1" in report + assert "Diverged 1" in report + assert "Replay failures" in report.split("Divergence", 1)[1] + assert "#7 worker process exited 1" in report + assert "TOOL_HISTORY_MISS" in report + + +def test_second_run_async_does_not_reuse_a_closed_loop_client(monkeypatch: pytest.MonkeyPatch) -> None: + from tracegrad.integrations.kitaru import backend as backend_mod + + created: list[object] = [] + + class LoopBoundGateway: + def __init__(self, *args: object, **kwargs: object) -> None: + self._loop: asyncio.AbstractEventLoop | None = None + self._closed = False + created.append(self) + + async def probe(self) -> str: + loop = asyncio.get_running_loop() + if self._closed or ( + self._loop is not None and (self._loop is not loop or self._loop.is_closed()) + ): + raise RuntimeError("reused client bound to a closed loop") + self._loop = loop + return "ok" + + async def close(self) -> None: + self._closed = True + + monkeypatch.setattr(backend_mod, "require_kitaru", lambda: None) + monkeypatch.setattr(backend_mod, "KitaruGateway", LoopBoundGateway) + + backend = backend_mod.KitaruVerificationBackend() + + async def probe() -> str: + return await backend._gw().probe() + + assert backend._run_async(probe()) == "ok" + assert backend._run_async(probe()) == "ok" + assert len(created) == 2 + assert created[0] is not created[1] + assert created[0]._closed is True + assert created[1]._closed is True + + +def test_collect_fetches_replay_payloads_with_bounded_concurrency() -> None: + n_success = FETCH_JOBS + 4 + fetched: list[str] = [] + + class CountingGateway: + def __init__(self) -> None: + self.http_in_flight = 0 + self.max_http_in_flight = 0 + self._replay_counts: dict[str, int] = {} + self.max_replay_in_flight = 0 + self._lock = asyncio.Lock() + + def _replay_key(self, session_id: str) -> str: + if session_id.startswith("r") and session_id[1:].isdigit(): + return f"s{session_id[1:]}" + return session_id + + async def _track(self, session_id: str) -> None: + if session_id in {"fail-session", "hist-session"}: + raise AssertionError(f"failed replay {session_id} must not be fetched") + key = self._replay_key(session_id) + async with self._lock: + self.http_in_flight += 1 + self.max_http_in_flight = max(self.max_http_in_flight, self.http_in_flight) + self._replay_counts[key] = self._replay_counts.get(key, 0) + 1 + self.max_replay_in_flight = max( + self.max_replay_in_flight, len(self._replay_counts) + ) + fetched.append(session_id) + try: + await asyncio.sleep(0.05) + finally: + async with self._lock: + self.http_in_flight -= 1 + self._replay_counts[key] -= 1 + if self._replay_counts[key] == 0: + del self._replay_counts[key] + + async def wait_for_experiment_run( + self, run_id: str, timeout: float | None = None + ) -> object: + return SimpleNamespace(status="completed") + + async def list_replays(self, experiment_run_id: str) -> list[object]: + rows: list[object] = [ + SimpleNamespace( + baseline_session_id="fail-session", + result_session_id=None, + status="failed", + error="worker crashed", + ), + SimpleNamespace( + baseline_session_id="hist-session", + result_session_id=None, + status="failed", + error="No history result for tool 'search'", + ), + ] + for index in range(n_success): + rows.append( + SimpleNamespace( + baseline_session_id=f"s{index}", + result_session_id=f"r{index}", + status="completed", + error=None, + ) + ) + return rows + + async def session_nodes(self, session_id: str) -> tuple[object, ...]: + await self._track(session_id) + system = CANDIDATE if session_id.startswith("r") else PROMPT + return ( + SimpleNamespace( + index=0, + parent_index=None, + secondary_parent_indexes=[], + node_type="llm_call", + inputs={"system": system}, + system_prompt_selector="/system", + ), + ) + + async def evaluations_for(self, session_id: str) -> list[object]: + await self._track(session_id) + passed = session_id.startswith("r") + return [ + SimpleNamespace( + name="quality", + evaluator_version=3, + passed=passed, + score=0.9 if passed else 0.2, + id=session_id, + ) + ] + + async def evaluation_aggregates(self, experiment_run_id: str) -> list[object]: + return [ + { + "name": "quality", + "baseline": {"count": n_success, "mean": 0.2, "pass_rate": 0.0}, + "result": {"count": n_success, "mean": 0.9, "pass_rate": 1.0}, + } + ] + + gateway = CountingGateway() + backend = KitaruVerificationBackend(gateway=gateway) + request = _request(session_numbers={f"s{index}": index for index in range(n_success)}) + result = backend.collect( + request, SubmittedVerification(experiment_id="exp-1", experiment_run_id="erun-1") + ) + + assert gateway.max_replay_in_flight > 1 + assert gateway.max_replay_in_flight <= FETCH_JOBS + assert gateway.max_http_in_flight > 1 + assert result.improved_sessions == [f"s{index}" for index in range(n_success)] + assert result.replay_failures == [ + ReplayFailure(session_id="fail-session", error="worker crashed", number=None) + ] + assert result.diverged_sessions[0].kind == DIVERGENCE_HISTORY + assert result.diverged_sessions[0].session_id == "hist-session" + assert "fail-session" not in fetched + assert "hist-session" not in fetched + assert result.baseline_count == n_success + assert result.candidate_count == n_success