From 2fb20e22d6498f638798c4eeb749448609d6d63c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 12:14:24 +0000 Subject: [PATCH 01/23] feat(kitaru): optional trace source and replay verification Land ADRs 0001-0010 and CONTEXT.md, add the kitaru extra (pinned >=0.22,<0.23), map frozen cohorts onto the JSONL ingest already reads, and gate apply on a hash-matching verification. Co-authored-by: Dickson Neoh --- CONTEXT.md | 139 ++++++ README.md | 87 +++- ...ubmits-verification-it-does-not-host-it.md | 49 ++ ...02-prompt-lineage-from-recorded-prompts.md | 44 ++ ...-judge-scores-are-accepted-not-rescaled.md | 47 ++ ...e-kitaru-source-snapshots-mapped-traces.md | 54 +++ ...05-traces-are-extracted-from-root-nodes.md | 53 +++ ...verride-scope-is-verified-not-predicted.md | 46 ++ ...-requires-a-single-agent-version-cohort.md | 43 ++ ...-inspection-is-handed-off-by-identifier.md | 37 ++ ...ply-is-gated-on-a-matching-verification.md | 45 ++ ...0010-no-tracesource-protocol-in-phase-1.md | 43 ++ pyproject.toml | 3 + src/tracegrad/apply.py | 23 + src/tracegrad/cli.py | 173 ++++++- src/tracegrad/config.py | 10 + src/tracegrad/integrations/__init__.py | 6 + src/tracegrad/integrations/kitaru/__init__.py | 43 ++ .../integrations/kitaru/accounting.py | 32 ++ src/tracegrad/integrations/kitaru/backend.py | 391 ++++++++++++++++ src/tracegrad/integrations/kitaru/client.py | 301 +++++++++++++ src/tracegrad/integrations/kitaru/errors.py | 47 ++ src/tracegrad/integrations/kitaru/graph.py | 96 ++++ src/tracegrad/integrations/kitaru/mapping.py | 270 +++++++++++ src/tracegrad/integrations/kitaru/pointer.py | 66 +++ src/tracegrad/integrations/kitaru/policy.py | 43 ++ src/tracegrad/integrations/kitaru/require.py | 24 + src/tracegrad/integrations/kitaru/scores.py | 134 ++++++ src/tracegrad/integrations/kitaru/snapshot.py | 193 ++++++++ src/tracegrad/integrations/kitaru/source.py | 233 ++++++++++ src/tracegrad/ports.py | 17 + src/tracegrad/state.py | 15 +- src/tracegrad/verify.py | 422 ++++++++++++++++++ tests/test_config.py | 15 + tests/test_kitaru_mapping.py | 304 +++++++++++++ tests/test_kitaru_optional.py | 143 ++++++ tests/test_kitaru_snapshot.py | 91 ++++ tests/test_kitaru_verify.py | 364 +++++++++++++++ 38 files changed, 4133 insertions(+), 13 deletions(-) create mode 100644 CONTEXT.md create mode 100644 docs/adr/0001-tracegrad-submits-verification-it-does-not-host-it.md create mode 100644 docs/adr/0002-prompt-lineage-from-recorded-prompts.md create mode 100644 docs/adr/0003-judge-scores-are-accepted-not-rescaled.md create mode 100644 docs/adr/0004-the-kitaru-source-snapshots-mapped-traces.md create mode 100644 docs/adr/0005-traces-are-extracted-from-root-nodes.md create mode 100644 docs/adr/0006-override-scope-is-verified-not-predicted.md create mode 100644 docs/adr/0007-verification-requires-a-single-agent-version-cohort.md create mode 100644 docs/adr/0008-inspection-is-handed-off-by-identifier.md create mode 100644 docs/adr/0009-apply-is-gated-on-a-matching-verification.md create mode 100644 docs/adr/0010-no-tracesource-protocol-in-phase-1.md create mode 100644 src/tracegrad/integrations/__init__.py create mode 100644 src/tracegrad/integrations/kitaru/__init__.py create mode 100644 src/tracegrad/integrations/kitaru/accounting.py create mode 100644 src/tracegrad/integrations/kitaru/backend.py create mode 100644 src/tracegrad/integrations/kitaru/client.py create mode 100644 src/tracegrad/integrations/kitaru/errors.py create mode 100644 src/tracegrad/integrations/kitaru/graph.py create mode 100644 src/tracegrad/integrations/kitaru/mapping.py create mode 100644 src/tracegrad/integrations/kitaru/pointer.py create mode 100644 src/tracegrad/integrations/kitaru/policy.py create mode 100644 src/tracegrad/integrations/kitaru/require.py create mode 100644 src/tracegrad/integrations/kitaru/scores.py create mode 100644 src/tracegrad/integrations/kitaru/snapshot.py create mode 100644 src/tracegrad/integrations/kitaru/source.py create mode 100644 src/tracegrad/verify.py create mode 100644 tests/test_kitaru_mapping.py create mode 100644 tests/test_kitaru_optional.py create mode 100644 tests/test_kitaru_snapshot.py create mode 100644 tests/test_kitaru_verify.py diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 0000000..183913f --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,139 @@ +# 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`. Both divergence kinds +are **incomparable** — not improved, not regressed (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..d7dd3cd 100644 --- a/src/tracegrad/apply.py +++ b/src/tracegrad/apply.py @@ -277,6 +277,29 @@ def snapshot_template(project_root: str | Path, run_id: str, template: Path) -> return target +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 full proposal) and the apply gate (the selection + about to be written) so both hash the same way. + """ + + 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) + + def apply_proposal( project_root: str | Path, proposal: Proposal, diff --git a/src/tracegrad/cli.py b/src/tracegrad/cli.py index b9f3478..f4b6555 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,59 @@ 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) + 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 0 if result.status != "failed" else 1 + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( prog="tracegrad", @@ -381,7 +507,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 +569,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 +579,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 +607,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..885d421 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -0,0 +1,391 @@ +"""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 uuid +from typing import Any + +from tracegrad.verify import ( + DIVERGENCE_HISTORY, + DIVERGENCE_SCOPE, + Divergence, + ReplayFailure, + SubmittedVerification, + VerificationRequest, + VerificationResult, + verification_fingerprint_for, +) + +from .client import KitaruGateway, run_async, worker_covers_agent_version +from .errors import KitaruVerifyError +from .graph import index_nodes, is_root_llm_node, llm_nodes, node_index +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: str | None) -> bool: + if not error: + return False + text = error.lower() + needles = ( + "tool_history_miss", + "tool history miss", + "history miss", + "on_miss", + "no recorded call", + "recorded history", + ) + return any(needle in text for needle in needles) + + +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.""" + + 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_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): + if isinstance(base_passed, bool) and isinstance(cand_passed, bool): + return "unchanged" + return None + if cand_score > base_score: + return "improved" + if cand_score < base_score: + return "regressed" + return "unchanged" + + +class KitaruVerificationBackend: + """Submit and collect a Kitaru experiment run for one candidate prompt.""" + + name = "kitaru" + + def __init__(self, gateway: KitaruGateway | None = None) -> None: + require_kitaru() + self._gateway = gateway + self._owns = gateway is None + + def _gw(self) -> KitaruGateway: + if self._gateway is None: + self._gateway = KitaruGateway() + return self._gateway + + def preflight(self, request: VerificationRequest) -> None: + 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)) + workers = await gateway.list_live_workers() + 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 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 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)], + ) + ) + run = await 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, + ), + ) + return SubmittedVerification( + experiment_id=str(experiment.id), + experiment_run_id=str(run.id), + ) + + def collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + return run_async(self._collect(request, submitted)) + + async def _collect( + self, request: VerificationRequest, submitted: SubmittedVerification + ) -> VerificationResult: + gateway = self._gw() + run = await gateway.wait_for_experiment_run(submitted.experiment_run_id) + status = str(getattr(getattr(run, "status", None), "value", getattr(run, "status", "completed"))) + replays = await gateway.list_replays(submitted.experiment_run_id) + improved: list[str] = [] + regressed: list[str] = [] + unchanged: list[str] = [] + diverged: list[Divergence] = [] + failures: list[ReplayFailure] = [] + + 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 + baseline_nodes = await gateway.session_nodes(str(replay.baseline_session_id)) + result_nodes = await gateway.session_nodes(str(result_id)) + 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 + baseline_eval = select_evaluation( + await gateway.evaluations_for(str(replay.baseline_session_id)), + request.evaluation_name, + ) + candidate_eval = select_evaluation( + await gateway.evaluations_for(str(result_id)), + request.evaluation_name, + ) + if isinstance(baseline_eval, str) or isinstance(candidate_eval, str): + continue + if evaluator_version_of(baseline_eval) != request.evaluator_version: + continue + if evaluator_version_of(candidate_eval) != request.evaluator_version: + continue + verdict = classify_scores(baseline_eval, candidate_eval) + if verdict == "improved": + improved.append(session_id) + elif verdict == "regressed": + regressed.append(session_id) + elif verdict == "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 gateway.evaluation_aggregates(submitted.experiment_run_id) + 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, + ) + if self._owns and self._gateway is not None: + await self._gateway.close() + self._gateway = None + return result + + +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..321e1d1 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/client.py @@ -0,0 +1,301 @@ +"""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 uuid +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from .errors import KitaruSourceError, KitaruVerifyError +from .require import require_kitaru + + +@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 + try: + _uuid(version_ref) + except ValueError: + return None + 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 = 8 + ) -> 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)) + + return list(await asyncio.gather(*(one(session) for session in sessions))) + + 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 + + params = WorkerListParams(include_stale=False) + return [item async for item in self._client.workers.iter(params)] + + 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: + from kitaru.client.client import KitaruClient + + wrapper = KitaruClient(api_client=self._client) + return await wrapper.wait_for_experiment_run(_uuid(run_id), timeout=timeout) + + 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..2942e83 --- /dev/null +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -0,0 +1,193 @@ +"""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 ( + StateLayout, + atomic_write, + atomic_write_json, + initialize, + 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 + 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 snapshot_dir(layout: StateLayout, cohort_version_id: str) -> Path: + return kitaru_root(layout) / cohort_version_id + + +def snapshot_exists(layout: StateLayout, cohort_version_id: str) -> bool: + target = snapshot_dir(layout, cohort_version_id) + return (target / BATCH_FILENAME).is_file() and (target / FINGERPRINT_FILENAME).is_file() + + +def write_snapshot( + layout: StateLayout, + *, + fingerprint: SourceFingerprint, + meta: SourceMeta, + mapped: Sequence[MappedTrace], + dropped: Sequence[SourceDrop], +) -> Path: + """Write JSONL + fingerprint + meta under ``.tracegrad/sources/kitaru/``.""" + + target = snapshot_dir(layout, fingerprint.cohort_version_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 "")) + atomic_write_json( + kitaru_root(layout) / LATEST_POINTER, + { + "cohort_version_id": fingerprint.cohort_version_id, + "batch": str(target / BATCH_FILENAME), + }, + ) + return target + + +def load_fingerprint(layout: StateLayout, cohort_version_id: str) -> SourceFingerprint: + path = snapshot_dir(layout, cohort_version_id) / FINGERPRINT_FILENAME + return SourceFingerprint.model_validate_json(path.read_text(encoding="utf-8")) + + +def load_meta(layout: StateLayout, cohort_version_id: str) -> SourceMeta: + path = snapshot_dir(layout, cohort_version_id) / META_FILENAME + return SourceMeta.model_validate_json(path.read_text(encoding="utf-8")) + + +def load_source_drops(layout: StateLayout, cohort_version_id: str) -> tuple[SourceDrop, ...]: + path = snapshot_dir(layout, cohort_version_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, cohort_version_id: str) -> Path: + return snapshot_dir(layout, cohort_version_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 load_run_source(project_root: str | Path | StateLayout, run_id: str) -> dict[str, Any] | None: + layout = initialize(project_root) + target = layout.runs / validate_run_id(run_id) / RUN_SOURCE_FILENAME + if not target.exists(): + return None + return json.loads(target.read_text(encoding="utf-8")) + + +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 + return str(stored.cohort_version_id) == str(requested.get("cohort_version_id")) diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py new file mode 100644 index 0000000..867155e --- /dev/null +++ b/src/tracegrad/integrations/kitaru/source.py @@ -0,0 +1,233 @@ +"""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 + +from tracegrad.schema import Manifest, TemplateEngine +from tracegrad.state import 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, + map_batch, +) +from .require import require_kitaru +from .scores import judge_fingerprint_for +from .snapshot import ( + SourceFingerprint, + SourceMeta, + batch_path, + fingerprints_compatible, + load_fingerprint, + load_meta, + load_source_drops, + persist_run_source, + snapshot_exists, + 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 _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 _fetch_and_map( + *, + gateway: Any, + resolution: Any, + evaluation_name: str, +) -> tuple[Any, Any, Any, Any]: + sessions = await gateway.list_sessions(resolution.cohort_version_id) + 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 gateway.evaluator_id(evaluation_name) + 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, + 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 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.""" + + refuse_format_engine(manifest) + require_kitaru() + layout = initialize(project_root) + + from .client import KitaruGateway, run_async + + owns = gateway is None + gateway = gateway or KitaruGateway() + + async def _run() -> PreparedSource: + try: + resolution = await gateway.resolve_cohort(cohort_name, cohort_version) + requested = { + "evaluation_name": evaluation_name, + "cohort_version_id": resolution.cohort_version_id, + } + reused = ( + not refresh + and snapshot_exists(layout, resolution.cohort_version_id) + and fingerprints_compatible( + load_fingerprint(layout, resolution.cohort_version_id), requested + ) + ) + if reused: + fingerprint = load_fingerprint(layout, resolution.cohort_version_id) + meta = load_meta(layout, resolution.cohort_version_id) + dropped = load_source_drops(layout, resolution.cohort_version_id) + refreshed = False + else: + fingerprint, meta, mapped, dropped = await _fetch_and_map( + gateway=gateway, + resolution=resolution, + evaluation_name=evaluation_name, + ) + write_snapshot( + layout, fingerprint=fingerprint, meta=meta, mapped=mapped, dropped=dropped + ) + refreshed = True + + if meta.traces_mapped: + derived = judge_fingerprint_for( + meta.evaluator_name or evaluation_name, + fingerprint.evaluator_version, + ) + check_judge_fingerprint(manifest, derived) + + 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 reused else 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, fingerprint.cohort_version_id), + fingerprint=fingerprint, + meta=meta, + source_table=table, + refreshed=refreshed, + ) + 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..9dfa28c --- /dev/null +++ b/src/tracegrad/verify.py @@ -0,0 +1,422 @@ +"""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`` 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, candidate_prompt, 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" + +DIVERGENCE_HISTORY = "TOOL_HISTORY_MISS" +DIVERGENCE_SCOPE = "OVERRIDE_SCOPE_DIVERGENCE" +VERIFICATION_FILENAME = "state.json" + + +class VerifyError(ValueError): + """Verification cannot start or cannot gate apply.""" + + +class Divergence(BaseModel): + model_config = ConfigDict(extra="forbid") + + session_id: StrictStr + kind: Literal["TOOL_HISTORY_MISS", "OVERRIDE_SCOPE_DIVERGENCE"] + 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: + target = verification_path(layout, verification_id) + if not target.exists(): + return None + return VerificationState.model_validate_json(target.read_text(encoding="utf-8")) + + +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, +) -> VerificationState | None: + """A persisted verification of exactly this candidate text, if any.""" + + layout = initialize(project_root) + for state in list_verification_states(layout): + if state.candidate_prompt_hash == candidate_prompt_hash and state.experiment_run_id: + return state + return None + + +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 + + return json.loads(target.read_text(encoding="utf-8")) + + +def backend_is_configured(project_root: str | Path, run_id: str) -> bool: + """Whether this run originated from a Kitaru source (ADR 0009).""" + + return load_run_source_payload(project_root, run_id) is not None + + +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.""" + + if force or not backend_is_configured(project_root, run_id): + return + if matching_verification(project_root, candidate_prompt_hash) is None: + raise VerifyError( + "apply is gated on a hash-matching verification for this candidate. " + "Run `tracegrad verify --backend kitaru` first, or pass --force " + "to override. Matching the hash (not the run id) is what makes the " + "gate real: verify, hand-edit, and apply notices the text was never " + "verified. See ADR 0009." + ) + + +def build_request( + *, + project_root: str | Path, + run_id: str, + proposal: Proposal, + base_directory: str | Path = ".", + source: dict[str, Any], +) -> VerificationRequest: + template = contained_path(base_directory, proposal.template_file) + current = template.read_text(encoding="utf-8") + candidate = candidate_prompt( + current, proposal, range(len(proposal.edits)) + ) + fingerprint = source["fingerprint"] + meta = source["meta"] + 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() + }, + ) + + +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}", + ] + 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()) + 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_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..b415559 --- /dev/null +++ b/tests/test_kitaru_mapping.py @@ -0,0 +1,304 @@ +"""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 + +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() 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..6a11d16 --- /dev/null +++ b/tests/test_kitaru_snapshot.py @@ -0,0 +1,91 @@ +"""Snapshot reuse for the Kitaru source (ADR 0004).""" + +from __future__ import annotations + +from tracegrad.canonical import text_hash +from tracegrad.integrations.kitaru.accounting import format_source_table +from tracegrad.integrations.kitaru.mapping import MappedTrace, SourceDrop +from tracegrad.integrations.kitaru.snapshot import ( + SourceFingerprint, + SourceMeta, + fingerprints_compatible, + load_fingerprint, + load_source_drops, + snapshot_exists, + write_snapshot, +) +from tracegrad.schema import Trace +from tracegrad.state import initialize + + +def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: + layout = initialize(tmp_path) + trace = Trace( + trace_id="s1", + input="q", + output="a", + judge={"score": 0.2, "rationale": "needs a citation in the answer now"}, + prompt_hash=text_hash("prompt"), + ) + mapped = [ + MappedTrace( + trace=trace, + session_number=12, + evaluator_name="quality", + evaluator_version=3, + evaluator_version_id="ev", + multi_turn=False, + system_prompt="prompt", + ) + ] + dropped = [SourceDrop("s2", "system-prompt-unavailable", number=13)] + fingerprint = SourceFingerprint( + cohort_id="c", + cohort_version_id="cv", + evaluation_name="quality", + evaluator_id="eid", + evaluator_version=3, + agent_id="a", + ) + meta = SourceMeta( + cohort_name="support-production", + traces_mapped=1, + sessions_selected=2, + evaluator_name="quality", + ) + write_snapshot(layout, fingerprint=fingerprint, meta=meta, mapped=mapped, dropped=dropped) + + assert snapshot_exists(layout, "cv") + stored = load_fingerprint(layout, "cv") + 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"} + ) + drops = load_source_drops(layout, "cv") + assert drops[0].reason == "system-prompt-unavailable" + batch = (layout.sources / "kitaru" / "cv" / "batch.jsonl").read_text(encoding="utf-8") + assert "s1" in batch + + +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") diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py new file mode 100644 index 0000000..008fe54 --- /dev/null +++ b/tests/test_kitaru_verify.py @@ -0,0 +1,364 @@ +"""Phase 2 verification: gate, resume, policy, override scope (issue #9).""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tracegrad.apply import ( + Proposal, + ProposedEdit, + apply_proposal, + candidate_prompt, + save_proposal, +) +from tracegrad.canonical import text_hash +from tracegrad.edits import resolve_edits +from tracegrad.integrations.kitaru.backend import ( + assert_override_scope, + classify_scores, + is_tool_history_miss, + mixed_agent_version_message, +) +from tracegrad.integrations.kitaru.policy import ( + RECORDED_HISTORY_POLICY, + asserts_no_passthrough, + recorded_history_policy, +) +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 ( + SubmittedVerification, + VerificationRequest, + VerificationResult, + VerifyError, + matching_verification, + refuse_ungated_apply, + run_verification, + verification_id_for, +) + +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 _source_sidecar(project: Path, run_id: str = "run-0001") -> None: + layout = initialize(project) + atomic_write_json( + layout.runs / run_id / RUN_SOURCE_FILENAME, + { + "fingerprint": { + "source": "kitaru", + "cohort_id": "c1", + "cohort_version_id": "cv1", + "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("TOOL_HISTORY_MISS: search_account") + assert is_tool_history_miss("no recorded call matched under on_miss=fail") + assert not is_tool_history_miss("evaluator crashed") + + +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_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_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 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) is not None + + +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_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_verification_id_is_path_safe() -> None: + vid = verification_id_for("run-0001", "sha256:abcdef1234567890") + assert vid.startswith("verify-run-0001-") + assert "/" not in vid From 3dcd4e283c25f8d35b860b6618d5c6f54a3846f4 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 12:33:34 +0000 Subject: [PATCH 02/23] fix(kitaru): snapshot reuse, real history-miss match, apply result gate Match ToolPolicyMissError / "No history result for tool" so those replays are TOOL_HISTORY_MISS, not replay_failures. Read latest.json and reuse the local snapshot unless --refresh so re-runs stay offline. Require state.result before apply; a finished REVIEW/FAILED report still gates. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 30 ++- src/tracegrad/integrations/kitaru/snapshot.py | 112 ++++++++- src/tracegrad/integrations/kitaru/source.py | 147 ++++++----- src/tracegrad/verify.py | 13 +- tests/test_kitaru_snapshot.py | 230 +++++++++++++++++- tests/test_kitaru_verify.py | 61 ++++- 6 files changed, 510 insertions(+), 83 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index 885d421..dd4a88a 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -73,19 +73,25 @@ def _evaluator_config(request: VerificationRequest) -> Any: ) -def is_tool_history_miss(error: str | None) -> bool: - if not error: +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 - text = error.lower() - needles = ( - "tool_history_miss", - "tool history miss", - "history miss", - "on_miss", - "no recorded call", - "recorded history", - ) - return any(needle in text for needle in needles) + 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: diff --git a/src/tracegrad/integrations/kitaru/snapshot.py b/src/tracegrad/integrations/kitaru/snapshot.py index 2942e83..b5deb62 100644 --- a/src/tracegrad/integrations/kitaru/snapshot.py +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -115,6 +115,9 @@ def write_snapshot( { "cohort_version_id": fingerprint.cohort_version_id, "batch": str(target / BATCH_FILENAME), + "cohort_name": meta.cohort_name, + "evaluation_name": fingerprint.evaluation_name, + "display_version": meta.display_version, }, ) return target @@ -190,4 +193,111 @@ def fingerprints_compatible(stored: SourceFingerprint, requested: Mapping[str, A return False if stored.evaluation_name != requested.get("evaluation_name"): return False - return str(stored.cohort_version_id) == str(requested.get("cohort_version_id")) + 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 load_latest_pointer(layout: StateLayout) -> dict[str, Any] | None: + """Read ``latest.json`` if a previous source write left one.""" + + 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 list_snapshot_ids(layout: StateLayout) -> list[str]: + root = kitaru_root(layout) + if not root.is_dir(): + return [] + return [ + child.name + for child in sorted(root.iterdir()) + if child.is_dir() and snapshot_exists(layout, child.name) + ] + + +def snapshot_matches_request( + layout: StateLayout, + cohort_version_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, cohort_version_id): + return False + fingerprint = load_fingerprint(layout, cohort_version_id) + if not fingerprints_compatible(fingerprint, {"evaluation_name": evaluation_name}): + return False + meta = load_meta(layout, cohort_version_id) + if meta.cohort_name != cohort_name: + return False + if cohort_version is None: + return True + return cohort_version_id == cohort_version or meta.display_version == cohort_version + + +def find_local_snapshot( + layout: StateLayout, + *, + cohort_name: str, + evaluation_name: str, + cohort_version: str | None = None, +) -> str | None: + """Return a local ``cohort_version_id`` that satisfies this request. + + Prefers ``latest.json`` when it matches so a re-run keeps the pinned + version even if the server's latest has moved. Does not contact the + server. ``--refresh`` is the caller's decision not to call this. + """ + + if cohort_version is not None: + if snapshot_matches_request( + layout, + cohort_version, + cohort_name=cohort_name, + evaluation_name=evaluation_name, + cohort_version=cohort_version, + ): + return cohort_version + for snapshot_id in list_snapshot_ids(layout): + if snapshot_matches_request( + layout, + snapshot_id, + cohort_name=cohort_name, + evaluation_name=evaluation_name, + cohort_version=cohort_version, + ): + return snapshot_id + return None + + pointer = load_latest_pointer(layout) + if pointer is not None: + latest_id = str(pointer.get("cohort_version_id") or "") + if latest_id and snapshot_matches_request( + layout, + latest_id, + cohort_name=cohort_name, + evaluation_name=evaluation_name, + ): + return latest_id + for snapshot_id in list_snapshot_ids(layout): + if snapshot_matches_request( + layout, + snapshot_id, + cohort_name=cohort_name, + evaluation_name=evaluation_name, + ): + return snapshot_id + return None diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py index 867155e..3f03b5b 100644 --- a/src/tracegrad/integrations/kitaru/source.py +++ b/src/tracegrad/integrations/kitaru/source.py @@ -8,10 +8,10 @@ from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Sequence from tracegrad.schema import Manifest, TemplateEngine -from tracegrad.state import initialize +from tracegrad.state import StateLayout, initialize from .accounting import format_source_table from .errors import KitaruSourceError @@ -19,6 +19,7 @@ REASON_AMBIGUOUS_EVALUATION, REASON_FORMAT_ENGINE_REFUSED, REASON_JUDGE_FINGERPRINT_CONFLICT, + SourceDrop, map_batch, ) from .require import require_kitaru @@ -27,12 +28,11 @@ SourceFingerprint, SourceMeta, batch_path, - fingerprints_compatible, + find_local_snapshot, load_fingerprint, load_meta, load_source_drops, persist_run_source, - snapshot_exists, write_snapshot, ) @@ -147,6 +147,44 @@ async def _fetch_and_map( return fingerprint, meta, mapping.mapped, mapping.dropped +def _assembled_source( + layout: StateLayout, + *, + fingerprint: SourceFingerprint, + meta: SourceMeta, + dropped: Sequence[SourceDrop], + refreshed: bool, + manifest: Manifest, + evaluation_name: str, + run_id: str | None, +) -> PreparedSource: + if meta.traces_mapped: + derived = judge_fingerprint_for( + meta.evaluator_name or evaluation_name, + fingerprint.evaluator_version, + ) + check_judge_fingerprint(manifest, derived) + 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, fingerprint.cohort_version_id), + fingerprint=fingerprint, + meta=meta, + source_table=table, + refreshed=refreshed, + ) + + def prepare_kitaru_source( *, project_root: str | Path, @@ -158,73 +196,68 @@ def prepare_kitaru_source( gateway: Any | None = None, run_id: str | None = None, ) -> PreparedSource: - """Resolve, snapshot, and return the JSONL path ingest already reads.""" + """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) - require_kitaru() layout = initialize(project_root) - from .client import KitaruGateway, run_async + 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, + 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 - gateway = gateway or KitaruGateway() + if owns: + require_kitaru() + from .client import KitaruGateway + + gateway = KitaruGateway() async def _run() -> PreparedSource: try: resolution = await gateway.resolve_cohort(cohort_name, cohort_version) - requested = { - "evaluation_name": evaluation_name, - "cohort_version_id": resolution.cohort_version_id, - } - reused = ( - not refresh - and snapshot_exists(layout, resolution.cohort_version_id) - and fingerprints_compatible( - load_fingerprint(layout, resolution.cohort_version_id), requested - ) + fingerprint, meta, mapped, dropped = await _fetch_and_map( + gateway=gateway, + resolution=resolution, + evaluation_name=evaluation_name, ) - if reused: - fingerprint = load_fingerprint(layout, resolution.cohort_version_id) - meta = load_meta(layout, resolution.cohort_version_id) - dropped = load_source_drops(layout, resolution.cohort_version_id) - refreshed = False - else: - fingerprint, meta, mapped, dropped = await _fetch_and_map( - gateway=gateway, - resolution=resolution, - evaluation_name=evaluation_name, - ) - write_snapshot( - layout, fingerprint=fingerprint, meta=meta, mapped=mapped, dropped=dropped - ) - refreshed = True - - if meta.traces_mapped: - derived = judge_fingerprint_for( - meta.evaluator_name or evaluation_name, - fingerprint.evaluator_version, - ) - check_judge_fingerprint(manifest, derived) - - 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 reused else dropped, + write_snapshot( + layout, fingerprint=fingerprint, meta=meta, mapped=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, fingerprint.cohort_version_id), + return _assembled_source( + layout, fingerprint=fingerprint, meta=meta, - source_table=table, - refreshed=refreshed, + dropped=dropped, + refreshed=True, + manifest=manifest, + evaluation_name=evaluation_name, + run_id=run_id, ) finally: if owns: diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 9dfa28c..2b12666 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -171,11 +171,20 @@ def matching_verification( project_root: str | Path, candidate_prompt_hash: str, ) -> VerificationState | None: - """A persisted verification of exactly this candidate text, if any.""" + """A persisted verification of exactly this candidate text, if any. + + Requires a stored ``result``. A submit that never collected (``result`` + is ``None``) must not ungate apply. A finished REVIEW/FAILED report + still matches; status is not required to be ``completed``. + """ layout = initialize(project_root) for state in list_verification_states(layout): - if state.candidate_prompt_hash == candidate_prompt_hash and state.experiment_run_id: + if ( + state.candidate_prompt_hash == candidate_prompt_hash + and state.experiment_run_id + and state.result is not None + ): return state return None diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 6a11d16..5ab21c0 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -2,34 +2,44 @@ from __future__ import annotations +from pathlib import Path + from tracegrad.canonical import text_hash from tracegrad.integrations.kitaru.accounting import format_source_table +from tracegrad.integrations.kitaru.client import CohortResolution from tracegrad.integrations.kitaru.mapping import MappedTrace, SourceDrop from tracegrad.integrations.kitaru.snapshot import ( + LATEST_POINTER, SourceFingerprint, SourceMeta, + find_local_snapshot, fingerprints_compatible, + kitaru_root, load_fingerprint, + load_latest_pointer, load_source_drops, snapshot_exists, write_snapshot, ) -from tracegrad.schema import Trace +from tracegrad.integrations.kitaru.source import prepare_kitaru_source +from tracegrad.schema import Manifest, Trace from tracegrad.state import initialize -def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: - layout = initialize(tmp_path) - trace = Trace( +def _sample_trace() -> Trace: + return Trace( trace_id="s1", input="q", output="a", judge={"score": 0.2, "rationale": "needs a citation in the answer now"}, prompt_hash=text_hash("prompt"), ) - mapped = [ + + +def _sample_mapped() -> list[MappedTrace]: + return [ MappedTrace( - trace=trace, + trace=_sample_trace(), session_number=12, evaluator_name="quality", evaluator_version=3, @@ -38,8 +48,10 @@ def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: system_prompt="prompt", ) ] - dropped = [SourceDrop("s2", "system-prompt-unavailable", number=13)] - fingerprint = SourceFingerprint( + + +def _sample_fingerprint(**overrides: object) -> SourceFingerprint: + payload = dict( cohort_id="c", cohort_version_id="cv", evaluation_name="quality", @@ -47,12 +59,32 @@ def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: evaluator_version=3, agent_id="a", ) - meta = SourceMeta( + payload.update(overrides) + return SourceFingerprint.model_validate(payload) + + +def _sample_meta(**overrides: object) -> SourceMeta: + payload = dict( cohort_name="support-production", + display_version="week-34", traces_mapped=1, sessions_selected=2, evaluator_name="quality", ) + payload.update(overrides) + return SourceMeta.model_validate(payload) + + +def _manifest() -> Manifest: + return Manifest(template_file=Path("prompt.md"), judge_fingerprint="quality@3") + + +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, "cv") @@ -65,10 +97,22 @@ def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: 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, "cv") assert drops[0].reason == "system-prompt-unavailable" batch = (layout.sources / "kitaru" / "cv" / "batch.jsonl").read_text(encoding="utf-8") assert "s1" in batch + pointer = load_latest_pointer(layout) + assert pointer is not None + assert pointer["cohort_version_id"] == "cv" + assert pointer["cohort_name"] == "support-production" + assert pointer["evaluation_name"] == "quality" + assert pointer["display_version"] == "week-34" + assert find_local_snapshot( + layout, cohort_name="support-production", evaluation_name="quality" + ) == "cv" def test_source_and_batch_tables_are_not_merged() -> None: @@ -89,3 +133,171 @@ def test_source_and_batch_tables_are_not_merged() -> None: 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" / "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", + ) + assert by_id.fingerprint.cohort_version_id == "cv" + assert by_display.fingerprint.cohort_version_id == "cv" + assert by_id.refreshed is False + assert by_display.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) + assert pointer is not None + assert pointer["cohort_version_id"] == "cv-new" + + +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 diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 008fe54..80adc19 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -189,9 +189,16 @@ def test_passthrough_is_rejected() -> None: def test_tool_history_miss_is_typed() -> None: - assert is_tool_history_miss("TOOL_HISTORY_MISS: search_account") - assert is_tool_history_miss("no recorded call matched under on_miss=fail") + 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: @@ -358,6 +365,56 @@ def test_hand_edit_after_verify_misses_the_gate(tmp_path: Path) -> None: ) +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) 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_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) + 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_verification_id_is_path_safe() -> None: vid = verification_id_for("run-0001", "sha256:abcdef1234567890") assert vid.startswith("verify-run-0001-") From 82392d0b37f57b8a96148f1ff2351cb2620c3f0e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 12:49:12 +0000 Subject: [PATCH 03/23] fix(kitaru): isolate snapshots by evaluation and per-request latest Key on-disk snapshots by cohort_version_id/evaluation_name so a later --kitaru-evaluation cannot clobber another judge's JSONL. Store latest.json as one pointer per (cohort_name, evaluation_name), and if that pointer is missing pick the newest matching snapshot rather than the lexicographically first directory. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/snapshot.py | 251 +++++++++++++----- src/tracegrad/integrations/kitaru/source.py | 6 +- tests/test_kitaru_snapshot.py | 175 ++++++++++-- 3 files changed, 344 insertions(+), 88 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/snapshot.py b/src/tracegrad/integrations/kitaru/snapshot.py index b5deb62..f86109b 100644 --- a/src/tracegrad/integrations/kitaru/snapshot.py +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -14,9 +14,11 @@ from pydantic.types import StrictInt, StrictStr from tracegrad.state import ( + PathContainmentError, StateLayout, atomic_write, atomic_write_json, + contained_path, initialize, validate_run_id, ) @@ -68,15 +70,51 @@ def kitaru_root(layout: StateLayout) -> Path: return layout.sources / SOURCE_KIND -def snapshot_dir(layout: StateLayout, cohort_version_id: str) -> Path: - return kitaru_root(layout) / cohort_version_id +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_exists(layout: StateLayout, cohort_version_id: str) -> bool: - target = snapshot_dir(layout, cohort_version_id) + +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, *, @@ -87,7 +125,8 @@ def write_snapshot( ) -> Path: """Write JSONL + fingerprint + meta under ``.tracegrad/sources/kitaru/``.""" - target = snapshot_dir(layout, fingerprint.cohort_version_id) + 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=(",", ":")) @@ -110,31 +149,32 @@ def write_snapshot( for drop in dropped ] atomic_write(target / DROPS_FILENAME, "\n".join(drop_lines) + ("\n" if drop_lines else "")) - atomic_write_json( - kitaru_root(layout) / LATEST_POINTER, + _upsert_latest_pointer( + layout, { - "cohort_version_id": fingerprint.cohort_version_id, - "batch": str(target / BATCH_FILENAME), "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, + "batch": str(target / BATCH_FILENAME), }, ) return target -def load_fingerprint(layout: StateLayout, cohort_version_id: str) -> SourceFingerprint: - path = snapshot_dir(layout, cohort_version_id) / FINGERPRINT_FILENAME +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, cohort_version_id: str) -> SourceMeta: - path = snapshot_dir(layout, cohort_version_id) / META_FILENAME +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, cohort_version_id: str) -> tuple[SourceDrop, ...]: - path = snapshot_dir(layout, cohort_version_id) / DROPS_FILENAME +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] = [] @@ -153,8 +193,8 @@ def load_source_drops(layout: StateLayout, cohort_version_id: str) -> tuple[Sour return tuple(drops) -def batch_path(layout: StateLayout, cohort_version_id: str) -> Path: - return snapshot_dir(layout, cohort_version_id) / BATCH_FILENAME +def batch_path(layout: StateLayout, snapshot_id: str) -> Path: + return snapshot_dir(layout, snapshot_id) / BATCH_FILENAME def persist_run_source( @@ -199,9 +239,7 @@ def fingerprints_compatible(stored: SourceFingerprint, requested: Mapping[str, A return True -def load_latest_pointer(layout: StateLayout) -> dict[str, Any] | None: - """Read ``latest.json`` if a previous source write left one.""" - +def _read_pointer_file(layout: StateLayout) -> dict[str, Any] | None: path = kitaru_root(layout) / LATEST_POINTER if not path.is_file(): return None @@ -214,20 +252,71 @@ def load_latest_pointer(layout: StateLayout) -> dict[str, Any] | 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 [] - return [ - child.name - for child in sorted(root.iterdir()) - if child.is_dir() and snapshot_exists(layout, child.name) - ] + 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, - cohort_version_id: str, + snapshot_id: str, *, cohort_name: str, evaluation_name: str, @@ -235,17 +324,45 @@ def snapshot_matches_request( ) -> bool: """Whether this on-disk snapshot is the one the request asked for.""" - if not snapshot_exists(layout, cohort_version_id): + if not snapshot_exists(layout, snapshot_id): return False - fingerprint = load_fingerprint(layout, cohort_version_id) + fingerprint = load_fingerprint(layout, snapshot_id) if not fingerprints_compatible(fingerprint, {"evaluation_name": evaluation_name}): return False - meta = load_meta(layout, cohort_version_id) + meta = load_meta(layout, snapshot_id) if meta.cohort_name != cohort_name: return False if cohort_version is None: return True - return cohort_version_id == cohort_version or meta.display_version == cohort_version + return ( + fingerprint.cohort_version_id == cohort_version + or meta.display_version == cohort_version + or snapshot_id == 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( @@ -255,49 +372,47 @@ def find_local_snapshot( evaluation_name: str, cohort_version: str | None = None, ) -> str | None: - """Return a local ``cohort_version_id`` that satisfies this request. - - Prefers ``latest.json`` when it matches so a re-run keeps the pinned - version even if the server's latest has moved. Does not contact the - server. ``--refresh`` is the caller's decision not to call this. + """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. """ - if cohort_version is not None: - if snapshot_matches_request( + def matches(snapshot_id: str) -> bool: + return snapshot_matches_request( layout, - cohort_version, + snapshot_id, cohort_name=cohort_name, evaluation_name=evaluation_name, cohort_version=cohort_version, - ): - return cohort_version - for snapshot_id in list_snapshot_ids(layout): - if snapshot_matches_request( - layout, - snapshot_id, - cohort_name=cohort_name, - evaluation_name=evaluation_name, - cohort_version=cohort_version, - ): - return snapshot_id - return None + ) - pointer = load_latest_pointer(layout) + 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: - latest_id = str(pointer.get("cohort_version_id") or "") - if latest_id and snapshot_matches_request( - layout, - latest_id, - cohort_name=cohort_name, - evaluation_name=evaluation_name, - ): - return latest_id - for snapshot_id in list_snapshot_ids(layout): - if snapshot_matches_request( - layout, - snapshot_id, - cohort_name=cohort_name, - evaluation_name=evaluation_name, - ): - return snapshot_id - return 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 index 3f03b5b..7638185 100644 --- a/src/tracegrad/integrations/kitaru/source.py +++ b/src/tracegrad/integrations/kitaru/source.py @@ -33,6 +33,7 @@ load_meta, load_source_drops, persist_run_source, + snapshot_identity_key, write_snapshot, ) @@ -150,6 +151,7 @@ async def _fetch_and_map( def _assembled_source( layout: StateLayout, *, + snapshot_id: str, fingerprint: SourceFingerprint, meta: SourceMeta, dropped: Sequence[SourceDrop], @@ -177,7 +179,7 @@ def _assembled_source( "first-root input and last-root output" ) return PreparedSource( - traces_path=batch_path(layout, fingerprint.cohort_version_id), + traces_path=batch_path(layout, snapshot_id), fingerprint=fingerprint, meta=meta, source_table=table, @@ -220,6 +222,7 @@ def prepare_kitaru_source( dropped = load_source_drops(layout, local_id) return _assembled_source( layout, + snapshot_id=local_id, fingerprint=fingerprint, meta=meta, dropped=dropped, @@ -251,6 +254,7 @@ async def _run() -> PreparedSource: ) return _assembled_source( layout, + snapshot_id=snapshot_identity_key(fingerprint), fingerprint=fingerprint, meta=meta, dropped=dropped, diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 5ab21c0..9bd49bf 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os from pathlib import Path from tracegrad.canonical import text_hash @@ -19,6 +20,7 @@ load_latest_pointer, load_source_drops, snapshot_exists, + snapshot_key, write_snapshot, ) from tracegrad.integrations.kitaru.source import prepare_kitaru_source @@ -26,23 +28,25 @@ from tracegrad.state import initialize -def _sample_trace() -> Trace: - return Trace( - trace_id="s1", +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"), ) - - -def _sample_mapped() -> list[MappedTrace]: return [ MappedTrace( - trace=_sample_trace(), + trace=trace, session_number=12, - evaluator_name="quality", - evaluator_version=3, + evaluator_name=evaluator_name, + evaluator_version=evaluator_version, evaluator_version_id="ev", multi_turn=False, system_prompt="prompt", @@ -75,8 +79,8 @@ def _sample_meta(**overrides: object) -> SourceMeta: return SourceMeta.model_validate(payload) -def _manifest() -> Manifest: - return Manifest(template_file=Path("prompt.md"), judge_fingerprint="quality@3") +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: @@ -87,8 +91,8 @@ def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: meta = _sample_meta() write_snapshot(layout, fingerprint=fingerprint, meta=meta, mapped=mapped, dropped=dropped) - assert snapshot_exists(layout, "cv") - stored = load_fingerprint(layout, "cv") + 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( @@ -100,19 +104,24 @@ def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: assert not fingerprints_compatible( stored, {"evaluation_name": "quality", "cohort_version_id": "other"} ) - drops = load_source_drops(layout, "cv") + drops = load_source_drops(layout, snapshot_key("cv", "quality")) assert drops[0].reason == "system-prompt-unavailable" - batch = (layout.sources / "kitaru" / "cv" / "batch.jsonl").read_text(encoding="utf-8") + batch = ( + layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl" + ).read_text(encoding="utf-8") assert "s1" in batch - pointer = load_latest_pointer(layout) + 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" assert find_local_snapshot( layout, cohort_name="support-production", evaluation_name="quality" - ) == "cv" + ) == snapshot_key("cv", "quality") def test_source_and_batch_tables_are_not_merged() -> None: @@ -152,7 +161,7 @@ def test_prepare_reuses_latest_pointer_without_the_server(tmp_path) -> None: ) assert prepared.refreshed is False assert prepared.fingerprint.cohort_version_id == "cv" - assert prepared.traces_path == layout.sources / "kitaru" / "cv" / "batch.jsonl" + assert prepared.traces_path == layout.sources / "kitaru" / "cv" / "quality" / "batch.jsonl" def test_prepare_does_not_follow_moved_remote_latest(tmp_path) -> None: @@ -274,9 +283,12 @@ async def close(self) -> None: assert gateway.resolved == 1 assert prepared.refreshed is True assert prepared.fingerprint.cohort_version_id == "cv-new" - pointer = load_latest_pointer(layout) + 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_old_latest_pointer_without_extra_fields_still_reuses(tmp_path) -> None: @@ -301,3 +313,128 @@ def test_old_latest_pointer_without_extra_fields_still_reuses(tmp_path) -> None: ) 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") From 233ff6bc0e81c4d871c3cdf2b418b292a23be619 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 12:54:08 +0000 Subject: [PATCH 04/23] fix(kitaru): type select/version/score misses as diverged Do not silently drop replay sessions when select_evaluation fails, the evaluator_version mismatches, or classify_scores returns None. Record them as diverged with SELECT_EVALUATION_FAILED, EVALUATOR_VERSION_MISMATCH, and SCORE_UNCLASSIFIED. Leave them out of replay_failures. Headline aggregates still come from Kitaru. Co-authored-by: Dickson Neoh --- CONTEXT.md | 8 +- src/tracegrad/integrations/kitaru/backend.py | 77 ++++++++- src/tracegrad/verify.py | 13 +- tests/test_kitaru_verify.py | 165 +++++++++++++++++++ 4 files changed, 251 insertions(+), 12 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 183913f..7823352 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -102,8 +102,12 @@ 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`. Both divergence kinds -are **incomparable** — not improved, not regressed (ADR 0006). +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 diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index dd4a88a..a6d14dd 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -10,8 +10,11 @@ from typing import Any from tracegrad.verify import ( + DIVERGENCE_EVALUATOR_VERSION, DIVERGENCE_HISTORY, DIVERGENCE_SCOPE, + DIVERGENCE_SCORE, + DIVERGENCE_SELECT, Divergence, ReplayFailure, SubmittedVerification, @@ -130,6 +133,59 @@ def assert_override_scope( 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: @@ -322,18 +378,21 @@ async def _collect( await gateway.evaluations_for(str(result_id)), request.evaluation_name, ) - if isinstance(baseline_eval, str) or isinstance(candidate_eval, str): - continue - if evaluator_version_of(baseline_eval) != request.evaluator_version: - continue - if evaluator_version_of(candidate_eval) != request.evaluator_version: + outcome = classify_replay_session( + session_id=session_id, + number=number, + baseline_eval=baseline_eval, + candidate_eval=candidate_eval, + requested_evaluator_version=int(request.evaluator_version), + ) + if isinstance(outcome, Divergence): + diverged.append(outcome) continue - verdict = classify_scores(baseline_eval, candidate_eval) - if verdict == "improved": + if outcome == "improved": improved.append(session_id) - elif verdict == "regressed": + elif outcome == "regressed": regressed.append(session_id) - elif verdict == "unchanged": + elif outcome == "unchanged": unchanged.append(session_id) # Headline numbers from /api/v1/ui/experiment-runs/{id}/evaluation-aggregates diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 2b12666..4b33ef6 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -30,8 +30,19 @@ 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.""" @@ -41,7 +52,7 @@ class Divergence(BaseModel): model_config = ConfigDict(extra="forbid") session_id: StrictStr - kind: Literal["TOOL_HISTORY_MISS", "OVERRIDE_SCOPE_DIVERGENCE"] + kind: DivergenceKind detail: StrictStr = "" number: StrictInt | None = None diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 80adc19..ec0d3a6 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -19,6 +19,7 @@ from tracegrad.edits import resolve_edits from tracegrad.integrations.kitaru.backend import ( assert_override_scope, + classify_replay_session, classify_scores, is_tool_history_miss, mixed_agent_version_message, @@ -28,15 +29,24 @@ 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_SCORE, + DIVERGENCE_SELECT, + Divergence, SubmittedVerification, VerificationRequest, VerificationResult, VerifyError, + format_verification_report, matching_verification, refuse_ungated_apply, run_verification, @@ -301,6 +311,161 @@ def test_classify_fail_to_pass_is_improved() -> None: 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" + + +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) + 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() From dc99df5a04f0aa366aa9607942e8cae1cb6748ad Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 13:13:43 +0000 Subject: [PATCH 05/23] fix(kitaru): isolate verify event loops, report failures, lookup evaluator Close and drop the owned gateway after each run_async so preflight cannot leave an httpx client bound to a dead loop for submit. Print replay failure count and per-session errors next to Divergence. Look up evaluator_id by mapping.evaluator_name, falling back to the CLI flag. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 33 ++++++-- src/tracegrad/integrations/kitaru/source.py | 4 +- src/tracegrad/verify.py | 8 ++ tests/test_kitaru_mapping.py | 44 +++++++++++ tests/test_kitaru_verify.py | 81 ++++++++++++++++++++ 5 files changed, 162 insertions(+), 8 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index a6d14dd..84ef512 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -224,17 +224,39 @@ class KitaruVerificationBackend: name = "kitaru" def __init__(self, gateway: KitaruGateway | None = None) -> None: - require_kitaru() 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: - run_async(self._preflight(request)) + self._run_async(self._preflight(request)) async def _preflight(self, request: VerificationRequest) -> None: gateway = self._gw() @@ -274,7 +296,7 @@ async def _preflight(self, request: VerificationRequest) -> None: ) def submit(self, request: VerificationRequest) -> SubmittedVerification: - return run_async(self._submit(request)) + return self._run_async(self._submit(request)) async def _submit(self, request: VerificationRequest) -> SubmittedVerification: from kitaru.api_models.v1.experiment import ExperimentCreateRequest @@ -308,7 +330,7 @@ async def _submit(self, request: VerificationRequest) -> SubmittedVerification: def collect( self, request: VerificationRequest, submitted: SubmittedVerification ) -> VerificationResult: - return run_async(self._collect(request, submitted)) + return self._run_async(self._collect(request, submitted)) async def _collect( self, request: VerificationRequest, submitted: SubmittedVerification @@ -424,9 +446,6 @@ async def _collect( verification_fingerprint=verification_fingerprint_for(request), experiment_run_id=submitted.experiment_run_id, ) - if self._owns and self._gateway is not None: - await self._gateway.close() - self._gateway = None return result diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py index 7638185..45a3a91 100644 --- a/src/tracegrad/integrations/kitaru/source.py +++ b/src/tracegrad/integrations/kitaru/source.py @@ -110,7 +110,9 @@ async def _fetch_and_map( f"resolves to more than one evaluator_version across the cohort " f"({breakdown}). Refusing to mix. See ADR 0003." ) - evaluator_id = await gateway.evaluator_id(evaluation_name) + evaluator_id = await gateway.evaluator_id( + mapping.evaluator_name or evaluation_name + ) fingerprint = SourceFingerprint( source="kitaru", cohort_id=resolution.cohort_id, diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 4b33ef6..9afde6c 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -406,6 +406,7 @@ def _num(value: float | None) -> str: 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("") @@ -422,6 +423,13 @@ def _num(value: float | None) -> str: 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: diff --git a/tests/test_kitaru_mapping.py b/tests/test_kitaru_mapping.py index b415559..f3905f0 100644 --- a/tests/test_kitaru_mapping.py +++ b/tests/test_kitaru_mapping.py @@ -6,6 +6,7 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace import pytest @@ -302,3 +303,46 @@ def test_source_and_batch_reasons_stay_kebab_case() -> None: 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_verify.py b/tests/test_kitaru_verify.py index ec0d3a6..a7cac4f 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json from pathlib import Path from types import SimpleNamespace @@ -42,6 +43,7 @@ DIVERGENCE_SCORE, DIVERGENCE_SELECT, Divergence, + ReplayFailure, SubmittedVerification, VerificationRequest, VerificationResult, @@ -584,3 +586,82 @@ 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_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 From 540ffccb9cf3d679eafb694a057fd15b63065df8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 13:25:46 +0000 Subject: [PATCH 06/23] fix(kitaru): match numeric --kitaru-cohort-version on snapshot reuse Persist CohortResolution.version_number on SourceMeta and treat a numeric CLI ref the same as id/display_version, so a re-run with --kitaru-cohort-version 3 can stay offline. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/snapshot.py | 6 ++++ src/tracegrad/integrations/kitaru/source.py | 1 + tests/test_kitaru_snapshot.py | 29 +++++++++++++++++++ 3 files changed, 36 insertions(+) diff --git a/src/tracegrad/integrations/kitaru/snapshot.py b/src/tracegrad/integrations/kitaru/snapshot.py index f86109b..6736625 100644 --- a/src/tracegrad/integrations/kitaru/snapshot.py +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -56,6 +56,7 @@ class SourceMeta(BaseModel): 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) @@ -157,6 +158,7 @@ def write_snapshot( "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), }, ) @@ -338,6 +340,10 @@ def snapshot_matches_request( 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 + ) ) diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py index 45a3a91..315c9c2 100644 --- a/src/tracegrad/integrations/kitaru/source.py +++ b/src/tracegrad/integrations/kitaru/source.py @@ -134,6 +134,7 @@ async def _fetch_and_map( 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={ diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 9bd49bf..3e53ed7 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -18,6 +18,7 @@ kitaru_root, load_fingerprint, load_latest_pointer, + load_meta, load_source_drops, snapshot_exists, snapshot_key, @@ -71,6 +72,7 @@ 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", @@ -119,9 +121,27 @@ def test_snapshot_round_trips_and_is_reusable(tmp_path) -> None: 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: @@ -227,10 +247,19 @@ def test_prepare_reuses_explicit_version_and_display_version(tmp_path) -> None: 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: From 06905d4459ee5c97a262958342cce4e36cd5ac8e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 13:32:38 +0000 Subject: [PATCH 07/23] fix(verify): refuse stale proposals before a Kitaru run Match apply: if the template on disk no longer hashes to proposal.base_prompt_hash, mark stale and refuse with the same message. Do not build a request, submit, or create an experiment. Co-authored-by: Dickson Neoh --- src/tracegrad/cli.py | 8 ++++++ src/tracegrad/verify.py | 7 ++++- tests/test_kitaru_verify.py | 57 +++++++++++++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/src/tracegrad/cli.py b/src/tracegrad/cli.py index f4b6555..9e7612d 100644 --- a/src/tracegrad/cli.py +++ b/src/tracegrad/cli.py @@ -458,6 +458,14 @@ def command_verify(args: argparse.Namespace, out: TextIO) -> int: "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, diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 9afde6c..748e58e 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -15,7 +15,7 @@ from pydantic import BaseModel, ConfigDict, Field from pydantic.types import StrictFloat, StrictInt, StrictStr -from .apply import Proposal, candidate_prompt, load_proposal +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 ( @@ -247,6 +247,11 @@ def build_request( 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) current = template.read_text(encoding="utf-8") candidate = candidate_prompt( diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index a7cac4f..4ec1da8 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import io import json from pathlib import Path from types import SimpleNamespace @@ -12,6 +13,8 @@ from tracegrad.apply import ( Proposal, ProposedEdit, + StaleProposalError, + applied_history, apply_proposal, candidate_prompt, save_proposal, @@ -48,7 +51,9 @@ VerificationRequest, VerificationResult, VerifyError, + build_request, format_verification_report, + load_run_source_payload, matching_verification, refuse_ungated_apply, run_verification, @@ -588,6 +593,58 @@ def test_verification_id_is_path_safe() -> None: 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_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)) + + def test_report_lists_replay_failures_next_to_divergence() -> None: request = _request(session_numbers={"crash-session": 7}) result = VerificationResult( From 243f8f1b4fe68d74d5b8c7ddefbc3dd3696d6351 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 13:37:32 +0000 Subject: [PATCH 08/23] fix(verify): wrap unreadable template as VerifyError build_request's template.read_text can raise OSError after the stale check; catch it as VerifyError like apply_proposal, so it hits main's handler instead of a traceback. Co-authored-by: Dickson Neoh --- src/tracegrad/verify.py | 5 ++++- tests/test_kitaru_verify.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 748e58e..232414a 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -253,7 +253,10 @@ def build_request( "the proposal is stale — re-run tracegrad" ) template = contained_path(base_directory, proposal.template_file) - current = template.read_text(encoding="utf-8") + 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)) ) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 4ec1da8..8b8ce42 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -614,6 +614,25 @@ def test_build_request_refuses_a_stale_proposal_without_submitting(tmp_path: Pat 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_verify_cli_refuses_a_stale_proposal_before_kitaru(tmp_path: Path) -> None: from tracegrad import cli From f4629a39520b4330072af8f100079f77bb88baa5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 13:53:21 +0000 Subject: [PATCH 09/23] fix(kitaru): parallelize collect fetches; apply via candidate_prompt Collect was four sequential HTTP calls per replay after wait. Fetch per-replay payloads with the same jobs=8 semaphore as source mapping. apply_proposal now writes candidate_prompt's text so ADR 0009's hash gate hashes the same bytes that land on disk. Co-authored-by: Dickson Neoh --- src/tracegrad/apply.py | 29 ++--- src/tracegrad/integrations/kitaru/backend.py | 60 ++++++++-- src/tracegrad/integrations/kitaru/client.py | 4 +- tests/test_apply.py | 16 +++ tests/test_kitaru_verify.py | 120 +++++++++++++++++++ 5 files changed, 193 insertions(+), 36 deletions(-) diff --git a/src/tracegrad/apply.py b/src/tracegrad/apply.py index d7dd3cd..61840b1 100644 --- a/src/tracegrad/apply.py +++ b/src/tracegrad/apply.py @@ -284,8 +284,8 @@ def candidate_prompt( ) -> str: """The text that would be written if these indices were accepted. - Used by verify (the full proposal) and the apply gate (the selection - about to be written) so both hash the same way. + Used by verify, the apply gate, and apply_proposal so the ADR 0009 hash + gate hashes the same bytes that are written. """ selected = sorted({index for index in accepted_indices}) @@ -329,28 +329,15 @@ 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_prompt is that shared path. + updated = candidate_prompt(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, @@ -359,7 +346,6 @@ def apply_proposal( rejected=tuple(rejected), snapshot=None, unchanged=True, - resolution_rejections=tuple(item.reason for item in resolution.rejected), ) snapshot = snapshot_template(project_root, proposal.run_id, template) @@ -387,7 +373,6 @@ def apply_proposal( accepted=tuple(accepted), rejected=tuple(rejected), snapshot=snapshot, - resolution_rejections=tuple(item.reason for item in resolution.rejected), ) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index 84ef512..af89052 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -6,7 +6,9 @@ from __future__ import annotations +import asyncio import uuid +from collections.abc import Sequence from typing import Any from tracegrad.verify import ( @@ -23,7 +25,7 @@ verification_fingerprint_for, ) -from .client import KitaruGateway, run_async, worker_covers_agent_version +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 from .mapping import extract_system_prompt @@ -344,6 +346,7 @@ async def _collect( unchanged: list[str] = [] diverged: list[Divergence] = [] failures: list[ReplayFailure] = [] + fetch_targets: list[Any] = [] for replay in replays: session_id = str(replay.baseline_session_id) @@ -377,8 +380,13 @@ async def _collect( ReplayFailure(session_id=session_id, error="missing result session", number=number) ) continue - baseline_nodes = await gateway.session_nodes(str(replay.baseline_session_id)) - result_nodes = await gateway.session_nodes(str(result_id)) + 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 ) @@ -392,19 +400,11 @@ async def _collect( ) ) continue - baseline_eval = select_evaluation( - await gateway.evaluations_for(str(replay.baseline_session_id)), - request.evaluation_name, - ) - candidate_eval = select_evaluation( - await gateway.evaluations_for(str(result_id)), - request.evaluation_name, - ) outcome = classify_replay_session( session_id=session_id, number=number, - baseline_eval=baseline_eval, - candidate_eval=candidate_eval, + 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): @@ -449,6 +449,40 @@ async def _collect( 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. + """ + + 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 + + return list(await asyncio.gather(*(one(replay) for replay in replays))) + + def _maybe_float(value: Any) -> float | None: if isinstance(value, (int, float)) and not isinstance(value, bool): return float(value) diff --git a/src/tracegrad/integrations/kitaru/client.py b/src/tracegrad/integrations/kitaru/client.py index 321e1d1..2992790 100644 --- a/src/tracegrad/integrations/kitaru/client.py +++ b/src/tracegrad/integrations/kitaru/client.py @@ -15,6 +15,8 @@ from .errors import KitaruSourceError, KitaruVerifyError from .require import require_kitaru +FETCH_JOBS = 8 + @dataclass(frozen=True) class CohortResolution: @@ -188,7 +190,7 @@ async def _evaluations_for(self, session_id: uuid.UUID) -> list[Any]: return [item async for item in self._client.evaluations.iter(params)] async def fetch_records( - self, sessions: Sequence[Any], *, jobs: int = 8 + self, sessions: Sequence[Any], *, jobs: int = FETCH_JOBS ) -> list[tuple[Any, tuple[Any, ...], tuple[Any, ...]]]: semaphore = asyncio.Semaphore(max(1, jobs)) diff --git a/tests/test_apply.py b/tests/test_apply.py index 2bf088b..52d05c4 100644 --- a/tests/test_apply.py +++ b/tests/test_apply.py @@ -10,6 +10,7 @@ applied_history, apply_proposal, build_proposal, + candidate_prompt, current_baseline, is_stale, latest_run_id, @@ -151,6 +152,21 @@ 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_writes_the_template_and_records_the_baseline(tmp_path: Path) -> None: template = _template(tmp_path) outcome, _ = _outcome() diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 8b8ce42..7ef69ae 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -22,12 +22,14 @@ 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 from tracegrad.integrations.kitaru.policy import ( RECORDED_HISTORY_POLICY, asserts_no_passthrough, @@ -43,6 +45,7 @@ from tracegrad.state import atomic_write_json, initialize from tracegrad.verify import ( DIVERGENCE_EVALUATOR_VERSION, + DIVERGENCE_HISTORY, DIVERGENCE_SCORE, DIVERGENCE_SELECT, Divergence, @@ -741,3 +744,120 @@ async def probe() -> str: 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) + return () + + 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 From 6dde94e09c9f17cd6a06c4a7f9e297ef0be8f752 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 13:57:15 +0000 Subject: [PATCH 10/23] fix(verify): partial apply needs --force or --all; empty roots diverge refuse_ungated_apply no longer tells the operator to re-run verify. Verify hashes the full proposal, so a subset cannot be ungated that way. assert_override_scope treats a result graph with no root llm_call as OVERRIDE_SCOPE_DIVERGENCE instead of a vacuous pass (ADR 0006). Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 12 +- src/tracegrad/verify.py | 15 ++- tests/test_kitaru_verify.py | 132 ++++++++++++++++++- 3 files changed, 151 insertions(+), 8 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index af89052..2061d7f 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -27,7 +27,7 @@ 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 +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 @@ -114,8 +114,16 @@ def assert_override_scope( 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.""" + """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): diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 232414a..1025221 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -225,17 +225,22 @@ def refuse_ungated_apply( candidate_prompt_hash: str, force: bool, ) -> None: - """Refuse apply when a backend is configured and no matching verify exists.""" + """Refuse apply when a backend is configured and no matching verify exists. + + Verify hashes the full proposal. A selected subset therefore 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 if matching_verification(project_root, candidate_prompt_hash) is None: raise VerifyError( "apply is gated on a hash-matching verification for this candidate. " - "Run `tracegrad verify --backend kitaru` first, or pass --force " - "to override. Matching the hash (not the run id) is what makes the " - "gate real: verify, hand-edit, and apply notices the text was never " - "verified. See ADR 0009." + "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. " + "See ADR 0009." ) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 7ef69ae..10feca1 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -99,6 +99,46 @@ def _proposal(project: Path, prompt: str = PROMPT) -> 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") -> None: layout = initialize(project) atomic_write_json( @@ -307,6 +347,59 @@ def test_override_scope_divergence_on_non_root() -> 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 @@ -504,6 +597,33 @@ def test_apply_is_gated_on_a_matching_hash(tmp_path: Path) -> None: ) +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]) @@ -817,7 +937,17 @@ async def list_replays(self, experiment_run_id: str) -> list[object]: async def session_nodes(self, session_id: str) -> tuple[object, ...]: await self._track(session_id) - return () + 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) From 73dab55bdcaafa57c55d9c93d2ec1322134c5eaa Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 14:08:26 +0000 Subject: [PATCH 11/23] fix(kitaru): check judge fingerprint before snapshot write A --refresh whose mapped evaluator conflicts with the manifest used to write_snapshot (and latest.json) first, then error. Check the fetched mapping before writing so a failed refresh leaves the last-good snapshot. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/source.py | 36 +++++++-- tests/test_kitaru_snapshot.py | 89 +++++++++++++++++++++ 2 files changed, 119 insertions(+), 6 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py index 315c9c2..3aeef57 100644 --- a/src/tracegrad/integrations/kitaru/source.py +++ b/src/tracegrad/integrations/kitaru/source.py @@ -68,6 +68,22 @@ def check_judge_fingerprint(manifest: Manifest, derived: str) -> None: ) +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: @@ -163,12 +179,12 @@ def _assembled_source( evaluation_name: str, run_id: str | None, ) -> PreparedSource: - if meta.traces_mapped: - derived = judge_fingerprint_for( - meta.evaluator_name or evaluation_name, - fingerprint.evaluator_version, - ) - check_judge_fingerprint(manifest, derived) + _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( @@ -252,6 +268,14 @@ async def _run() -> PreparedSource: 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 ) diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 3e53ed7..6980155 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -4,6 +4,9 @@ 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 @@ -320,6 +323,92 @@ async def close(self) -> None: 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( From 20c64b42fd80f824fee4180e359ef6f638171754 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 14:16:46 +0000 Subject: [PATCH 12/23] fix(kitaru): lock fail-closed collect on fetch 404/timeout One session_nodes/evaluations error aborts the whole collect gather. Do not isolate as ReplayFailure, persist a result, or ungate apply. Resume keeps the experiment run id and redoes collect. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 4 + tests/test_kitaru_verify.py | 101 +++++++++++++++++++ 2 files changed, 105 insertions(+) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index 2061d7f..6225334 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -470,6 +470,9 @@ async def _fetch_replay_payloads( 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: @@ -488,6 +491,7 @@ async def one(replay: Any) -> _ReplayPayload: ) return baseline_nodes, result_nodes, baseline_evals, candidate_evals + # Fail-closed: gather raises on the first fetch error (no return_exceptions). return list(await asyncio.gather(*(one(replay) for replay in replays))) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 10feca1..9ff4d52 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -57,6 +57,7 @@ build_request, format_verification_report, load_run_source_payload, + load_verification_state, matching_verification, refuse_ungated_apply, run_verification, @@ -684,6 +685,106 @@ def collect( ) +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(TimeoutError, match="404/timeout"): + run_verification(tmp_path, request, FetchErrorBackend()) + + assert matching_verification(tmp_path, digest) 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 test_apply_gate_accepts_a_finished_failed_report(tmp_path: Path) -> None: """A finished REVIEW/FAILED report still allows apply; do not require completed.""" From 86f8a38283705fa83ed7575b88a41769b5fdd7f9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 14:19:19 +0000 Subject: [PATCH 13/23] fix(kitaru): wrap collect fetch errors as KitaruVerifyError A 404/timeout from the payload gather still aborts the whole collect fail-closed. Re-raise as KitaruVerifyError so the CLI prints an actionable message instead of a traceback: apply stays gated, resume redoes collect, --force if they must write anyway. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 11 ++++++++++- tests/test_kitaru_verify.py | 9 ++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index 6225334..fa911b3 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -492,7 +492,16 @@ async def one(replay: Any) -> _ReplayPayload: return baseline_nodes, result_nodes, baseline_evals, candidate_evals # Fail-closed: gather raises on the first fetch error (no return_exceptions). - return list(await asyncio.gather(*(one(replay) for replay in replays))) + try: + return list(await asyncio.gather(*(one(replay) for replay in replays))) + except KitaruVerifyError: + raise + except Exception as exc: + raise KitaruVerifyError( + "collect aborted: fetching a replay payload failed " + f"({type(exc).__name__}: {exc}). Apply stays gated. Resume redoes " + "collect. Pass --force on apply if you must write anyway." + ) from exc def _maybe_float(value: Any) -> float | None: diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 9ff4d52..4eeee12 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -30,6 +30,7 @@ mixed_agent_version_message, ) from tracegrad.integrations.kitaru.client import FETCH_JOBS +from tracegrad.integrations.kitaru.errors import KitaruVerifyError from tracegrad.integrations.kitaru.policy import ( RECORDED_HISTORY_POLICY, asserts_no_passthrough, @@ -769,8 +770,14 @@ def collect( request, submitted ) - with pytest.raises(TimeoutError, match="404/timeout"): + 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) is None state = load_verification_state( From c83838f1ce9831b8be1e18159761a8ebddd43cbb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 14:30:20 +0000 Subject: [PATCH 14/23] fix(kitaru): drop include_stale from WorkerListParams kitaru 0.22 WorkerListParams is FilterableListParams with extra='forbid', so include_stale=False raised ValidationError before the live-worker check. List with the 0.22 constructor; worker_covers_agent_version already skips live=False. Wrap list_live_workers like the other probes. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 9 +++- src/tracegrad/integrations/kitaru/client.py | 5 +- tests/test_kitaru_verify.py | 50 +++++++++++++++++++- 3 files changed, 60 insertions(+), 4 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index fa911b3..198284f 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -297,7 +297,14 @@ async def _preflight(self, request: VerificationRequest) -> None: only = next(iter(distinct)) if only != request.agent_version_id: raise KitaruVerifyError(mixed_agent_version_message(counts)) - workers = await gateway.list_live_workers() + 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 " diff --git a/src/tracegrad/integrations/kitaru/client.py b/src/tracegrad/integrations/kitaru/client.py index 2992790..c780adc 100644 --- a/src/tracegrad/integrations/kitaru/client.py +++ b/src/tracegrad/integrations/kitaru/client.py @@ -222,8 +222,9 @@ async def get_cohort_version(self, cohort_version_id: str) -> Any: async def list_live_workers(self) -> list[Any]: from kitaru.api_models.v1.worker import WorkerListParams - params = WorkerListParams(include_stale=False) - return [item async for item in self._client.workers.iter(params)] + # 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) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 4eeee12..04dc868 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -29,7 +29,7 @@ is_tool_history_miss, mixed_agent_version_message, ) -from tracegrad.integrations.kitaru.client import FETCH_JOBS +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, @@ -408,6 +408,54 @@ def test_mixed_agent_version_message_includes_a_breakdown() -> None: 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) From dde80bc63e59feba268a6165a964e48dc21ecfc0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 14:35:04 +0000 Subject: [PATCH 15/23] fix(kitaru): wrap sibling SDK calls; drop dead _pick_version UUID wait_for_experiment_run, list_replays, evaluation_aggregates, and create_experiment/start_run now raise KitaruVerifyError like payload fetch. fetch_records gather failures become KitaruSourceError. Both stays fail-closed. Delete _pick_version's trailing UUID try/except that always returned None. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 84 ++++++++----- src/tracegrad/integrations/kitaru/client.py | 18 ++- tests/test_kitaru_snapshot.py | 41 ++++++- tests/test_kitaru_verify.py | 121 +++++++++++++++++++ 4 files changed, 230 insertions(+), 34 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index 198284f..6e37f6e 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -228,6 +228,25 @@ def classify_scores( 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.""" @@ -321,23 +340,29 @@ async def _submit(self, request: VerificationRequest) -> SubmittedVerification: gateway = self._gw() digest = request.candidate_prompt_hash.removeprefix("sha256:")[:8] - experiment = await 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)], - ) + 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 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, + 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), @@ -353,9 +378,15 @@ async def _collect( self, request: VerificationRequest, submitted: SubmittedVerification ) -> VerificationResult: gateway = self._gw() - run = await gateway.wait_for_experiment_run(submitted.experiment_run_id) + 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 gateway.list_replays(submitted.experiment_run_id) + 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] = [] @@ -435,7 +466,10 @@ async def _collect( # 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 gateway.evaluation_aggregates(submitted.experiment_run_id) + 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" @@ -499,16 +533,10 @@ async def one(replay: Any) -> _ReplayPayload: return baseline_nodes, result_nodes, baseline_evals, candidate_evals # Fail-closed: gather raises on the first fetch error (no return_exceptions). - try: - return list(await asyncio.gather(*(one(replay) for replay in replays))) - except KitaruVerifyError: - raise - except Exception as exc: - raise KitaruVerifyError( - "collect aborted: fetching a replay payload failed " - f"({type(exc).__name__}: {exc}). Apply stays gated. Resume redoes " - "collect. Pass --force on apply if you must write anyway." - ) from exc + 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: diff --git a/src/tracegrad/integrations/kitaru/client.py b/src/tracegrad/integrations/kitaru/client.py index c780adc..05e568e 100644 --- a/src/tracegrad/integrations/kitaru/client.py +++ b/src/tracegrad/integrations/kitaru/client.py @@ -136,10 +136,6 @@ def _pick_version( return item if str(item.version) == version_ref: return item - try: - _uuid(version_ref) - except ValueError: - return None return None async def _list_versions(self, cohort_id: uuid.UUID) -> list[Any]: @@ -198,7 +194,19 @@ async def one(session: Any) -> tuple[Any, tuple[Any, ...], tuple[Any, ...]]: async with semaphore: return await self.session_bundle(str(session.id)) - return list(await asyncio.gather(*(one(session) for session in sessions))) + 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 diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 6980155..896133d 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import os from pathlib import Path from types import SimpleNamespace @@ -10,7 +11,8 @@ from tracegrad.canonical import text_hash from tracegrad.integrations.kitaru.accounting import format_source_table -from tracegrad.integrations.kitaru.client import CohortResolution +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 ( LATEST_POINTER, @@ -556,3 +558,40 @@ def test_missing_pointer_picks_newest_matching_snapshot(tmp_path) -> None: 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 diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 04dc868..12591f8 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -840,6 +840,127 @@ def collect( ) +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) 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.""" From d75d229826618911ec7c90802ac9c9accd8a9b76 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 14:49:24 +0000 Subject: [PATCH 16/23] fix(kitaru): skip corrupt snapshots and wrap remaining source IO find_local_snapshot no longer aborts reuse when a snapshot dir is unreadable. list_sessions, resolve_cohort, and evaluator_id raise KitaruSourceError like fetch_records. A corrupt kitaru-source.json does not traceback or ungate apply. Delete unused load_run_source. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/snapshot.py | 18 ++-- src/tracegrad/integrations/kitaru/source.py | 32 ++++++- src/tracegrad/verify.py | 15 +++- tests/test_kitaru_snapshot.py | 87 +++++++++++++++++++ tests/test_kitaru_verify.py | 18 ++++ 5 files changed, 151 insertions(+), 19 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/snapshot.py b/src/tracegrad/integrations/kitaru/snapshot.py index 6736625..0e4b629 100644 --- a/src/tracegrad/integrations/kitaru/snapshot.py +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -19,7 +19,6 @@ atomic_write, atomic_write_json, contained_path, - initialize, validate_run_id, ) @@ -218,14 +217,6 @@ def persist_run_source( return target -def load_run_source(project_root: str | Path | StateLayout, run_id: str) -> dict[str, Any] | None: - layout = initialize(project_root) - target = layout.runs / validate_run_id(run_id) / RUN_SOURCE_FILENAME - if not target.exists(): - return None - return json.loads(target.read_text(encoding="utf-8")) - - def fingerprints_compatible(stored: SourceFingerprint, requested: Mapping[str, Any]) -> bool: """Whether a snapshot can be reused for this request without refetching.""" @@ -328,10 +319,13 @@ def snapshot_matches_request( if not snapshot_exists(layout, snapshot_id): return False - fingerprint = load_fingerprint(layout, snapshot_id) - if not fingerprints_compatible(fingerprint, {"evaluation_name": evaluation_name}): + try: + fingerprint = load_fingerprint(layout, snapshot_id) + if not fingerprints_compatible(fingerprint, {"evaluation_name": evaluation_name}): + return False + meta = load_meta(layout, snapshot_id) + except (OSError, ValueError): return False - meta = load_meta(layout, snapshot_id) if meta.cohort_name != cohort_name: return False if cohort_version is None: diff --git a/src/tracegrad/integrations/kitaru/source.py b/src/tracegrad/integrations/kitaru/source.py index 3aeef57..3e7a46c 100644 --- a/src/tracegrad/integrations/kitaru/source.py +++ b/src/tracegrad/integrations/kitaru/source.py @@ -100,13 +100,33 @@ def _single_agent_version(counts: dict[str, int]) -> str | None: 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 gateway.list_sessions(resolution.cohort_version_id) + 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: @@ -126,8 +146,9 @@ async def _fetch_and_map( f"resolves to more than one evaluator_version across the cohort " f"({breakdown}). Refusing to mix. See ADR 0003." ) - evaluator_id = await gateway.evaluator_id( - mapping.evaluator_name or evaluation_name + evaluator_id = await _await_source( + gateway.evaluator_id(mapping.evaluator_name or evaluation_name), + what="looking up the evaluator", ) fingerprint = SourceFingerprint( source="kitaru", @@ -262,7 +283,10 @@ def prepare_kitaru_source( async def _run() -> PreparedSource: try: - resolution = await gateway.resolve_cohort(cohort_name, cohort_version) + 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, diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 1025221..3adcdf3 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -209,13 +209,22 @@ def load_run_source_payload(project_root: str | Path, run_id: str) -> dict[str, return None import json - return json.loads(target.read_text(encoding="utf-8")) + try: + payload = json.loads(target.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return payload if isinstance(payload, dict) else None def backend_is_configured(project_root: str | Path, run_id: str) -> bool: - """Whether this run originated from a Kitaru source (ADR 0009).""" + """Whether this run originated from a Kitaru source (ADR 0009). + + Presence of the sidecar is enough. A corrupt file must not ungate apply. + """ - return load_run_source_payload(project_root, run_id) is not None + layout = initialize(project_root) + target = layout.runs / validate_run_id(run_id) / RUN_SOURCE_FILENAME + return target.is_file() def refuse_ungated_apply( diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 896133d..2a6451d 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -595,3 +595,90 @@ async def session_bundle(self, session_id: str) -> tuple[object, ...]: 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_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()) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 12591f8..276a61f 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -55,6 +55,7 @@ VerificationRequest, VerificationResult, VerifyError, + backend_is_configured, build_request, format_verification_report, load_run_source_payload, @@ -1033,6 +1034,23 @@ def test_build_request_wraps_unreadable_template_as_verify_error( ) +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_verify_cli_refuses_a_stale_proposal_before_kitaru(tmp_path: Path) -> None: from tracegrad import cli From c66c930b98400d0958e7655b9f71c8576613214a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 15:03:00 +0000 Subject: [PATCH 17/23] fix(kitaru): restore apply rejections; reject incomplete sidecars candidate_prompt discarded resolve_edits rejections after apply started writing through it, so ApplyResult.resolution_rejections stayed empty. Share one helper that returns the candidate bytes plus rejection reasons so the ADR 0009 hash still matches the write, and thread those reasons into both ApplyResult paths as before. load_run_source_payload treated any JSON object as valid, so empty or partial sidecars reached build_request and KeyError/TypeError escaped the CLI handler. Reject payloads whose fingerprint/meta are not dicts with the required keys (load returns None; apply stays gated; verify raises VerifyError). Co-authored-by: Dickson Neoh --- src/tracegrad/apply.py | 40 ++++++++++++---- src/tracegrad/verify.py | 92 +++++++++++++++++++++++++------------ tests/test_apply.py | 30 +++++++++++- tests/test_kitaru_verify.py | 27 +++++++++++ 4 files changed, 150 insertions(+), 39 deletions(-) diff --git a/src/tracegrad/apply.py b/src/tracegrad/apply.py index 61840b1..8716589 100644 --- a/src/tracegrad/apply.py +++ b/src/tracegrad/apply.py @@ -277,15 +277,15 @@ def snapshot_template(project_root: str | Path, run_id: str, template: Path) -> return target -def candidate_prompt( +def _candidate_text_and_rejections( prompt: str, proposal: Proposal, accepted_indices: Iterable[int], -) -> str: - """The text that would be written if these indices were accepted. +) -> tuple[str, tuple[str, ...]]: + """The text that would be written, plus resolve_edits rejection reasons. - Used by verify, the apply gate, and apply_proposal so the ADR 0009 hash - gate hashes the same bytes that are written. + 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}) @@ -294,10 +294,28 @@ def candidate_prompt( raise ApplyError(f"no such edit index: {index}") accepted = [proposal.edits[index].edit for index in selected] if not accepted: - return prompt + return prompt, () inventory = build_inventory(prompt) resolution = resolve_edits(inventory, accepted) - return apply_resolved(prompt, resolution.resolved) + 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( @@ -330,8 +348,10 @@ def apply_proposal( selected = sorted({index for index in accepted_indices}) # ADR 0009's hash gate is only sound if apply writes the same bytes verify - # hashed. candidate_prompt is that shared path. - updated = candidate_prompt(current, proposal, selected) + # 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 = [ @@ -346,6 +366,7 @@ def apply_proposal( rejected=tuple(rejected), snapshot=None, unchanged=True, + resolution_rejections=resolution_rejections, ) snapshot = snapshot_template(project_root, proposal.run_id, template) @@ -373,6 +394,7 @@ def apply_proposal( accepted=tuple(accepted), rejected=tuple(rejected), snapshot=snapshot, + resolution_rejections=resolution_rejections, ) diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 3adcdf3..3ca3267 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -27,6 +27,14 @@ ) 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" @@ -213,7 +221,22 @@ def load_run_source_payload(project_root: str | Path, run_id: str) -> dict[str, payload = json.loads(target.read_text(encoding="utf-8")) except (OSError, ValueError): return None - return payload if isinstance(payload, dict) else 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: @@ -274,36 +297,47 @@ def build_request( candidate = candidate_prompt( current, proposal, range(len(proposal.edits)) ) - fingerprint = source["fingerprint"] - meta = source["meta"] - agent_version_id = meta.get("agent_version_id") - if not agent_version_id: + 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( - "verification requires a single agent_version_id on the originating " - "cohort; this run's source metadata does not have one. See ADR 0007." + f"run {run_id} source metadata is not a usable Kitaru sidecar" ) - 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() - }, - ) + 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: diff --git a/tests/test_apply.py b/tests/test_apply.py index 52d05c4..e47a94c 100644 --- a/tests/test_apply.py +++ b/tests/test_apply.py @@ -6,6 +6,8 @@ from tracegrad.apply import ( ApplyError, + Proposal, + ProposedEdit, StaleProposalError, applied_history, apply_proposal, @@ -21,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 @@ -167,6 +169,32 @@ def test_apply_proposal_writes_the_same_text_as_candidate_prompt(tmp_path: Path) 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_kitaru_verify.py b/tests/test_kitaru_verify.py index 276a61f..76577aa 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -1051,6 +1051,33 @@ def test_corrupt_run_source_sidecar_does_not_traceback_or_ungate(tmp_path: Path) ) +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 From fd863531c1eb917c94a9ee4aee6498c1487c6815 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 15:17:37 +0000 Subject: [PATCH 18/23] fix(kitaru): poll experiment_runs.get instead of wrapping KitaruClient KitaruGateway.wait_for_experiment_run constructed KitaruClient and called wait_for_experiment_run, which is not on the 0.22 API client (close / context-manager only). Wrapping KitaruClient is also unsafe: close() shuts the shared API client. Poll self._client.experiment_runs.get until completed/failed/canceled, matching kitaru CLI poll_run. Timeouts still raise so collect stays fail-closed. Confirmed against pinned kitaru 0.22. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/client.py | 50 +++++++++++- tests/test_kitaru_snapshot.py | 90 +++++++++++++++++++++ 2 files changed, 137 insertions(+), 3 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/client.py b/src/tracegrad/integrations/kitaru/client.py index 05e568e..d7efe8b 100644 --- a/src/tracegrad/integrations/kitaru/client.py +++ b/src/tracegrad/integrations/kitaru/client.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import time import uuid from collections.abc import Sequence from dataclasses import dataclass @@ -16,6 +17,8 @@ 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) @@ -244,10 +247,51 @@ 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: - from kitaru.client.client import KitaruClient + """Poll ``experiment_runs.get`` until completed, failed, or canceled. - wrapper = KitaruClient(api_client=self._client) - return await wrapper.wait_for_experiment_run(_uuid(run_id), timeout=timeout) + 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 diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index 2a6451d..da3b163 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import inspect import os from pathlib import Path from types import SimpleNamespace @@ -682,3 +683,92 @@ async def close(self) -> None: 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, + ) + ) From 91583a1d4f18ba6ec81f552a1c820c1332a5dd73 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 15:29:39 +0000 Subject: [PATCH 19/23] fix(kitaru): skip snapshots with corrupt source-drops find_local_snapshot already skipped unreadable fingerprint/meta, but prepare_kitaru_source then loaded source-drops.jsonl with no guard. Malformed JSON or a missing session_id/reason raised JSONDecodeError or KeyError instead of a named source error. Treat a bad drops file like a corrupt snapshot in snapshot_matches_request so reuse picks the next good snapshot, or falls through to a named fetch error when none remain. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/snapshot.py | 5 +- tests/test_kitaru_snapshot.py | 60 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/src/tracegrad/integrations/kitaru/snapshot.py b/src/tracegrad/integrations/kitaru/snapshot.py index 0e4b629..4a615b3 100644 --- a/src/tracegrad/integrations/kitaru/snapshot.py +++ b/src/tracegrad/integrations/kitaru/snapshot.py @@ -324,7 +324,10 @@ def snapshot_matches_request( if not fingerprints_compatible(fingerprint, {"evaluation_name": evaluation_name}): return False meta = load_meta(layout, snapshot_id) - except (OSError, ValueError): + # 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 diff --git a/tests/test_kitaru_snapshot.py b/tests/test_kitaru_snapshot.py index da3b163..e0c653b 100644 --- a/tests/test_kitaru_snapshot.py +++ b/tests/test_kitaru_snapshot.py @@ -16,6 +16,7 @@ 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, @@ -626,6 +627,65 @@ def test_corrupt_snapshot_dir_is_skipped_during_reuse(tmp_path) -> 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", From 64986733e8062bd6034293399ec3e16284815929 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 16:25:04 +0000 Subject: [PATCH 20/23] fix(kitaru): gate apply on hash and cohort, not run_id Captain option (c): matching_verification requires candidate_prompt_hash and the sidecar cohort_version_id. Hash-only would ungate a new cohort without re-verify; matching apply's run_id would miss a valid verify of the same candidate on the same cohort. refuse_ungated_apply still receives run_id, but only to load the source sidecar for that cohort check. Same hash + same cohort ungates; same hash + a new cohort stays gated; a hand-edit (hash change) stays gated. Co-authored-by: Dickson Neoh --- src/tracegrad/verify.py | 61 +++++++++++++++++++++++++-------- tests/test_kitaru_verify.py | 67 ++++++++++++++++++++++++++++++++----- 2 files changed, 106 insertions(+), 22 deletions(-) diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index 3ca3267..ddb2fc6 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -3,8 +3,8 @@ 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`` lives here so the core of ``apply`` stays a writer, -not a Kitaru client. +``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 @@ -189,18 +189,24 @@ def list_verification_states(layout: StateLayout) -> list[VerificationState]: def matching_verification( project_root: str | Path, candidate_prompt_hash: str, + *, + cohort_version_id: str, ) -> VerificationState | None: - """A persisted verification of exactly this candidate text, if any. + """A persisted verification of this candidate on this cohort, if any. - Requires a stored ``result``. A submit that never collected (``result`` - is ``None``) must not ungate apply. A finished REVIEW/FAILED report - still matches; status is not required to be ``completed``. + 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 ): @@ -208,6 +214,21 @@ def matching_verification( 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.""" @@ -259,20 +280,32 @@ def refuse_ungated_apply( ) -> None: """Refuse apply when a backend is configured and no matching verify exists. - Verify hashes the full proposal. A selected subset therefore cannot be - ungated by re-running verify; the operator must ``apply --all`` or - ``--force`` (ADR 0009). + 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 - if matching_verification(project_root, candidate_prompt_hash) is None: + 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. " - "Verify always hashes the full proposal. A subset of a verified " - "proposal (interactive or --accept) needs --force or " + "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. " - "See ADR 0009." + "A new cohort needs `tracegrad verify` even when the candidate " + "text is unchanged. See ADR 0009." ) diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 76577aa..7b6a519 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -142,7 +142,9 @@ def _multi_edit_proposal(project: Path, prompt: str = PROMPT) -> Proposal: return proposal -def _source_sidecar(project: Path, run_id: str = "run-0001") -> None: +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, @@ -150,7 +152,7 @@ def _source_sidecar(project: Path, run_id: str = "run-0001") -> None: "fingerprint": { "source": "kitaru", "cohort_id": "c1", - "cohort_version_id": "cv1", + "cohort_version_id": cohort_version_id, "evaluation_name": "quality", "evaluator_id": "ev", "evaluator_version": 3, @@ -604,7 +606,11 @@ def collect( "EVALUATOR_VERSION_MISMATCH", "SCORE_UNCLASSIFIED", } - matched = matching_verification(tmp_path, request.candidate_prompt_hash) + 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" @@ -629,7 +635,14 @@ def test_interrupted_verify_does_not_duplicate_the_experiment(tmp_path: Path) -> 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) is not None + assert ( + matching_verification( + tmp_path, + request.candidate_prompt_hash, + cohort_version_id=request.cohort_version_id, + ) + is not None + ) def test_apply_is_gated_on_a_matching_hash(tmp_path: Path) -> None: @@ -728,7 +741,7 @@ def collect( 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) 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 @@ -828,7 +841,7 @@ def collect( assert "TimeoutError" in message assert caught.value.__cause__ is not None - assert matching_verification(tmp_path, digest) is 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) ) @@ -931,7 +944,7 @@ def collect( 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) is 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) ) @@ -979,7 +992,7 @@ def collect( request = _request(candidate_prompt=written, candidate_prompt_hash=digest) run_verification(tmp_path, request, FailedBackend()) - matched = matching_verification(tmp_path, digest) + 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" @@ -988,6 +1001,44 @@ def collect( ) +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-") From 5047472b244282cccb9e0937ae6578ae22b810bb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 17:35:11 +0000 Subject: [PATCH 21/23] fix(kitaru): treat corrupt verification state as missing load_verification_state raised on unreadable state.json, so run_verification could not resubmit. Swallow OSError/ValueError (including pydantic ValidationError) and return None, matching list_verification_states and snapshot reuse. Apply stays gated until a usable result is stored; hash+cohort match is unchanged. Co-authored-by: Dickson Neoh --- src/tracegrad/verify.py | 12 +++++++++++- tests/test_kitaru_verify.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/src/tracegrad/verify.py b/src/tracegrad/verify.py index ddb2fc6..b10eb7b 100644 --- a/src/tracegrad/verify.py +++ b/src/tracegrad/verify.py @@ -170,10 +170,20 @@ def save_verification_state(layout: StateLayout, state: VerificationState) -> Pa 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 - return VerificationState.model_validate_json(target.read_text(encoding="utf-8")) + 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]: diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 7b6a519..57b2c8f 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -64,6 +64,7 @@ refuse_ungated_apply, run_verification, verification_id_for, + verification_path, ) PROMPT = "Rules:\n- Be concise.\n- Cite the doc.\n" @@ -645,6 +646,40 @@ def test_interrupted_verify_does_not_duplicate_the_experiment(tmp_path: Path) -> ) +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) From 250a2e1713ff1f40983196c3bbce743ef882bd97 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 17:47:09 +0000 Subject: [PATCH 22/23] fix(kitaru): unclassifiable scores are SCORE_UNCLASSIFIED classify_scores treated a map_score drop-reason as unchanged when both passed flags were bools. That put incomparable scores in the comparable unchanged bucket. Return None so classify_replay_session records SCORE_UNCLASSIFIED. Pass-flip still classifies earlier. Co-authored-by: Dickson Neoh --- src/tracegrad/integrations/kitaru/backend.py | 2 -- tests/test_kitaru_verify.py | 15 +++++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/src/tracegrad/integrations/kitaru/backend.py b/src/tracegrad/integrations/kitaru/backend.py index 6e37f6e..637dd4c 100644 --- a/src/tracegrad/integrations/kitaru/backend.py +++ b/src/tracegrad/integrations/kitaru/backend.py @@ -218,8 +218,6 @@ def classify_scores( base_score = map_score(baseline) cand_score = map_score(candidate) if isinstance(base_score, str) or isinstance(cand_score, str): - if isinstance(base_passed, bool) and isinstance(cand_passed, bool): - return "unchanged" return None if cand_score > base_score: return "improved" diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 57b2c8f..05abea6 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -542,6 +542,21 @@ def test_unclassified_scores_are_score_unclassified() -> None: 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( From 2e22d25e6de6942124b1236f0e14d4a2327061ad Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 26 Aug 2026 17:56:13 +0000 Subject: [PATCH 23/23] fix(kitaru): verify exits 0 only when status is completed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Partial previously shared the success path with completed, so verify && apply could write from a half-run. Exit 0 only for completed; partial, canceled, failed, and incomplete stay non-zero. The hash+cohort apply gate is unchanged — a human can still apply after REVIEW when that policy allows. Co-authored-by: Dickson Neoh --- src/tracegrad/cli.py | 14 ++++++- tests/test_kitaru_verify.py | 79 +++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/src/tracegrad/cli.py b/src/tracegrad/cli.py index 9e7612d..30e5552 100644 --- a/src/tracegrad/cli.py +++ b/src/tracegrad/cli.py @@ -484,7 +484,19 @@ def command_verify(args: argparse.Namespace, out: TextIO) -> int: ), file=out, ) - return 0 if result.status != "failed" else 1 + 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: diff --git a/tests/test_kitaru_verify.py b/tests/test_kitaru_verify.py index 05abea6..f9c9f74 100644 --- a/tests/test_kitaru_verify.py +++ b/tests/test_kitaru_verify.py @@ -1210,6 +1210,85 @@ def test_verify_cli_refuses_a_stale_proposal_before_kitaru(tmp_path: Path) -> No 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(