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