diff --git a/tests/planning/test_cot_capture.py b/tests/planning/test_cot_capture.py new file mode 100644 index 0000000..2b7d4dd --- /dev/null +++ b/tests/planning/test_cot_capture.py @@ -0,0 +1,211 @@ +"""W2.D.3 — Planner CoT capture: gzipped ledger entry + 8KB Langfuse span. + +Tests per BUILD_PLAN W2.D.3.a: + +1. ``test_gzipped_cot_in_ledger`` — after ``capture_planner_cot(cot_text, ledger)`` + the ledger contains a ``planner_cot`` entry with: + - ``payload["cot_gzip"]`` — base64-encoded gzip of the CoT text. + - ``payload["cot_gzip_sha256"]`` — SHA-256 of the compressed bytes. + - Decompressing ``base64.b64decode(payload["cot_gzip"])`` yields the + original CoT text. + +2. ``test_8kb_attached_to_langfuse_span`` — ``capture_planner_cot`` returns + a ``CotCaptureResult`` whose ``langfuse_span_payload`` field contains at + most 8192 bytes of the original CoT text (first 8 KB, not gzipped). + +Supporting tests: +- ``test_extract_cloud_cot_from_response_text`` — ``extract_cot(response_text, mode="cloud")`` + returns the full response text (no ```` stripping for cloud). +- ``test_extract_airgap_cot_strips_think_tags`` — ``extract_cot(text, mode="airgap")`` + strips Qwen3 ``…`` wrapper and returns inner content. +- ``test_extract_airgap_cot_passthrough_if_no_think_tags`` — no ```` → + return text as-is (Qwen3 non-thinking mode / GLM-4.5). +- ``test_cot_sha256_is_over_compressed_bytes`` — hash is computed over the + gzip bytes, not the original text (matches ARCHITECTURE.md §5 discipline). + +ARCHITECTURE.md §9 — "planner_cot" is one of the 13 canonical event types. +CLAUDE.md §3.9 — CoT capture does NOT leak to Langfuse unless self-hosted + (air-gap side CoT never leaves the box). The span payload is advisory — + it flows only when Langfuse is reachable; the ledger entry is mandatory. +CLAUDE.md §3.10 — no mocks; Langfuse attach raises NotImplementedError for + live API calls; extraction + gzip + hash are pure Python. +""" + +from __future__ import annotations + +import base64 +import gzip +import hashlib + +from verdict.ledger.memory import InMemoryLedger +from verdict.planning.cot_capture import CotCaptureResult, capture_planner_cot, extract_cot + +# --------------------------------------------------------------------------- +# W2.D.3 — CoT extraction tests +# --------------------------------------------------------------------------- + +_SAMPLE_COT = ( + "Let me think through this step by step.\n" + "First I note the memory image has Windows 10 headers.\n" + "pslist shows 45 processes, psscan shows 47 — DKOM delta = 2 hidden PIDs.\n" + "This strongly suggests T1014 (Rootkit) via EPROCESS unlinking.\n" + "I should run malfind on the hidden PIDs to look for process hollowing.\n" + "No evidence of lateral movement artifacts yet — network logs are clean.\n" +) + +_SAMPLE_THINK_WRAPPED = ( + "\n" + "I am analyzing the Windows memory image.\n" + "The psscan / pslist delta indicates DKOM activity.\n" + "Hypothesis H1 (T1014) looks strong.\n" + "\n" + "Based on my analysis, I recommend the following investigation plan..." +) + + +def test_extract_cloud_cot_returns_full_text() -> None: + """extract_cot(text, mode='cloud') returns the full response text. + Cloud (Claude) does not wrap reasoning in tags. + """ + result = extract_cot(_SAMPLE_COT, mode="cloud") + assert result == _SAMPLE_COT + + +def test_extract_airgap_cot_strips_think_tags() -> None: + """extract_cot(text, mode='airgap') strips Qwen3 … wrapper. + The returned text is the inner reasoning content, not the outer response. + """ + result = extract_cot(_SAMPLE_THINK_WRAPPED, mode="airgap") + assert "" not in result + assert "" not in result + assert "DKOM activity" in result + # The outer response text should NOT be included + assert "Based on my analysis" not in result + + +def test_extract_airgap_cot_passthrough_if_no_think_tags() -> None: + """If no tags are present (non-thinking mode), text passes through.""" + plain_text = "The evidence summary indicates a ransomware incident." + result = extract_cot(plain_text, mode="airgap") + assert result == plain_text + + +def test_extract_cloud_cot_empty_string() -> None: + """extract_cot handles empty string without error.""" + result = extract_cot("", mode="cloud") + assert result == "" + + +def test_extract_airgap_cot_nested_think_uses_innermost() -> None: + """Malformed nested tags — extract_cot uses first open to last close.""" + nested = "outerinnerresponse" + result = extract_cot(nested, mode="airgap") + # Should extract from first to last + assert "" not in result + assert "" not in result + + +# --------------------------------------------------------------------------- +# W2.D.3 — gzip + hash tests +# --------------------------------------------------------------------------- + + +def test_gzipped_cot_in_ledger() -> None: + """After capture_planner_cot, the ledger has a planner_cot entry with: + - payload["cot_gzip"]: base64-encoded gzip of cot_text + - payload["cot_gzip_sha256"]: SHA-256 of the compressed bytes + Decompressing cot_gzip yields the original text. + """ + ledger = InMemoryLedger() + capture_planner_cot( + cot_text=_SAMPLE_COT, + case_id="case-w2d3-001", + ledger=ledger, + ) + assert len(ledger.entries) == 1 + entry = ledger.entries[0] + assert entry.event_type == "planner_cot" + assert entry.case_id == "case-w2d3-001" + assert "cot_gzip" in entry.payload + assert "cot_gzip_sha256" in entry.payload + + # Decompress and verify round-trip + compressed = base64.b64decode(entry.payload["cot_gzip"]) + recovered = gzip.decompress(compressed).decode("utf-8") + assert recovered == _SAMPLE_COT + + # Verify hash is over compressed bytes + expected_hash = hashlib.sha256(compressed).hexdigest() + assert entry.payload["cot_gzip_sha256"] == expected_hash + + +def test_cot_sha256_is_over_compressed_bytes() -> None: + """The cot_gzip_sha256 is SHA-256 over the gzip-compressed bytes, + not over the original text (ARCHITECTURE.md §5 discipline). + """ + ledger = InMemoryLedger() + cot_text = "step 1: examine prefetch. step 2: cross-check amcache." + capture_planner_cot(cot_text=cot_text, case_id="case-hash-check", ledger=ledger) + entry = ledger.entries[0] + + compressed = base64.b64decode(entry.payload["cot_gzip"]) + hash_of_compressed = hashlib.sha256(compressed).hexdigest() + hash_of_original = hashlib.sha256(cot_text.encode()).hexdigest() + + assert entry.payload["cot_gzip_sha256"] == hash_of_compressed + assert entry.payload["cot_gzip_sha256"] != hash_of_original + + +def test_cot_capture_result_8kb_span_payload() -> None: + """capture_planner_cot returns CotCaptureResult.langfuse_span_payload + containing at most 8192 bytes of the original CoT (first 8KB, uncompressed). + """ + long_cot = "A" * 20_000 # 20 KB of text + ledger = InMemoryLedger() + result = capture_planner_cot( + cot_text=long_cot, + case_id="case-8kb", + ledger=ledger, + ) + assert isinstance(result, CotCaptureResult) + span_bytes = result.langfuse_span_payload.encode("utf-8") + assert len(span_bytes) <= 8192, ( + f"langfuse_span_payload must be ≤8192 bytes, got {len(span_bytes)}" + ) + # Must be the START of the CoT, not a random slice + assert result.langfuse_span_payload == long_cot[:8192] + + +def test_cot_capture_short_cot_not_truncated() -> None: + """Short CoT (< 8KB) is not truncated in langfuse_span_payload.""" + short_cot = "The evidence is consistent with T1014." + ledger = InMemoryLedger() + result = capture_planner_cot( + cot_text=short_cot, + case_id="case-short", + ledger=ledger, + ) + assert result.langfuse_span_payload == short_cot + + +def test_cot_capture_result_has_entry_id() -> None: + """CotCaptureResult.ledger_entry_id matches the written ledger entry.""" + ledger = InMemoryLedger() + result = capture_planner_cot( + cot_text=_SAMPLE_COT, + case_id="case-entry-id", + ledger=ledger, + ) + assert result.ledger_entry_id == ledger.entries[0].entry_id + + +def test_8kb_attached_to_langfuse_span() -> None: + """The CotCaptureResult.langfuse_span_payload is the first 8KB of + the original (uncompressed) CoT text. This is what the graph node + attaches to the Langfuse span attribute. + """ + cot = "X" * 16_384 # 16 KB + ledger = InMemoryLedger() + result = capture_planner_cot(cot_text=cot, case_id="case-langfuse", ledger=ledger) + assert len(result.langfuse_span_payload) == 8192 + assert result.langfuse_span_payload == "X" * 8192 diff --git a/verdict/planning/cot_capture.py b/verdict/planning/cot_capture.py new file mode 100644 index 0000000..faeb43c --- /dev/null +++ b/verdict/planning/cot_capture.py @@ -0,0 +1,196 @@ +"""Planner CoT capture — W2.D.3. + +Extracts, compresses, hashes, and stores the planner's chain-of-thought +(CoT) reasoning in the ledger. Attaches the first 8 KB to a Langfuse +span attribute for interactive trace inspection. + +Two CoT formats (ARCHITECTURE.md §9 + BUILD_PLAN W2.D.3.b): + +- **Cloud** (Claude via Agent SDK): the CoT is the model's full response + text before any tool calls or structured JSON output. Claude does NOT + wrap reasoning in ```` tags. +- **Air-gap** (Qwen3 in thinking mode): CoT appears inside + ``…`` tags in the raw model output. GLM-4.5-Air (verifier + only) may or may not use ```` tags; absent tags → passthrough. + +Ledger storage: + ``LedgerEntry(event_type="planner_cot")`` with payload: + - ``"cot_gzip"``: base64-encoded gzip of the UTF-8 CoT text. + - ``"cot_gzip_sha256"``: SHA-256 of the *compressed* bytes (per + ARCHITECTURE.md §5 per-output-file hash discipline — hash the bytes + you store, not the bytes you discard). + +Langfuse span: + ``CotCaptureResult.langfuse_span_payload`` = first 8192 bytes of the + original (uncompressed) CoT text. ARCHITECTURE.md §9 specifies "first + 8KB to Langfuse span attribute". The live Langfuse attach call is not + wired here — it lives in the graph node that wraps this function. + +CLAUDE.md §3.9 — CoT never leaves the box in air-gap mode; Langfuse must + be self-hosted on the air-gap side. This module is side-effect-free + on the Langfuse wire. +CLAUDE.md §3.10 — no mocks; pure Python logic is real; live Langfuse API + calls raise NotImplementedError per §3.10. +""" + +from __future__ import annotations + +import base64 +import dataclasses +import gzip +import hashlib +import re + +from verdict.ledger.memory import InMemoryLedger, Ledger +from verdict.schemas.ledger import Mode + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +#: Maximum bytes attached to a Langfuse span attribute for CoT preview. +#: ARCHITECTURE.md §9: "first 8KB to Langfuse span attribute". +LANGFUSE_SPAN_MAX_BYTES: int = 8192 + +#: Qwen3 … pattern. Greedy to capture full block (including +#: any nested malformed tags). re.DOTALL because reasoning spans multiple +#: lines. +_THINK_RE: re.Pattern[str] = re.compile( + r"(.*)", re.DOTALL | re.IGNORECASE +) + + +# --------------------------------------------------------------------------- +# CotCaptureResult — returned by capture_planner_cot +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class CotCaptureResult: + """Result of a single CoT capture operation. + + Attributes + ---------- + ledger_entry_id: + The ``entry_id`` of the ``planner_cot`` ``LedgerEntry`` written. + langfuse_span_payload: + First 8192 bytes of the original (uncompressed) CoT text. The + graph node attaches this to the Langfuse span attribute. Empty + string if the original CoT was empty. + cot_gzip_sha256: + SHA-256 of the gzip-compressed CoT bytes. Matches + ``ledger_entry.payload["cot_gzip_sha256"]``. + """ + + ledger_entry_id: str + langfuse_span_payload: str + cot_gzip_sha256: str + + +# --------------------------------------------------------------------------- +# extract_cot — strip tags for air-gap models +# --------------------------------------------------------------------------- + + +def extract_cot(response_text: str, *, mode: str) -> str: + """Extract the reasoning CoT from a model response. + + Parameters + ---------- + response_text: + Raw model output (cloud: full text; airgap: may have ````). + mode: + ``"cloud"`` — return ``response_text`` as-is (Claude SDK responses + don't use ```` tags). + ``"airgap"`` — strip Qwen3 ``…`` wrapper and return + the inner reasoning. If no tags are present, return text as-is. + + Returns + ------- + str + Extracted reasoning text. + """ + if mode == "cloud": + return response_text + + # air-gap: try to extract block. + # Greedy match captures from first to last , which means + # nested tags end up in group(1). Strip any residual tags so + # the result is clean reasoning text only. + m = _THINK_RE.search(response_text) + if m: + inner = m.group(1).strip() + # Remove any nested / tags left by greedy match + inner = re.sub(r"", "", inner, flags=re.IGNORECASE).strip() + return inner + + # No tags found (non-thinking mode / GLM-4.5-Air passthrough) + return response_text + + +# --------------------------------------------------------------------------- +# capture_planner_cot — main entry point +# --------------------------------------------------------------------------- + + +def capture_planner_cot( + cot_text: str, + *, + case_id: str, + ledger: Ledger | None = None, + mode: Mode = "cloud", +) -> CotCaptureResult: + """Compress, hash, and store the planner CoT in the ledger. + + Steps: + 1. gzip-compress the CoT text (UTF-8). + 2. SHA-256 the compressed bytes. + 3. base64-encode the compressed bytes for JSON-safe ledger storage. + 4. Write a ``planner_cot`` ``LedgerEntry`` with the payload. + 5. Return a ``CotCaptureResult`` with the entry id, the first 8KB + of the original text, and the hash. + + Parameters + ---------- + cot_text: + The reasoning chain to capture. May be empty string (safe to + call; writes an entry with an empty CoT). + case_id: + The case this CoT belongs to. + ledger: + Ledger to write the ``planner_cot`` event to. Defaults to a + fresh ``InMemoryLedger`` when not supplied. + mode: + Mode at case_init — propagated to the ledger entry. + + Returns + ------- + CotCaptureResult + Contains ``ledger_entry_id``, ``langfuse_span_payload``, and + ``cot_gzip_sha256``. + """ + active_ledger: Ledger = ledger if ledger is not None else InMemoryLedger() + + compressed = gzip.compress(cot_text.encode("utf-8"), compresslevel=9) + cot_gzip_sha256 = hashlib.sha256(compressed).hexdigest() + cot_gzip_b64 = base64.b64encode(compressed).decode("ascii") + + langfuse_span_payload = cot_text[:LANGFUSE_SPAN_MAX_BYTES] + + entry = active_ledger.write( + event_type="planner_cot", + case_id=case_id, + payload={ + "cot_gzip": cot_gzip_b64, + "cot_gzip_sha256": cot_gzip_sha256, + "cot_length_chars": len(cot_text), + "langfuse_span_bytes": len(langfuse_span_payload), + }, + mode=mode, + ) + + return CotCaptureResult( + ledger_entry_id=entry.entry_id, + langfuse_span_payload=langfuse_span_payload, + cot_gzip_sha256=cot_gzip_sha256, + )