From 074a27ada7b86ae435e1e43fb93b74c4c6f3addc Mon Sep 17 00:00:00 2001 From: tzhouam Date: Mon, 28 Sep 2026 21:34:28 +0800 Subject: [PATCH] Add trace schema v1, replay and dataset export Every model call, decision and later outcome is now kept durably, so the expensive models' work can be replayed and turned into datasets for cheaper models plus workflows. The copilot and the review bot share one schema. - trace_store (trace/1): append-only JSONL records (the source of truth), content-addressed gzip blobs for prompts and replies (records hold only sha256 refs; reads re-verify the hash), and a SQLite index by kind, run, repo, change set, PR, rule ID, role, model and outcome. The index can be rebuilt from the JSONL. trace_context binds run/playbook/repo/change-set context through contextvars. Every string of a record (keys, model and usage included) and every blob is redacted before it is hashed or written: known token shapes, private-key blocks, and the value of every environment variable whose name marks it as a secret. It uses only the standard library and is exported through sdk.v1 (TraceStore, TRACE_SCHEMA, trace_context, redact) without loading private modules. - knowledge service: - ModelGateway records every call with the pinned provider/model/effort and the call's final verdict: transport failures and truncated, empty, unparseable or schema-failing replies are recorded as failures. - gate_and_stage allocates the change set ID first, judges under its context, then writes a decision record. - merge.advance writes outcomes (merged, closed_unmerged, head_changed). - retirements write rule_retired. - intake, sweep and calibration are bound to their playbooks. Generator calls precede their change set, so each drafting has a unique draft key (event, or release + page) and each call its attempt number; the decision lists only the ACCEPTED attempt, so rejected replies and earlier tries never inherit its labels. - Tracing never breaks the service. - replay: re-asks a recorded call with a substitute model and compares the decision-bearing fields. It touches no GitHub, outbox or ledger; the substitute call and the comparison are traced too. - export: judge or generator rows of (input, expensive output, gate decision, outcomes, later rule outcomes), generator rows joined through draft keys. Calibration and replay calls, failed calls, and any prompt that cites a calibration case's evidence reference anywhere or its title are dropped. - CLI: kb traces, kb replay --record ID --model P:M[:E], kb export --out F. Co-Authored-By: Claude Opus 5.5 --- doc/architecture/SPEC/kb_service.md | 14 + doc/architecture/SPEC/sdk.md | 4 + doc/architecture/SPEC/trace_store.md | 24 ++ doc/features/knowledge-service.md | 11 + .../kb_service/calibration.py | 8 + src/infermatrix_copilot/kb_service/cli.py | 51 ++++ src/infermatrix_copilot/kb_service/intake.py | 8 +- src/infermatrix_copilot/kb_service/merge.py | 23 ++ src/infermatrix_copilot/kb_service/models.py | 35 ++- src/infermatrix_copilot/kb_service/replay.py | 132 +++++++++ src/infermatrix_copilot/kb_service/runtime.py | 101 +++++-- src/infermatrix_copilot/kb_service/sweep.py | 46 ++- src/infermatrix_copilot/sdk/v1/__init__.py | 5 + src/infermatrix_copilot/trace_store.py | 280 ++++++++++++++++++ test/test_kb_intake_gate.py | 2 +- test/test_trace_store.py | 256 ++++++++++++++++ 16 files changed, 956 insertions(+), 44 deletions(-) create mode 100644 doc/architecture/SPEC/trace_store.md create mode 100644 src/infermatrix_copilot/kb_service/replay.py create mode 100644 src/infermatrix_copilot/trace_store.py create mode 100644 test/test_trace_store.py diff --git a/doc/architecture/SPEC/kb_service.md b/doc/architecture/SPEC/kb_service.md index cbd7f08e..b26932bd 100644 --- a/doc/architecture/SPEC/kb_service.md +++ b/doc/architecture/SPEC/kb_service.md @@ -81,3 +81,17 @@ stdlib + PyYAML + `cryptography`（`kb` extra）+ `.adapters` + `.knowledge_serv 既不执行也不回执。每轮与每个决定都写入 `traces/publisher.jsonl`。 - 测试：`test_kb_publisher.py`（端到端 open_pr → ack → pr_open、dry-run、崩溃后不重复执行、伪造/未配置/过期、 越界路径、verdict/入队/暂停/关闭、暂停期间仍执行暂停、SSH 引号与回执、ack 字段）。 + +## 2026-09-28 留痕、回放与数据集 +- 运行时持有 `TraceStore(state_dir/traces)`：`ModelGateway` 的每次调用（含失败）写为 `model_call`； + `gate_and_stage` 先分配变更集 ID，判定期间绑定 `changeset_id`/`rule_ids`，暂存后写 `decision`； + `merge.advance` 写 `outcome`（merged / closed_unmerged / head_changed），退役写 `rule_retired`。 + intake、sweep、calibration 分别绑定 playbook `kb-intake`/`kb-sweep`/`kb-calibrate`。留痕失败从不中断服务。 +- `replay`：用记录中的原始 system/prompt 询问替代模型，比较决定性字段（`dimensions` 或 `verdict`），写 + `replay` 记录；替代模型的调用本身也记为 playbook `kb-replay`。只调用模型，不触及 GitHub、outbox 或账本。 +- `export_dataset`：按角色导出（输入、高成本模型输出、变更集判定、事后结果、规则后续结果）；丢弃校准与回放 + 调用、失败调用（截断、空、无法解析或不符 schema 的回复都记为失败），以及 prompt 中任意位置出现校准用例证据 + 引用（如 `PR #8107`，按边界匹配）或标题的调用。生成器调用发生在变更集之前：每次起草有唯一 `draft_key`（intake 为事件、 + 巡检为发版 + 页面，均带随机后缀），只有被采纳的那次尝试（`draft_key#attempt`）列入之后 `decision` 记录的 + `draft_keys`，被拒绝的尝试与失败的重试不继承判定。记录的每个字符串（含键、model、usage）都经过脱敏。 +- CLI：`kb traces`、`kb replay --record ID --model P:M[:E]`、`kb export --out FILE`。 diff --git a/doc/architecture/SPEC/sdk.md b/doc/architecture/SPEC/sdk.md index 66c6802b..62e511b3 100644 --- a/doc/architecture/SPEC/sdk.md +++ b/doc/architecture/SPEC/sdk.md @@ -124,3 +124,7 @@ Direct 1.1 / Strict 1.3 accept typed `CarriedFinding` inputs (unique IDs, source ## 2026-09-28 退役规则不再提供 `DirectClient` 返回的文档内容与 `sha256` 均基于去掉退役规则后的服务文本（无退役规则时与磁盘字节一致）。 + +## 2026-09-28 trace/1 导出 +`sdk.v1` 导出 `TraceStore`、`TRACE_SCHEMA`、`trace_context`、`redact`（见 `trace_store.md`），RB 用同一 schema +记录 review 的模型调用、判定与结果；导入它们不加载 `config`/`contract`/`mcp_server` 等私有模块。 diff --git a/doc/architecture/SPEC/trace_store.md b/doc/architecture/SPEC/trace_store.md new file mode 100644 index 00000000..9ee8c563 --- /dev/null +++ b/doc/architecture/SPEC/trace_store.md @@ -0,0 +1,24 @@ +# trace_store.py —— 规范 + + + +`LOC ~260 · trace/1：Copilot 与 RB 共用的模型调用/判定/结果记录 · refactor-status: stable` + +## 职责 +- `TraceStore(root)`：追加式 `records/<日期>.jsonl`（事实来源）+ 内容寻址 blob + （`blobs//.gz`，gzip，读取时复核哈希）+ SQLite 索引 `index.db` + （按 kind、run、仓库、变更集、PR、规则 ID、角色、模型、结果查询；`rebuild_index` 可从 JSONL 重建）。 +- 记录（`schema: trace/1`）：`kind` ∈ `model_call`/`decision`/`outcome`/`replay`；`context`（run_id、 + playbook、step、repo、changeset_id、pr、rule_ids…）、`model`（role、provider、model、effort、served_model）、 + `usage`、`seconds`、`inputs`/`outputs`（只存 blob 引用，从不内联正文）、`result`、`error`、`env` + （copilot 版本与 SHA）。`outcome` 是事后金标（merged、closed_unmerged、head_changed、rule_retired）。 +- `trace_context(**fields)`：用 contextvars 绑定上下文，嵌套叠加，线程/协程安全。 + +## 不变量 +- 写入前统一脱敏：已知令牌形态（GitHub、Anthropic、OpenAI、AWS、Slack、Bearer、私钥块）以及名字表明是 + 密钥的环境变量的值（路径形式的值除外）。脱敏发生在哈希之前，blob 与 JSONL 中都不会出现密钥。 +- 只依赖标准库；经 `sdk.v1` 导出（`TraceStore`、`TRACE_SCHEMA`、`trace_context`、`redact`）时不加载私有模块。 + +## 测试 +`test_trace_store.py`（schema 与去重 blob、索引与重建、脱敏、网关记录成功与失败调用、回放、导出的判定/结果 +关联与校准集泄漏过滤、SDK 导出、知识服务的 decision/outcome 记录）。 diff --git a/doc/features/knowledge-service.md b/doc/features/knowledge-service.md index cdce9ffd..5d34c668 100644 --- a/doc/features/knowledge-service.md +++ b/doc/features/knowledge-service.md @@ -106,3 +106,14 @@ ALLOW_POST=1 ALLOW_PUSH=1 infermatrix-copilot kb publish --remote bot-host:/path 没有 `ALLOW_POST=1` 时只把将要执行的动作写入 `$KB_PUBLISHER_STATE/traces/publisher.jsonl`；推送分支还需要 `ALLOW_PUSH=1`。控制记录超过 10 分钟未更新（服务可能宕机）时整轮不执行任何动作。 + +## 留痕、回放与数据集（trace/1） + +知识服务的每次模型调用（含失败）、每个质量门判定与每个事后结果都写入 `$KB_STATE_DIR/traces/` +（JSONL + 内容寻址 blob + SQLite 索引，写入前脱敏）。RB 通过 `sdk.v1.TraceStore` 使用同一 schema。 + +```bash +infermatrix-copilot kb traces --changeset # 查询记录 +infermatrix-copilot kb replay --record --model codex:gpt-6-mini:low # 用便宜模型重放一次调用 +infermatrix-copilot kb export --out judge.jsonl # 导出数据集（自动剔除校准集相关调用） +``` diff --git a/src/infermatrix_copilot/kb_service/calibration.py b/src/infermatrix_copilot/kb_service/calibration.py index 34e5cdcc..f5bdfb28 100644 --- a/src/infermatrix_copilot/kb_service/calibration.py +++ b/src/infermatrix_copilot/kb_service/calibration.py @@ -67,6 +67,14 @@ def load_cases(directory: str | Path) -> list[dict]: def run_calibration(directory: str | Path, *, gateway, judge) -> CalibrationReport: + """Traced as playbook ``kb-calibrate``: dataset export never uses these calls.""" + from ..trace_store import trace_context + + with trace_context(playbook="kb-calibrate", run_id=f"calibrate-{Path(directory).name}"): + return _run_calibration(directory, gateway=gateway, judge=judge) + + +def _run_calibration(directory: str | Path, *, gateway, judge) -> CalibrationReport: details = [] bad_total = bad_caught = good_total = good_rejected = 0 for case in load_cases(directory): diff --git a/src/infermatrix_copilot/kb_service/cli.py b/src/infermatrix_copilot/kb_service/cli.py index 65efd611..4d831523 100644 --- a/src/infermatrix_copilot/kb_service/cli.py +++ b/src/infermatrix_copilot/kb_service/cli.py @@ -10,6 +10,9 @@ kb serve [--once] the scheduler (holds the single-writer lease) kb activate activate the knowledge snapshot of the repository's main now kb rollback --to SHA point `active` back at an earlier snapshot + kb traces [--kind K] [--changeset ID] [--rule ID] [--limit N] query trace/1 records + kb replay --record ID --model PROVIDER:MODEL[:EFFORT] re-ask a recorded call + kb export --out FILE [--role judge|generator] dataset (calibration-safe) kb publish (--remote HOST:/STATE_DIR | --local DIR) [--once] the publisher (GPU box, owner's gh login): perform the signed outbox items; writes need ALLOW_POST=1 @@ -110,6 +113,18 @@ def main(argv: list[str] | None = None) -> int: sub.add_parser("activate") rollback = sub.add_parser("rollback") rollback.add_argument("--to", required=True) + traces = sub.add_parser("traces") + traces.add_argument("--kind") + traces.add_argument("--changeset") + traces.add_argument("--rule") + traces.add_argument("--run") + traces.add_argument("--limit", type=int, default=50) + replay = sub.add_parser("replay") + replay.add_argument("--record", required=True) + replay.add_argument("--model", required=True, help="provider:model[:effort], e.g. codex:gpt-6-mini:low") + export = sub.add_parser("export") + export.add_argument("--out", required=True) + export.add_argument("--role", default="judge", choices=("judge", "generator")) publish = sub.add_parser("publish") where = publish.add_mutually_exclusive_group(required=True) where.add_argument("--remote", help="host:/absolute/path of the service state directory (over ssh)") @@ -129,6 +144,8 @@ def main(argv: list[str] | None = None) -> int: if args.command == "publish": return _publish(args) + if args.command in {"traces", "replay", "export"}: + return _traces_command(args, _state_dir(args.state_dir)) state_dir = _state_dir(args.state_dir) ledger = _ledger(state_dir) @@ -271,3 +288,37 @@ def _publish(args) -> int: return 0 publisher.serve(interval=args.interval) return 0 + + +def _traces_command(args, state_dir: Path) -> int: + """Read-only over the trace store (replay also calls the substitute model).""" + from ..trace_store import TraceStore + + store = TraceStore(state_dir / "traces") + if args.command == "traces": + records = store.query(kind=args.kind, changeset_id=args.changeset, rule_id=args.rule, + run_id=args.run, limit=args.limit) + for record in records: + print(json.dumps({k: record[k] for k in ("id", "kind", "at", "context", "model", "result", "error")}, + ensure_ascii=False, sort_keys=True)) + return 0 + if args.command == "replay": + from ..config import Settings + from .models import ModelGateway, ModelRole + from .replay import replay + from .runtime import trace_recorder + + role_name = (store.get(args.record).get("model") or {}).get("role") or "judge" + gateway = ModelGateway(Settings(), recorder=trace_recorder(store)) # the replay call is traced too + result = replay(store, args.record, gateway, ModelRole.parse(role_name, args.model)) + json.dump(result, sys.stdout, ensure_ascii=False, indent=2, sort_keys=True) + print() + return 0 + from ..sdk._resources import adapters_root + from .replay import export_dataset + + adapters = Path(os.environ.get("ADAPTERS_DIR") or adapters_root()) + calibration = [p for p in sorted(adapters.glob("*/kb-calibration")) if p.is_dir()] + json.dump(export_dataset(store, args.out, role=args.role, calibration_dirs=calibration), sys.stdout) + print() + return 0 diff --git a/src/infermatrix_copilot/kb_service/intake.py b/src/infermatrix_copilot/kb_service/intake.py index 571980b6..bd7dd36e 100644 --- a/src/infermatrix_copilot/kb_service/intake.py +++ b/src/infermatrix_copilot/kb_service/intake.py @@ -166,10 +166,13 @@ def draft_changes(*, repo: str, repo_dir: str, event_id: int, evidence: dict, prompt = draft_prompt(repo, evidence, files, repo_dir) attempts: list[dict] = [] feedback = "" + from ..trace_store import accept_attempt, trace_context + for attempt in range(MAX_REPAIRS + 1): try: - reply = gateway.call_json(generator, system=SYSTEM, prompt=prompt + feedback, - validate=_validate_reply) + with trace_context(attempt=attempt): + reply = gateway.call_json(generator, system=SYSTEM, prompt=prompt + feedback, + validate=_validate_reply) except ModelUnavailable as exc: if "failed its schema" not in str(exc): raise # the model itself is unavailable: the event waits @@ -194,6 +197,7 @@ def draft_changes(*, repo: str, repo_dir: str, event_id: int, evidence: dict, "Fix exactly that and answer again with the full JSON object.") attempts.append({"attempt": attempt, "error": str(exc)}) continue + accept_attempt(attempt) # only this call's reply became the change return Draft([event_id], operations, result, str(reply.data.get("rationale") or ""), attempts) attempts.append({"attempt": attempt, "error": feedback.strip()}) return Draft([event_id], [], None, "rejected after repairs", attempts, rejected=True) diff --git a/src/infermatrix_copilot/kb_service/merge.py b/src/infermatrix_copilot/kb_service/merge.py index d69f5e97..d0a2eec1 100644 --- a/src/infermatrix_copilot/kb_service/merge.py +++ b/src/infermatrix_copilot/kb_service/merge.py @@ -159,14 +159,17 @@ def advance(rt, lifecycle) -> list[str]: rt.ledger.update_changeset(changeset["id"], status="merged", merge_sha=str(pr.get("merge_commit_sha") or "")) record_retirements(rt, changeset) + _outcome(rt, changeset, "merged", merge_sha=str(pr.get("merge_commit_sha") or "")) events.append(f"merged {changeset['id']}") continue if pr.get("state") == "closed": rt.ledger.update_changeset(changeset["id"], status="closed") + _outcome(rt, changeset, "closed_unmerged") # overturned by people: a negative label events.append(f"closed {changeset['id']}") continue head = str((pr.get("head") or {}).get("sha") or "") if head != changeset["head_sha"]: + _outcome(rt, changeset, "head_changed", head_sha=head) rt.ledger.update_changeset(changeset["id"], status="head_changed") rt.ledger.enqueue_human(lifecycle.repo, f"PR #{number} head changed after signing", changeset["id"]) events.append(f"head_changed {changeset['id']}") @@ -192,11 +195,31 @@ def advance(rt, lifecycle) -> list[str]: return events +def _rule_ids(changeset: dict) -> list[str]: + return sorted({op.get("new_rule_id") or op["rule_id"] for op in changeset["detail"].get("operations", []) + if op.get("new_rule_id") or op.get("rule_id")}) + + +def _outcome(rt, changeset: dict, outcome: str, **result) -> None: + """What finally happened to a change set: the gold label its decision is scored against.""" + trace = getattr(rt, "trace", None) + if trace is not None: + trace("outcome", context={"repo": changeset["repo"], "changeset_id": changeset["id"], + "pr": changeset.get("pr_number"), "rule_ids": _rule_ids(changeset)}, + result={"outcome": outcome, **result}) + + def record_retirements(rt, changeset: dict) -> None: release = changeset["detail"].get("release", "") for op in changeset["detail"].get("operations", []): if op["kind"] in ("retire", "replace"): rt.ledger.record_retirement(changeset["repo"], op["rule_id"], op["page"], release) + trace = getattr(rt, "trace", None) + if trace is not None: # the rule's earlier admission now has a later-retired label + trace("outcome", context={"repo": changeset["repo"], "changeset_id": changeset["id"], + "rule_ids": [op["rule_id"]]}, + result={"outcome": "rule_retired", "rule_id": op["rule_id"], "release": release, + "reason": op.get("reason", "")}) elif op["kind"] == "purge": rt.ledger.mark_purged(changeset["repo"], op["rule_id"]) diff --git a/src/infermatrix_copilot/kb_service/models.py b/src/infermatrix_copilot/kb_service/models.py index c7ce5ca5..8c5f4094 100644 --- a/src/infermatrix_copilot/kb_service/models.py +++ b/src/infermatrix_copilot/kb_service/models.py @@ -104,29 +104,46 @@ def call_json(self, role: ModelRole, *, system: str, prompt: str, validate: Callable[[dict], None] | None = None) -> ModelReply: transport = self._transport(role.provider) started = time.time() + identity = {"role": role.name, "requested": role.label(), "provider": role.provider, + "model": role.model, "effort": role.effort} try: reply = transport.complete( system=system, messages=[{"role": "user", "content": prompt}], model=role.model, effort=role.effort, role=role.name) except Exception as exc: # the transport's own failure modes + if self._recorder is not None: + self._recorder({**identity, "served_model": "", "stop_reason": "", "usage": {}, + "seconds": round(time.time() - started, 3), "system": system, + "prompt": prompt, "reply": "", "error": str(exc)[:2000]}) raise ModelUnavailable(f"{role.label()} failed: {exc}") from exc seconds = time.time() - started text = "".join(getattr(block, "text", "") or "" for block in getattr(reply, "blocks", [])) record = { - "role": role.name, "requested": role.label(), + **identity, "served_model": getattr(reply, "model", "") or "", "stop_reason": getattr(reply, "stop_reason", ""), "usage": dict(getattr(reply, "usage", {}) or {}), "seconds": round(seconds, 3), "system": system, "prompt": prompt, "reply": text, } - if self._recorder is not None: - self._recorder(record) + # the record carries the call's FINAL verdict: a truncated, empty, + # unparseable or schema-failing reply is recorded as a failure, so it + # never becomes a training example + failure: ModelUnavailable | None = None + data: dict = {} if getattr(reply, "stop_reason", "") == "max_tokens" or not text.strip(): - raise ModelUnavailable(f"{role.label()} timed out or returned nothing") - data = parse_json_object(text) - if validate is not None: + failure = ModelUnavailable(f"{role.label()} timed out or returned nothing") + else: try: - validate(data) - except Exception as exc: # any malformed shape is a controlled refusal - raise ModelUnavailable(f"{role.label()} reply failed its schema: {exc!r}") from exc + data = parse_json_object(text) + if validate is not None: + try: + validate(data) + except Exception as exc: # any malformed shape is a controlled refusal + raise ModelUnavailable(f"{role.label()} reply failed its schema: {exc!r}") from exc + except ModelUnavailable as exc: + failure = exc + if self._recorder is not None: + self._recorder({**record, "error": str(failure) if failure else ""}) + if failure is not None: + raise failure return ModelReply(role, data, text, record["served_model"], record["usage"], seconds) diff --git a/src/infermatrix_copilot/kb_service/replay.py b/src/infermatrix_copilot/kb_service/replay.py new file mode 100644 index 00000000..0e1cd447 --- /dev/null +++ b/src/infermatrix_copilot/kb_service/replay.py @@ -0,0 +1,132 @@ +"""Replay recorded model calls with another model, and export datasets. + +Both work only from the trace store (``trace/1``): a recorded ``model_call`` +carries its exact system prompt and prompt as blobs, so a replay sends the +same input to a substitute model and compares structured answers. Nothing here +can reach GitHub, the outbox or the ledger; the only side effects are the +substitute model call itself and a ``replay`` record. + +Dataset export turns judge (or generator) calls into rows of (input, the +expensive model's output, the gate decision, the later outcome) for prompt +compression, few-shot selection or distillation. Rows that could leak the +calibration set are dropped before anything is written: calls made by a +calibration run, and prompts citing a calibration case's evidence. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Iterable + +from .calibration import load_cases +from .models import ModelRole, ModelUnavailable, parse_json_object + +REPLAY_PLAYBOOK = "kb-replay" +EXCLUDED_PLAYBOOKS = ("kb-calibrate", REPLAY_PLAYBOOK) + + +def comparable(data: dict) -> dict: + """The decision-bearing part of a structured reply.""" + if "dimensions" in data: + return {"dimensions": data.get("dimensions")} + if "verdict" in data: + return {"verdict": data.get("verdict")} + return data + + +def replay(store, record_id: str, gateway, role: ModelRole) -> dict: + from ..trace_store import trace_context + + record = store.get(record_id) + if record["kind"] != "model_call": + raise ValueError(f"{record_id} is a {record['kind']} record, not a model call") + system, prompt = store.blob(record["inputs"]["system"]), store.blob(record["inputs"]["prompt"]) + try: + recorded = None if record.get("error") else \ + comparable(parse_json_object(store.blob(record["outputs"]["reply"]))) + except (ModelUnavailable, ValueError, KeyError): + recorded = None # the recorded call failed: the replay still runs and is recorded + with trace_context(playbook=REPLAY_PLAYBOOK, replay_of=record_id, run_id=f"replay-{record_id}"): + try: + reply = gateway.call_json(role, system=system, prompt=prompt) + replayed, error = comparable(reply.data), "" + except ModelUnavailable as exc: + replayed, error = None, str(exc) + result = {"replay_of": record_id, "recorded_model": record.get("model", {}), "replay_model": role.label(), + "recorded": recorded, "replayed": replayed, + "agree": recorded is not None and replayed is not None and recorded == replayed} + store.append("replay", context={"replay_of": record_id, "playbook": REPLAY_PLAYBOOK}, + model={"role": role.name, "provider": role.provider, "model": role.model, "effort": role.effort}, + result=result, error=error) + return {**result, "error": error} + + +def calibration_fingerprints(directories: Iterable[str | Path]) -> list[re.Pattern]: + """Patterns whose match in a prompt means it overlaps a calibration case: + each case's evidence reference ANYWHERE in the text (a rule citing + ``^[PR #8107]`` counts, ``PR #81070`` does not) and its evidence title.""" + marks: list[re.Pattern] = [] + for directory in directories: + for case in load_cases(directory): + for item in case.get("evidence") or []: + reference = str(item.get("source_reference") or "").strip() + if reference: + marks.append(re.compile(rf"(?= 20: + marks.append(re.compile(re.escape(title))) + return marks + + +def export_dataset(store, out_path: str | Path, *, role: str = "judge", + calibration_dirs: Iterable[str | Path] = ()) -> dict: + """Write one JSON line per usable model call of ``role``; returns counts.""" + marks = calibration_fingerprints(calibration_dirs) + decisions: dict[str, dict] = {} + by_draft: dict[str, dict] = {} # generator calls precede their change set: joined by draft key + for record in store.query(kind="decision", limit=10**6): + decisions[record["context"].get("changeset_id")] = record["result"] + for key in record["context"].get("draft_keys") or []: + by_draft[key] = record["context"] + outcomes: dict[str, list[str]] = {} + rule_outcomes: dict[str, list[str]] = {} # e.g. rule_retired later, by another change set + for record in store.query(kind="outcome", limit=10**6): + result = record["result"] + if result.get("outcome") == "rule_retired": + rule_outcomes.setdefault(str(result.get("rule_id")), []).append("rule_retired") + else: + outcomes.setdefault(record["context"].get("changeset_id") or "", []).append(result["outcome"]) + counts = {"written": 0, "leak_dropped": 0, "calibration_dropped": 0, "failed_dropped": 0} + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as handle: + for record in store.query(kind="model_call", role=role, limit=10**6): + context = record.get("context") or {} + if context.get("playbook") in EXCLUDED_PLAYBOOKS: + counts["calibration_dropped"] += 1 + continue + if record.get("error"): + counts["failed_dropped"] += 1 + continue + system, prompt = store.blob(record["inputs"]["system"]), store.blob(record["inputs"]["prompt"]) + if any(mark.search(prompt) or mark.search(system) for mark in marks): + counts["leak_dropped"] += 1 + continue + call_key = f"{context['draft_key']}#{context['attempt']}" \ + if context.get("draft_key") and "attempt" in context else "" + staged = by_draft.get(call_key, {}) + changeset_id = context.get("changeset_id") or staged.get("changeset_id") or "" + context = {**context, "rule_ids": context.get("rule_ids") or staged.get("rule_ids") or []} + handle.write(json.dumps({ + "id": record["id"], "task": f"{role}:{context.get('step') or context.get('playbook') or ''}", + "repo": context.get("repo", ""), "changeset_id": changeset_id, + "rule_ids": context.get("rule_ids") or [], "model": record.get("model", {}), + "system": system, "prompt": prompt, "output": store.blob(record["outputs"]["reply"]), + "decision": decisions.get(changeset_id), "outcomes": outcomes.get(changeset_id, []), + "rule_outcomes": {rid: rule_outcomes[rid] for rid in context.get("rule_ids") or [] + if rid in rule_outcomes}, + }, ensure_ascii=False) + "\n") + counts["written"] += 1 + return counts diff --git a/src/infermatrix_copilot/kb_service/runtime.py b/src/infermatrix_copilot/kb_service/runtime.py index 1845d2a0..5602e733 100644 --- a/src/infermatrix_copilot/kb_service/runtime.py +++ b/src/infermatrix_copilot/kb_service/runtime.py @@ -12,6 +12,7 @@ import json import os import time +import uuid from dataclasses import dataclass, field from pathlib import Path from typing import Callable @@ -40,6 +41,7 @@ class KbRuntime: clock: Callable[[], float] = time.time lease_owner: str | None = None # set by `kb serve`, which holds the lease for its lifetime publisher_public_key: object | None = None + traces: object | None = None # trace_store.TraceStore: decisions and outcomes (model calls via the gateway) release_label: Callable[[str], str] = field(default=lambda repo: "") @classmethod @@ -66,10 +68,13 @@ def from_env(cls, settings, *, state_dir: Path | None = None) -> "KbRuntime": from ..knowledge_service.signing import load_public_key publisher_key = load_public_key(Path(os.environ["KB_PUBLISHER_PUBKEY"]).read_text(encoding="utf-8")) + from ..trace_store import TraceStore + + traces = TraceStore(state_dir / "traces") return cls( - publisher_public_key=publisher_key, + publisher_public_key=publisher_key, traces=traces, state_dir=state_dir, ledger=ledger, registry=registry, - gateway=ModelGateway(settings, recorder=_trace_recorder(state_dir)), + gateway=ModelGateway(settings, recorder=trace_recorder(traces)), generator=generator, judge=judge, knowledge=KnowledgeRepo(Path(os.environ.get("KB_KNOWLEDGE_CLONE") or state_dir / "knowledge-repo")), github=GitHubReader(), outbox=outbox, @@ -77,6 +82,15 @@ def from_env(cls, settings, *, state_dir: Path | None = None) -> "KbRuntime": # -- helpers ------------------------------------------------------------- + def trace(self, kind: str, **fields) -> None: + """Append a trace/1 record; never lets tracing break the service.""" + if self.traces is None: + return + try: + self.traces.append(kind, **fields) + except Exception: # noqa: BLE001 - a full disk must not stop the gate + pass + def today(self) -> str: return dt.datetime.fromtimestamp(self.clock(), dt.timezone.utc).date().isoformat() @@ -95,14 +109,21 @@ def load_changeset_files(self, changeset_id: str) -> dict: return json.loads(self.changeset_path(changeset_id).read_text(encoding="utf-8")) -def _trace_recorder(state_dir: Path): - """Append every model call (inputs, outputs, usage) to the service trace.""" - path = state_dir / "traces" / "model_calls.jsonl" +def trace_recorder(traces): + """Every model call (inputs, outputs, usage, failure) as a trace/1 + ``model_call`` record, under the caller's bound trace context.""" def record(entry: dict) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("a", encoding="utf-8") as handle: - handle.write(json.dumps({"at": time.time(), **entry}, ensure_ascii=False) + "\n") + try: + traces.append( + "model_call", + inputs={"system": entry["system"], "prompt": entry["prompt"]}, + outputs={"reply": entry["reply"]}, + model={k: entry.get(k, "") for k in ("role", "provider", "model", "effort", "served_model")}, + usage=entry.get("usage") or {}, seconds=entry.get("seconds"), + result={"stop_reason": entry.get("stop_reason", "")}, error=entry.get("error", "")) + except Exception: # noqa: BLE001 - tracing never breaks a model call + pass return record @@ -145,11 +166,14 @@ def run_intake(rt: KbRuntime, lifecycle: RepoLifecycle, *, max_events: int = 10) request duplicate PRs).""" from .ledger import LeaseError + from ..trace_store import trace_context + try: - if rt.lease_owner: - return _run_intake_locked(rt, lifecycle, max_events=max_events, owner=rt.lease_owner) - with rt.ledger.lease() as owner: - return _run_intake_locked(rt, lifecycle, max_events=max_events, owner=owner) + with trace_context(playbook="kb-intake", repo=lifecycle.repo, run_id=f"intake-{uuid.uuid4().hex[:12]}"): + if rt.lease_owner: + return _run_intake_locked(rt, lifecycle, max_events=max_events, owner=rt.lease_owner) + with rt.ledger.lease() as owner: + return _run_intake_locked(rt, lifecycle, max_events=max_events, owner=owner) except LeaseError: return None @@ -166,13 +190,23 @@ def _run_intake_locked(rt: KbRuntime, lifecycle: RepoLifecycle, *, max_events: i return None # retried next tick; events stay pending release, today = rt.release_for(lifecycle.repo), rt.today() drafts = [] + from ..trace_store import accepted_key, trace_context + + accepted: dict[int, str] = {} # event -> the accepted generator call's key for event in events: rt.ledger.heartbeat(owner) # model calls are slow; keep the lease live try: - drafts.append(draft_changes( - repo=lifecycle.repo, repo_dir=lifecycle.knowledge_dir, event_id=event["id"], - evidence=event["payload"], files=base, gateway=rt.gateway, - generator=rt.generator, release=release, today=today)) + # the change set does not exist yet: the draft key and the accepted + # attempt link exactly the call whose reply became the change to + # the decision that later stages it + key, holder = draft_key_for_event(lifecycle.repo, event["id"]), {} + with trace_context(draft_key=key, _accepted=holder, step="draft"): + drafts.append(draft_changes( + repo=lifecycle.repo, repo_dir=lifecycle.knowledge_dir, event_id=event["id"], + evidence=event["payload"], files=base, gateway=rt.gateway, + generator=rt.generator, release=release, today=today)) + if accepted_key(key, holder): + accepted[event["id"]] = accepted_key(key, holder) except ModelUnavailable as exc: # fenced: if this worker lost its lease meanwhile, the new holder may # already have staged the event; an unfenced reset would undo that @@ -201,24 +235,35 @@ def _run_intake_locked(rt: KbRuntime, lifecycle: RepoLifecycle, *, max_events: i evidence = [e["payload"] for e in events if e["id"] in kept_ids] return gate_and_stage(rt, lifecycle, owner, kind="intake", base=base, base_sha=base_sha, external=external, operations=operations, result=result, - evidence=evidence, event_ids=sorted(kept_ids), release=release) + evidence=evidence, event_ids=sorted(kept_ids), release=release, + draft_keys=[accepted[i] for i in sorted(kept_ids) if i in accepted]) + + +def draft_key_for_event(repo: str, event_id) -> str: + """Unique per drafting: a redrafted event never shares a key with an earlier try.""" + return f"event:{repo}:{event_id}:{uuid.uuid4().hex[:8]}" def gate_and_stage(rt: KbRuntime, lifecycle: RepoLifecycle, owner: str, *, kind: str, base: dict, base_sha: str, external: dict, operations, result, evidence: list[dict], event_ids: list[int], release: str, force_human: str = "", - hold: bool = False) -> str: + hold: bool = False, draft_keys: list[str] = ()) -> str: """Run the quality gate on one change set and stage it (files first, then every ledger write in one transaction fenced on the lease). Used by intake, sweep and purge. ``force_human`` routes a passing change set to people (e.g. a sweep-wide circuit breaker).""" + from ..trace_store import trace_context + head = {**base, **result.files} rt.ledger.heartbeat(owner) - decision = run_gate( - base=base, head=head, changes=changes_between(base, head), external_texts=external, - evidence=evidence, gateway=rt.gateway, judge=rt.judge, release=release, - repo_dir=lifecycle.knowledge_dir, protected_rules=lifecycle.protected_rules, - retire_ratio=lifecycle.retire_ratio, max_files=lifecycle.max_files) + changeset_id = rt.ledger.new_changeset_id(lifecycle.repo, kind) + rule_ids = sorted({op.new_rule_id or op.rule_id for op in operations if (op.new_rule_id or op.rule_id)}) + with trace_context(changeset_id=changeset_id, rule_ids=rule_ids, step="gate"): + decision = run_gate( + base=base, head=head, changes=changes_between(base, head), external_texts=external, + evidence=evidence, gateway=rt.gateway, judge=rt.judge, release=release, + repo_dir=lifecycle.knowledge_dir, protected_rules=lifecycle.protected_rules, + retire_ratio=lifecycle.retire_ratio, max_files=lifecycle.max_files) if force_human and decision.status == "pass": decision.status = "human" decision.reasons.append(force_human) @@ -228,7 +273,6 @@ def gate_and_stage(rt: KbRuntime, lifecycle: RepoLifecycle, owner: str, *, kind: verdicts = [{"layer": "L2", "verdict": b.verdict, "block_id": b.block.block_id, "model": b.model, "detail": b.to_dict()} for b in decision.blocks] verdicts.append({"layer": "gate", "verdict": decision.status, "detail": {"reasons": decision.reasons}}) - changeset_id = rt.ledger.new_changeset_id(lifecycle.repo, kind) rt.save_changeset_files(changeset_id, {"base_sha": base_sha, "files": result.files, "deleted": [], "evidence": evidence}) try: @@ -244,6 +288,15 @@ def gate_and_stage(rt: KbRuntime, lifecycle: RepoLifecycle, owner: str, *, kind: except Exception: rt.changeset_path(changeset_id).unlink(missing_ok=True) raise + rt.trace("decision", context={"changeset_id": changeset_id, "rule_ids": rule_ids, "step": "gate", + "draft_keys": list(draft_keys)}, + model={"role": "judge", "model": rt.judge.label()}, + result={"status": decision.status, "staged_as": status, "reasons": decision.reasons, + "blocks": [{"block_id": b.block.block_id, "rule_id": b.block.rule_id, "op": b.block.op, + "verdict": b.verdict, "dimensions": b.dimensions} for b in decision.blocks], + "consistency": [{"owner_dir": c["owner_dir"], "verdict": c["verdict"]} + for c in decision.consistency], + "l1_issues": [i.to_dict() for i in decision.l1.issues]}) return changeset_id diff --git a/src/infermatrix_copilot/kb_service/sweep.py b/src/infermatrix_copilot/kb_service/sweep.py index 93873390..ed48c05b 100644 --- a/src/infermatrix_copilot/kb_service/sweep.py +++ b/src/infermatrix_copilot/kb_service/sweep.py @@ -20,6 +20,7 @@ import json import subprocess +import uuid from pathlib import Path, PurePosixPath from types import SimpleNamespace @@ -198,10 +199,13 @@ def sweep_page(rt, lifecycle, *, page: str, files: dict[str, str], diff: str, hi "<", "\\u003c") + "\n\n" feedback = "" last_error = "" - for _attempt in range(MAX_REPAIRS + 1): + from ..trace_store import accept_attempt, trace_context + + for attempt in range(MAX_REPAIRS + 1): try: - reply = rt.gateway.call_json(rt.generator, system=SWEEP_SYSTEM, prompt=prompt + feedback, - validate=_validate_sweep) + with trace_context(attempt=attempt): + reply = rt.gateway.call_json(rt.generator, system=SWEEP_SYSTEM, prompt=prompt + feedback, + validate=_validate_sweep) except ModelUnavailable as exc: if "failed its schema" not in str(exc): raise @@ -217,7 +221,9 @@ def sweep_page(rt, lifecycle, *, page: str, files: dict[str, str], diff: str, hi last_error = feedback.strip() continue try: - return operations, apply_operations(files, operations, release=release, today=today) + result = apply_operations(files, operations, release=release, today=today) + accept_attempt(attempt) + return operations, result except LifecycleError as exc: feedback = f"\n\nYour previous answer was rejected by the knowledge base: {exc}. Fix exactly that." last_error = str(exc) @@ -272,6 +278,25 @@ def _progress(rt, lifecycle, sweep: dict) -> dict: def run_sweep(rt, lifecycle, owner: str, sweep: dict, upstream: UpstreamRepo, audit_hints: dict[str, list[dict]] | None = None) -> dict: + """Traced: every model call and decision of the sweep carries its release.""" + from ..trace_store import trace_context + + with trace_context(playbook="kb-sweep", repo=lifecycle.repo, release=str(sweep.get("tag") or ""), + run_id=f"sweep-{lifecycle.repo}-{sweep.get('tag') or ''}"): + return _run_sweep(rt, lifecycle, owner, sweep, upstream, audit_hints) + + +def draft_key_for_page(repo: str, sweep: dict, page: str) -> str: + """Unique per evaluation: a retried page never shares a key with a failed try.""" + return f"sweep:{repo}:{sweep.get('tag') or ''}:{page}:{uuid.uuid4().hex[:8]}" + + +def _run_sweep(rt, lifecycle, owner: str, sweep: dict, upstream: UpstreamRepo, + audit_hints: dict[str, list[dict]] | None = None) -> dict: + from ..trace_store import accepted_key, trace_context + + accepted: dict[str, str] = {} # page -> the accepted generator call's key + """Run (or resume) one sweep. Progress is per page: a page is done once its answer is staged or it is kept; a page whose evaluation failed is retried on the next due run and handed to people after MAX_PAGE_ATTEMPTS. The baseline @@ -306,9 +331,13 @@ def run_sweep(rt, lifecycle, owner: str, sweep: dict, upstream: UpstreamRepo, diff_cache[prefixes] = upstream.diff(sweep["from_sha"], sweep["to_sha"], list(prefixes)) rt.ledger.heartbeat(owner) try: - outcome = sweep_page(rt, lifecycle, page=page, files=base, diff=diff_cache[prefixes], - hints=(audit_hints or {}).get(page, []), sweep=sweep, - release=release, today=today) + key, holder = draft_key_for_page(lifecycle.repo, sweep, page), {} + with trace_context(draft_key=key, _accepted=holder, step="draft"): + outcome = sweep_page(rt, lifecycle, page=page, files=base, diff=diff_cache[prefixes], + hints=(audit_hints or {}).get(page, []), sweep=sweep, + release=release, today=today) + if accepted_key(key, holder): + accepted[page] = accepted_key(key, holder) except SweepPageFailed as exc: attempts = progress["attempts"].get(page, 0) + 1 progress["attempts"][page] = attempts @@ -335,7 +364,8 @@ def run_sweep(rt, lifecycle, owner: str, sweep: dict, upstream: UpstreamRepo, kind = "purge" if operations and all(op.kind == "purge" for op in operations) else "sweep" changeset_id = gate_and_stage(rt, lifecycle, owner, kind=kind, base=base, base_sha=base_sha, external=external, operations=operations, result=result, - evidence=evidence, event_ids=[], release=release, hold=True) + evidence=evidence, event_ids=[], release=release, hold=True, + draft_keys=[accepted[key]] if key in accepted else []) progress.setdefault("held", []).append(changeset_id) if key == "#t1": progress["t1_done"] = True diff --git a/src/infermatrix_copilot/sdk/v1/__init__.py b/src/infermatrix_copilot/sdk/v1/__init__.py index ae83b91c..d6846e68 100644 --- a/src/infermatrix_copilot/sdk/v1/__init__.py +++ b/src/infermatrix_copilot/sdk/v1/__init__.py @@ -46,6 +46,7 @@ UnsupportedRepositoryError, ) from .strict import StrictRuntime +from ...trace_store import SCHEMA as TRACE_SCHEMA, TraceStore, redact, trace_context __all__ = [ "DIRECT_API_VERSION", @@ -53,6 +54,7 @@ "QUALITY_API_VERSION", "SDK_API_VERSION", "STRICT_API_VERSION", + "TRACE_SCHEMA", "Capabilities", "ChangedPath", "CarriedFinding", @@ -92,6 +94,9 @@ "StrictRunHandle", "StrictRuntimeConfig", "StrictRuntime", + "TraceStore", "UnsupportedRepositoryError", "get_capabilities", + "redact", + "trace_context", ] diff --git a/src/infermatrix_copilot/trace_store.py b/src/infermatrix_copilot/trace_store.py new file mode 100644 index 00000000..761ba9b2 --- /dev/null +++ b/src/infermatrix_copilot/trace_store.py @@ -0,0 +1,280 @@ +"""Trace schema v1: durable records of model calls, decisions and outcomes. + +Shared by the copilot and the review bot, so the expensive models' work can be +replayed and turned into datasets for cheaper models plus workflows. Layout +under a trace root:: + + records/.jsonl one JSON record per line (append-only) + blobs//.gz content-addressed payloads (prompts, replies) + index.db SQLite index: id, kind, run, repo, PR, rule + IDs, model, outcome -> the record itself + +A record (``schema: trace/1``):: + + {"schema", "id", "kind", "at", + "context": {run_id, playbook, step, repo, changeset_id, pr, rule_ids, ...}, + "model": {role, provider, model, effort, served_model}, # model_call + "usage", "seconds", + "inputs": {name: "sha256:"}, # blobs, never inline text + "outputs": {name: "sha256:"}, + "result": {...}, # decision / outcome / replay + "error": "", + "env": {copilot_version, copilot_sha, ...}} + +``kind`` is ``model_call``, ``decision`` (a gate/review verdict), ``outcome`` +(what happened later: merged, closed, retired, accepted...) or ``replay``. +Outcomes are the gold labels for evaluating a cheaper model. + +Everything is redacted before it is hashed or written: known token shapes +(GitHub, Anthropic, OpenAI, AWS, bearer headers, private keys) and the value +of every environment variable whose name marks it as a secret. + +Standard library only; importing it loads nothing else from the package. +""" + +from __future__ import annotations + +import contextvars +import datetime as dt +import gzip +import hashlib +import json +import os +import re +import sqlite3 +import subprocess +import threading +import time +import uuid +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, Mapping + +SCHEMA = "trace/1" +KINDS = ("model_call", "decision", "outcome", "replay") +_CONTEXT: contextvars.ContextVar[dict] = contextvars.ContextVar("trace_context", default={}) + +_SECRET_SHAPES = re.compile( + r"gh[pousr]_[A-Za-z0-9]{20,}" + r"|github_pat_[A-Za-z0-9_]{20,}" + r"|sk-ant-[A-Za-z0-9_-]{16,}" + r"|sk-(?:proj-)?[A-Za-z0-9_-]{20,}" + r"|(?:AKIA|ASIA)[A-Z0-9]{16}" + r"|xox[abpr]-[A-Za-z0-9-]{10,}" + r"|(?i:bearer)\s+[A-Za-z0-9._~+/=-]{16,}" + r"|-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----" +) +_SECRET_NAME = re.compile(r"(TOKEN|SECRET|PASSWORD|PASSWD|API_KEY|_KEY|_PAT|CREDENTIALS?)$", re.I) +REDACTED = "[REDACTED]" + + +def _secret_values(environ: Mapping[str, str]) -> list[str]: + values = {v for k, v in environ.items() if _SECRET_NAME.search(k) and len(v) >= 8 + and not v.startswith(("/", "~", "."))} # paths to key files are not secrets + return sorted(values, key=len, reverse=True) + + +def redact(text: str, environ: Mapping[str, str] | None = None) -> str: + """``text`` with secrets replaced by ``[REDACTED]``.""" + for value in _secret_values(os.environ if environ is None else environ): + text = text.replace(value, REDACTED) + return _SECRET_SHAPES.sub(REDACTED, text) + + +def _redact_obj(value: Any, environ: Mapping[str, str] | None = None) -> Any: + if isinstance(value, str): + return redact(value, environ) + if isinstance(value, dict): + return {(redact(k, environ) if isinstance(k, str) else k): _redact_obj(v, environ) + for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_redact_obj(v, environ) for v in value] + return value + + +# -- context ----------------------------------------------------------------------- + +@contextmanager +def trace_context(**fields: Any) -> Iterator[dict]: + """Bind context fields (run_id, playbook, repo, changeset_id, ...) to every + record written inside the block, nested blocks adding to outer ones.""" + merged = {**_CONTEXT.get(), **{k: v for k, v in fields.items() if v is not None}} + token = _CONTEXT.set(merged) + try: + yield merged + finally: + _CONTEXT.reset(token) + + +def current_context() -> dict: + """The bound context as recorded (private ``_`` fields are not recorded).""" + return {k: v for k, v in _CONTEXT.get().items() if not k.startswith("_")} + + +def accept_attempt(attempt: int) -> None: + """Mark which attempt of the current draft was accepted. The caller binds a + holder as ``_accepted={}``; the accepted call is then identified by its + ``draft_key`` and ``attempt``, and only that call inherits the decision.""" + holder = _CONTEXT.get().get("_accepted") + if isinstance(holder, dict): + holder["attempt"] = int(attempt) + + +def accepted_key(draft_key: str, holder: Mapping[str, Any]) -> str | None: + return f"{draft_key}#{holder['attempt']}" if "attempt" in holder else None + + +_ENV: dict | None = None + + +def environment_fingerprint() -> dict: + """Which code produced a record (computed once per process).""" + global _ENV + if _ENV is None: + from . import __version__ + + sha = "" + try: + proc = subprocess.run(["git", "-C", str(Path(__file__).resolve().parent), "rev-parse", "HEAD"], + capture_output=True, text=True, timeout=5, check=False) + sha = proc.stdout.strip() if proc.returncode == 0 else "" + except (OSError, subprocess.SubprocessError): + pass + _ENV = {"copilot_version": __version__, "copilot_sha": sha} + return dict(_ENV) + + +# -- the store ------------------------------------------------------------------------ + +class TraceStore: + def __init__(self, root: str | Path, *, clock=time.time, environ: Mapping[str, str] | None = None): + self.root = Path(root) + self._clock = clock + self._environ = environ + self._lock = threading.Lock() + + # blobs ----------------------------------------------------------------------- + def put_blob(self, content: str | bytes) -> str: + """Store redacted text; bytes must be UTF-8 text (redaction needs text).""" + if isinstance(content, bytes): + try: + content = content.decode("utf-8") + except UnicodeDecodeError as exc: + raise TypeError("trace blobs are text: binary payloads cannot be redacted") from exc + data = redact(content, self._environ).encode("utf-8") + digest = hashlib.sha256(data).hexdigest() + path = self.root / "blobs" / digest[:2] / f"{digest}.gz" + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(f".{uuid.uuid4().hex}.tmp") + tmp.write_bytes(gzip.compress(data, mtime=0)) + os.replace(tmp, path) + return f"sha256:{digest}" + + def blob(self, ref: str) -> str: + digest = ref.removeprefix("sha256:") + if not re.fullmatch(r"[0-9a-f]{64}", digest): + raise KeyError(ref) + data = gzip.decompress((self.root / "blobs" / digest[:2] / f"{digest}.gz").read_bytes()) + if hashlib.sha256(data).hexdigest() != digest: + raise ValueError(f"blob {ref} is corrupt") + return data.decode("utf-8") + + # records --------------------------------------------------------------------- + def _db(self) -> sqlite3.Connection: + self.root.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(self.root / "index.db", timeout=30) + conn.execute("""CREATE TABLE IF NOT EXISTS records ( + id TEXT PRIMARY KEY, kind TEXT, at REAL, run_id TEXT, playbook TEXT, repo TEXT, + changeset_id TEXT, pr INTEGER, rule_ids TEXT, role TEXT, model TEXT, outcome TEXT, + record TEXT)""") + for column in ("kind", "run_id", "repo", "changeset_id", "pr", "model", "outcome"): + conn.execute(f"CREATE INDEX IF NOT EXISTS records_{column} ON records({column})") + return conn + + def append(self, kind: str, *, inputs: Mapping[str, str] | None = None, + outputs: Mapping[str, str] | None = None, model: Mapping[str, Any] | None = None, + result: Mapping[str, Any] | None = None, usage: Mapping[str, Any] | None = None, + seconds: float | None = None, error: str = "", context: Mapping[str, Any] | None = None, + ) -> dict: + """Write one record; text in ``inputs``/``outputs`` goes to blobs.""" + if kind not in KINDS: + raise ValueError(f"unknown trace record kind: {kind}") + now = self._clock() + record = { + "schema": SCHEMA, "id": f"{int(now * 1000)}-{uuid.uuid4().hex[:12]}", "kind": kind, "at": now, + "context": {**current_context(), **(context or {})}, + "model": dict(model or {}), "usage": dict(usage or {}), + "seconds": seconds, + "inputs": {k: self.put_blob(v) for k, v in (inputs or {}).items()}, + "outputs": {k: self.put_blob(v) for k, v in (outputs or {}).items()}, + "result": dict(result or {}), + "error": error or "", + "env": environment_fingerprint(), + } + # every string of the record, keys included, before it is serialized or indexed + record = _redact_obj(record, self._environ) + line = json.dumps(record, ensure_ascii=False, sort_keys=True) + day = dt.datetime.fromtimestamp(now, dt.timezone.utc).date().isoformat() + with self._lock: + path = self.root / "records" / f"{day}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(line + "\n") + with self._db() as conn: + conn.execute("INSERT INTO records VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)", _row(record, line)) + return record + + def get(self, record_id: str) -> dict: + with self._db() as conn: + row = conn.execute("SELECT record FROM records WHERE id=?", (record_id,)).fetchone() + if row is None: + raise KeyError(record_id) + return json.loads(row[0]) + + def query(self, *, kind: str | None = None, run_id: str | None = None, repo: str | None = None, + changeset_id: str | None = None, pr: int | None = None, rule_id: str | None = None, + role: str | None = None, model: str | None = None, playbook: str | None = None, + limit: int = 1000) -> list[dict]: + clauses, args = [], [] + for column, value in (("kind", kind), ("run_id", run_id), ("repo", repo), + ("changeset_id", changeset_id), ("pr", pr), ("role", role), + ("model", model), ("playbook", playbook)): + if value is not None: + clauses.append(f"{column}=?") + args.append(value) + if rule_id is not None: + clauses.append("rule_ids LIKE ?") + args.append(f"%,{rule_id},%") + sql = "SELECT record FROM records" + (" WHERE " + " AND ".join(clauses) if clauses else "") \ + + " ORDER BY at, id LIMIT ?" + with self._db() as conn: + return [json.loads(row[0]) for row in conn.execute(sql, (*args, int(limit)))] + + def rebuild_index(self) -> int: + """Re-create index.db from the JSONL records (the source of truth).""" + (self.root / "index.db").unlink(missing_ok=True) + count = 0 + with self._db() as conn: + for path in sorted((self.root / "records").glob("*.jsonl")): + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + conn.execute("INSERT OR REPLACE INTO records VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)", + _row(json.loads(line), line)) + count += 1 + return count + + +def _row(record: dict, line: str) -> tuple: + ctx = record.get("context") or {} + model = record.get("model") or {} + result = record.get("result") or {} + rule_ids = ctx.get("rule_ids") or [] + pr = str(ctx.get("pr") or "") + return (record["id"], record["kind"], record["at"], str(ctx.get("run_id") or ""), + str(ctx.get("playbook") or ""), str(ctx.get("repo") or ""), str(ctx.get("changeset_id") or ""), + int(pr) if pr.isdigit() else None, + "," + ",".join(map(str, rule_ids)) + "," if rule_ids else "", + str(model.get("role") or ""), str(model.get("model") or ""), + str(result.get("outcome") or result.get("status") or ""), line) diff --git a/test/test_kb_intake_gate.py b/test/test_kb_intake_gate.py index 30f3ab43..24add7bb 100644 --- a/test/test_kb_intake_gate.py +++ b/test/test_kb_intake_gate.py @@ -261,7 +261,7 @@ def test_intake_end_to_end_shadow_records_and_publishes_nothing(tmp_path): assert "DEMO-2a" in files[PAGE] assert publish(rt, lifecycle, changeset_id) == "shadow_recorded" assert not (tmp_path / "state" / "outbox").exists() - calls = (tmp_path / "state" / "traces" / "model_calls.jsonl") + calls = (tmp_path / "state" / "traces" / "records") assert not calls.exists() # scripted gateway bypasses the recorder; real one records diff --git a/test/test_trace_store.py b/test/test_trace_store.py new file mode 100644 index 00000000..0d3f43fd --- /dev/null +++ b/test/test_trace_store.py @@ -0,0 +1,256 @@ +"""trace/1: the shared trace store, the knowledge service's records, replay and export.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from infermatrix_copilot.kb_service.models import ModelGateway, ModelRole +from infermatrix_copilot.kb_service.replay import export_dataset, replay +from infermatrix_copilot.kb_service.runtime import trace_recorder +from infermatrix_copilot.trace_store import REDACTED, SCHEMA, TraceStore, redact, trace_context + +ROOT = Path(__file__).resolve().parents[1] +JUDGE = ModelRole("judge", "codex", "gpt-6-sol", "medium") +CHEAP = ModelRole("judge", "codex", "gpt-6-mini", "low") + + +class Transport: + """A harness transport answering with scripted text (or raising).""" + + def __init__(self, answer): + self.answer = answer + + def complete(self, *, system, messages, model, effort, role): + text = self.answer(model, messages[0]["content"]) + if isinstance(text, Exception): + raise text + return SimpleNamespace(blocks=[SimpleNamespace(text=text)], model=model, stop_reason="end_turn", + usage={"input_tokens": 10, "output_tokens": 5}) + + +def _gateway(store, answer): + return ModelGateway(None, transport_factory=lambda provider: Transport(answer), recorder=trace_recorder(store)) + + +def test_records_are_schema_v1_with_content_addressed_blobs_and_an_index(tmp_path): + store = TraceStore(tmp_path / "traces", environ={}) + with trace_context(run_id="r1", playbook="kb-intake", repo="demo"): + with trace_context(changeset_id="cs1", rule_ids=["DEMO-1a"]): + first = store.append("model_call", inputs={"prompt": "same text"}, outputs={"reply": "{}"}, + model={"role": "judge", "model": "gpt-6-sol"}) + second = store.append("decision", inputs={"prompt": "same text"}, result={"status": "pass"}) + assert first["schema"] == SCHEMA and first["context"]["changeset_id"] == "cs1" + assert "changeset_id" not in second["context"] and second["context"]["repo"] == "demo" + assert first["inputs"]["prompt"] == second["inputs"]["prompt"] # stored once + assert store.blob(first["inputs"]["prompt"]) == "same text" + assert len(list((tmp_path / "traces" / "blobs").rglob("*.gz"))) == 2 + assert [r["id"] for r in store.query(rule_id="DEMO-1a")] == [first["id"]] + assert [r["id"] for r in store.query(run_id="r1", kind="decision")] == [second["id"]] + assert store.get(first["id"]) == first + (tmp_path / "traces" / "index.db").unlink() + assert store.rebuild_index() == 2 and store.get(second["id"]) == second + assert first["env"]["copilot_version"] + + +def test_secrets_never_reach_records_or_blobs(tmp_path): + env = {"GH_TOKEN": "tok-value-1234567890", "KB_SIGNING_KEY": "/path/to/key.pem"} + store = TraceStore(tmp_path / "traces", environ=env) + secrets = ["ghp_" + "a" * 36, "github_pat_" + "b" * 40, "sk-ant-api03-" + "c" * 30, + "AKIA" + "D" * 16, "Bearer " + "e" * 30, "tok-value-1234567890"] + text = "headers: " + " | ".join(secrets) + " key file /path/to/key.pem" + record = store.append("model_call", inputs={"prompt": text}, outputs={"reply": text}, + result={"note": text, secrets[0]: 1}, error=text, context={"note": text}, + model={"served_model": secrets[1]}, usage={secrets[2]: 3}) + raw = b"".join(p.read_bytes() for p in (tmp_path / "traces").rglob("*.jsonl")) + raw += (tmp_path / "traces" / "index.db").read_bytes() + import gzip + raw += b"".join(gzip.decompress(p.read_bytes()) for p in (tmp_path / "traces").rglob("*.gz")) + for secret in secrets: + assert secret.encode() not in raw + assert REDACTED in store.blob(record["inputs"]["prompt"]) + assert "/path/to/key.pem" in store.blob(record["inputs"]["prompt"]) # key PATHS are not secrets + assert redact("nothing here", env) == "nothing here" + as_bytes = store.put_blob(("token " + secrets[0]).encode()) + assert secrets[0] not in store.blob(as_bytes) and REDACTED in store.blob(as_bytes) + with pytest.raises(TypeError): + store.put_blob(b"\xff\xfe binary") + + +def test_the_gateway_records_every_call_including_failures(tmp_path): + store = TraceStore(tmp_path / "traces", environ={}) + gateway = _gateway(store, lambda model, prompt: RuntimeError("quota") if "fail" in prompt else '{"verdict": "consistent"}') + with trace_context(playbook="kb-intake", changeset_id="cs1"): + gateway.call_json(JUDGE, system="sys", prompt="ok") + with pytest.raises(Exception): + gateway.call_json(JUDGE, system="sys", prompt="please fail") + ok, failed = store.query(kind="model_call") + assert ok["model"] == {"role": "judge", "provider": "codex", "model": "gpt-6-sol", "effort": "medium", + "served_model": "gpt-6-sol"} + assert ok["usage"]["input_tokens"] == 10 and ok["context"]["changeset_id"] == "cs1" + assert store.blob(ok["outputs"]["reply"]) == '{"verdict": "consistent"}' + assert "quota" in failed["error"] and store.blob(failed["inputs"]["prompt"]) == "please fail" + + +@pytest.mark.parametrize("reply", ["", "not json", '{"unexpected": 1}']) +def test_unusable_replies_are_recorded_as_failures_and_never_exported(tmp_path, reply): + store = TraceStore(tmp_path / "traces", environ={}) + gateway = _gateway(store, lambda model, prompt: reply if model == "gpt-6-sol" else '{"verdict": "consistent"}') + + def validate(data): + if "verdict" not in data: + raise ValueError("verdict missing") + with trace_context(playbook="kb-intake"), pytest.raises(Exception): + gateway.call_json(JUDGE, system="s", prompt="q", validate=validate) + (call,) = store.query(kind="model_call") + assert call["error"] + assert export_dataset(store, tmp_path / "out.jsonl")["written"] == 0 + result = replay(store, call["id"], gateway, CHEAP) # a failed call can still be replayed + assert result["recorded"] is None and result["replayed"] == {"verdict": "consistent"} + assert store.query(kind="replay")[0]["result"]["replay_of"] == call["id"] + + +def test_replay_asks_a_substitute_model_the_recorded_question(tmp_path): + store = TraceStore(tmp_path / "traces", environ={}) + answers = {"gpt-6-sol": '{"dimensions": {"correct": "yes"}}', "gpt-6-mini": '{"dimensions": {"correct": "no"}}'} + gateway = _gateway(store, lambda model, prompt: answers[model]) + with trace_context(playbook="kb-intake"): + gateway.call_json(JUDGE, system="sys", prompt="is it correct?") + (call,) = store.query(kind="model_call") + result = replay(store, call["id"], gateway, CHEAP) + assert result["recorded"] == {"dimensions": {"correct": "yes"}} + assert result["replayed"] == {"dimensions": {"correct": "no"}} and result["agree"] is False + (record,) = store.query(kind="replay") + assert record["result"]["replay_of"] == call["id"] and record["result"]["replay_model"] == CHEAP.label() + replayed_call = store.query(kind="model_call", model="gpt-6-mini")[0] + assert replayed_call["context"]["playbook"] == "kb-replay" + assert store.blob(replayed_call["inputs"]["prompt"]) == "is it correct?" + answers["gpt-6-mini"] = answers["gpt-6-sol"] + assert replay(store, call["id"], gateway, CHEAP)["agree"] is True + + +def test_export_joins_decisions_and_outcomes_and_drops_calibration_leaks(tmp_path): + store = TraceStore(tmp_path / "traces", environ={}) + calib = tmp_path / "calib" + (calib / "cases").mkdir(parents=True) + (calib / "cases" / "c1.json").write_text(json.dumps({ + "id": "c1", "expected": "reject", "head": {"x.md": "x"}, + "evidence": [{"source_reference": "PR #8107", "title": "Run the perf step via run_benchmark.py"}]})) + gateway = _gateway(store, lambda model, prompt: '{"verdict": "consistent"}') + with trace_context(playbook="kb-intake", changeset_id="cs1", rule_ids=["DEMO-2a"], step="gate"): + gateway.call_json(JUDGE, system="s", prompt=json.dumps({"source_reference": "PR #11"})) + with trace_context(playbook="kb-intake", changeset_id="cs2", step="gate"): + gateway.call_json(JUDGE, system="s", prompt=json.dumps({"source_reference": "PR #8107"}, indent=1)) + with trace_context(playbook="kb-intake", changeset_id="cs3", step="consistency"): + gateway.call_json(JUDGE, system="s", prompt="rule text ... ^[PR #8107]") # cited in rule text + with trace_context(playbook="kb-intake", changeset_id="cs4", step="gate"): + gateway.call_json(JUDGE, system="s", prompt="unrelated PR #81070") + with trace_context(playbook="kb-calibrate"): + gateway.call_json(JUDGE, system="s", prompt="calibration case") + with trace_context(playbook="kb-replay"): + gateway.call_json(JUDGE, system="s", prompt="a replay") + store.append("decision", context={"changeset_id": "cs1"}, result={"status": "pass"}) + store.append("outcome", context={"changeset_id": "cs1"}, result={"outcome": "merged"}) + store.append("outcome", context={"changeset_id": "cs9", "rule_ids": ["DEMO-2a"]}, + result={"outcome": "rule_retired", "rule_id": "DEMO-2a"}) + counts = export_dataset(store, tmp_path / "out.jsonl", calibration_dirs=[calib]) + assert counts == {"written": 2, "leak_dropped": 2, "calibration_dropped": 2, "failed_dropped": 0} + row, other = [json.loads(line) for line in (tmp_path / "out.jsonl").read_text().splitlines()] + assert other["changeset_id"] == "cs4" + assert row["changeset_id"] == "cs1" and row["decision"] == {"status": "pass"} + assert row["outcomes"] == ["merged"] and row["rule_outcomes"] == {"DEMO-2a": ["rule_retired"]} + assert row["output"] == '{"verdict": "consistent"}' and row["task"] == "judge:gate" + + +def test_generator_calls_of_a_real_intake_join_their_change_set(tmp_path): + from test_kb_flow import _flow_runtime + from test_kb_intake_gate import _generator_then_judge + from infermatrix_copilot.kb_service.runtime import collect_events, run_intake + + rt, lifecycle = _flow_runtime(tmp_path) + store = TraceStore(tmp_path / "traces", environ={}) + answer = _generator_then_judge() + rt.traces = store + rt.gateway = ModelGateway(None, recorder=trace_recorder(store), transport_factory=lambda provider: Transport( + lambda model, prompt: json.dumps(answer(SimpleNamespace( + name="generator" if model == rt.generator.model else "judge"), prompt)))) + collect_events(rt, lifecycle) + changeset_id = run_intake(rt, lifecycle) + (generator_call,) = store.query(kind="model_call", role="generator") + assert "changeset_id" not in generator_call["context"] and generator_call["context"]["draft_key"] + assert "_accepted" not in generator_call["context"] and generator_call["context"]["attempt"] == 0 + counts = export_dataset(store, tmp_path / "gen.jsonl", role="generator") + (row,) = [json.loads(line) for line in (tmp_path / "gen.jsonl").read_text().splitlines()] + assert counts["written"] == 1 and row["changeset_id"] == changeset_id + assert row["decision"]["status"] == "pass" and row["rule_ids"] == ["DEMO-2a"] + judged = export_dataset(store, tmp_path / "judge.jsonl") + assert judged["written"] >= 2 + assert all(json.loads(line)["changeset_id"] == changeset_id + for line in (tmp_path / "judge.jsonl").read_text().splitlines()) + + +def test_only_the_accepted_generator_attempt_inherits_the_decision(tmp_path): + """A first reply the knowledge base rejects is repaired; only the repaired + call is labelled with the staged change set's decision.""" + from test_kb_flow import _flow_runtime + from test_kb_intake_gate import PAGE, _generator_then_judge + from infermatrix_copilot.kb_service.runtime import collect_events, run_intake + + rt, lifecycle = _flow_runtime(tmp_path) + store = TraceStore(tmp_path / "traces", environ={}) + answer, calls = _generator_then_judge(), [] + + def reply(model, prompt): + if model == rt.generator.model: + calls.append(prompt) + if len(calls) == 1: # a bad first draft: an operation on a page that does not exist + return json.dumps({"operations": [{"kind": "edit_same_meaning", "page": PAGE, + "rule_id": "NOPE-1a", "section_markdown": "x"}], + "rationale": "r"}) + return json.dumps(answer(SimpleNamespace(name="generator" if model == rt.generator.model else "judge"), + prompt)) + rt.traces = store + rt.gateway = ModelGateway(None, recorder=trace_recorder(store), + transport_factory=lambda provider: Transport(reply)) + collect_events(rt, lifecycle) + changeset_id = run_intake(rt, lifecycle) + assert changeset_id and len(calls) == 2 + export_dataset(store, tmp_path / "gen.jsonl", role="generator") + rows = [json.loads(line) for line in (tmp_path / "gen.jsonl").read_text().splitlines()] + labelled = {row["id"]: row["changeset_id"] for row in rows} + first, second = store.query(kind="model_call", role="generator") + assert labelled == {first["id"]: "", second["id"]: changeset_id} + + +def test_the_sdk_exports_the_trace_store_without_private_modules(): + code = ("import sys; from infermatrix_copilot.sdk.v1 import TraceStore, TRACE_SCHEMA, trace_context, redact;" + "bad = {'infermatrix_copilot.config', 'infermatrix_copilot.contract', 'infermatrix_copilot.mcp_server'};" + "print(sorted(bad & set(sys.modules)), TRACE_SCHEMA)") + out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True, + env={"PYTHONPATH": str(ROOT / "src"), "PATH": "/usr/bin:/bin"}).stdout.strip() + assert out == "[] trace/1" + + +def test_the_knowledge_service_records_decisions_and_outcomes(tmp_path): + from test_kb_flow import _flow_runtime, _open_pr + from infermatrix_copilot.kb_service import merge + + rt, lifecycle = _flow_runtime(tmp_path) + rt.traces = TraceStore(tmp_path / "traces", environ={}) + changeset_id, _publisher = _open_pr(tmp_path, rt, lifecycle) + (decision,) = rt.traces.query(kind="decision") + assert decision["context"]["changeset_id"] == changeset_id + assert decision["context"]["playbook"] == "kb-intake" and decision["context"]["repo"] == lifecycle.repo + assert decision["result"]["status"] == "pass" and decision["result"]["blocks"] + assert decision["context"]["rule_ids"] + rt.ledger.update_changeset(changeset_id, status="pr_open", pr_number=42, head_sha="d" * 40) + rt.github.prs[42] = {"merged": True, "merge_commit_sha": "e" * 40, "state": "closed"} + merge.advance(rt, lifecycle) + (outcome,) = rt.traces.query(kind="outcome", changeset_id=changeset_id) + assert outcome["result"] == {"outcome": "merged", "merge_sha": "e" * 40} + assert outcome["context"]["pr"] == 42