diff --git a/docs/week4-cow-radix-extension-plan.md b/docs/week4-cow-radix-extension-plan.md new file mode 100644 index 00000000..112faf39 --- /dev/null +++ b/docs/week4-cow-radix-extension-plan.md @@ -0,0 +1,88 @@ +# COW / Radix Cache Extension Plan (design only) + +Status: **design only — not implemented.** This document preserves the deferred +copy-on-write / radix-cache plan so it can be built as a Week 5 module or a +Week 4 extension after the core Week 4 refsol is reviewed. It intentionally +implements nothing; the core Week 4 stack uses sequential rewind/do-again +(Day 4) and covers every learner scenario without concurrent forks. + +## Why this is deferred + +The core Week 4 thesis is that four planes must agree: durable events/evidence, +model-visible context, derived KV state, and world/approval authority. The +sequential model (checkpoint -> act -> rewind -> do-again) teaches the derived +KV-state plane honestly: reuse the unchanged prefix, recompute only the +divergent suffix. Concurrent COW forks and cross-session prefix sharing are a +*scaling* mechanism, not a correctness one — they answer "how does a production +server serve many sessions/subagents cheaply?" and fit the Week-5-style +optimization arc (like quantization, flash attention, and paged attention did +for earlier weeks). + +## What COW would add (capabilities the sequential model cannot provide) + +1. **Concurrent divergence** — two live branches sharing immutable KV pages + while both continue decoding (parallel what-if exploration, concurrent + subagents). Sequential rewind covers "try A, reverse, try B"; COW covers + "run A and B at the same time." +2. **Cross-session prefix sharing** — a radix/prefix registry that shares one + physical copy of a common prompt prefix across many sessions, with + refcounted immutable pages and copy-on-write tails. + +## Design sketch (for a future implementer) + +### Page model + +- Immutable full pages, content-addressed by exact token block. +- Per-page reference counts; a page is freed only when its last reference + releases it. +- Copy-on-write partial tail: divergence allocates a private page; the shared + prefix pages are never mutated. + +### Registry API (proposed) + +```text +PrefixRegistry(model_hash, tokenizer_hash, block_size, page_budget) + publish(token_ids) -> page_ids # idempotent, content-addressed + acquire(token_ids) -> ForkedPrefix # longest published block-aligned prefix + release(page_ids) # refcounted, fails on double-free + stats() -> PrefixStats # live/shared/private/refs/budget +ForkedPrefix + fork(boundary) -> ForkedPrefix # block-aligned COW child + append(token_ids) -> ForkedPrefix # private divergent tail + close() # releases shared refs exactly once +``` + +### Safety invariants + +- A cache hit requires exact model, tokenizer, and token identity; registry is + model-scoped and fails closed on identity mismatch. +- Pages are immutable after publish; closing a child can never alter another + fork's cache. +- Reference counts cannot leak or double-free; a memory budget evicts only + unreferenced pages. +- Child policies can only narrow parent authority; private child material is + never placed in a shared prefix. +- Branch summaries stay event-tree-level (Day 3) and need no COW. + +### Honest accounting + +Report cold vs optimized wall time separately, exact reused/rewound/prefilled/ +generated/discarded token counts, metadata-copy and KV-page-copy bytes, +live/shared/private page counts and peak KV bytes, and branch-discard costs. +Never call reduced visible bytes or more cache hits a speedup unless end-to-end +latency improves without changing the accepted action or losing evidence. + +### Evaluation + +A read-only fake workspace forks two action plans from one prefix, scores a +deterministic expected result, discards one, and proves no external mutation. +Negative tests attempt a write, change the workspace version, and exhaust the +page budget. Real-model runs compare aggregate time-to-first-token and peak KV +bytes after a warm-up publisher. + +## Source influence + +The design direction is informed by the Week 4 research lanes (Oracle #196, +Tuner #197) and by production radix-cache systems; the core Week 4 stack keeps +the four-plane thesis without this mechanism. No COW behavior is implied in any +core Week 4 feature. diff --git a/docs/week4-day-split.md b/docs/week4-day-split.md index 3797da20..bf06c277 100644 --- a/docs/week4-day-split.md +++ b/docs/week4-day-split.md @@ -1,29 +1,63 @@ # Week 4 Day Split (reference for reviewers) -Status: decided by Forge per Chi's instruction ("7/10/x days your call"); the -refsol stack implements features one PR per day so reviewers can see exactly -what belongs to each day. Curriculum prose (book) is not part of this stack. +Status: decided by Forge per Chi's instruction ("7/10/x days your call"). The +refsol stack is feature-based (one PR per feature so reviewers see exactly +what belongs to each review boundary); the 7-day split groups adjacent +features into teaching days. Curriculum prose (book) is not part of this +stack. ## 7-day structure | Day | Theme | Features (PRs) | Modules | |---|---|---|---| -| 1 | Validated agent loop + tool protocol | feat1 | `protocol.py`, `loop.py` | -| 2 | Effect receipts (durable evidence) | feat2 | `receipts.py`, `workspace.py` wiring | +| 1 | Validated agent loop + tool protocol | feat1 | `protocol.py`, `loop.py`, `generation.py` (minimal) | +| 2 | Authorize effects + durable receipts | feat2 | `workspace.py` (simplified), `receipts.py` | | 3 | Session tree (event-level id/parentId) | feat3 | `session.py` | -| 4 | KV checkpoint/resume + sequential rewind | feat4+5 | `checkpoint.py`, `branch.py`, `generation.py` | -| 5 | Receipt-backed compaction | feat6 | `compaction.py`, `context.py` | -| 6 | Steering + public status + exactly-once reconcile | feat7+8 | `control.py`, `status.py`, `reconcile.py` | -| 7 | Equivalence harness | feat9 | `harness.py`, `evaluation.py` | +| 4 | Derived KV state: checkpoint/resume + sequential rewind | feat4 + feat5 (two PRs) | `checkpoint.py`, `branch.py`, `generation.py` | +| 5 | Receipt-backed compaction | feat6 | `compaction.py` | +| 6 | Control/inspect/recover: steering + status, then exactly-once reconcile | feat7 + feat8 (two PRs) | `control.py`, `status.py`, `reconcile.py` | +| 7 | Equivalence harness | feat9 | `harness.py` | Extension (not a day): COW/radix cache — `docs/week4-cow-radix-extension-plan.md`. -## Stack shape +The old static-held-out grader (`evaluation.py` with `TaskPackage`/`StagedTask`) +is not part of the new course: Day 7's `harness.py` measures the integrated +system (warm/cold, fork/cold, compact/full, crash/resume equivalence) over +the three planes, which replaces the old standalone grader as the evaluation +story. -- PR 1: remove the old 7-day refsol; create the starter skeleton mapped to the - new refsol (this is the map reviewers read). -- PRs 2-8: implement each day's feature(s) in the refsol + focused tests. -- PR 9: COW/radix plan (design only). +## Design note: simplified workspace (Day 2) + +The old 7-day workspace carried a write-ahead mutation journal and undo +machinery (old Day 6 content). The new design drops that machinery: the +workspace keeps the authorization core (bounds, protected paths, observed +digests, approvals, atomic writes) plus effect receipts; crash/effect +recovery is taught by Day 6's exactly-once reconcile instead. This keeps each +day's surface small and matches the "start simple, extend" arc. + +The old summarizer-based `ContextManager` (whole-history compaction with a +model summary) is not part of the new course: Day 5's receipt-backed +compaction replaces it, keeping the durable trace untouched and re-expanding +verified ranges on demand. This removes a large control-coupled module and +keeps Day 5 self-contained. + +## Stack shape (11 PRs) + +1. reset: remove the old 7-day refsol; create the starter skeleton mapped to + the new refsol (this is the map reviewers read). +2. loop + tool protocol (Day 1) +3. effect receipts (Day 2) +4. session tree (Day 3) +5. KV checkpoint (Day 4a) +6. sequential rewind (Day 4b) +7. receipt-backed compaction (Day 5) +8. steering/status (Day 6a) +9. exactly-once reconcile (Day 6b) +10. equivalence harness (Day 7) +11. COW/radix plan (design only, non-day extension) + +Each feature PR is independently reviewable; the day mapping above shows how +adjacent features group into teaching days. ## Why 7 days diff --git a/src/tiny_llm_ref/agent/__init__.py b/src/tiny_llm_ref/agent/__init__.py index c19de7a8..2eb43931 100644 --- a/src/tiny_llm_ref/agent/__init__.py +++ b/src/tiny_llm_ref/agent/__init__.py @@ -1,6 +1,34 @@ # WARNING: Under review - generated by LLM. -from .generation import GenerationStats, generate_response, initial_messages +from .branch import BranchStats, RewindError, SequentialBranch +from .checkpoint import ( + CacheManifest, + ManifestError, + export_cache_manifest, + validate_resume, +) +from .compaction import ( + CompactionError, + CompactionResult, + compact_tool_results, + expand_receipt_range, + reexpand_receipt_message, +) +from .control import AgentInterrupted, CancellationToken, SteeringHandle +from .harness import ( + EquivalenceReport, + PlaneResult, + RunSnapshot, + compare_runs, + snapshot_run, +) +from .generation import ( + GenerationSession, + GenerationStats, + generate_response, + initial_messages, +) +from .session import SessionEvent, SessionLog, SessionStore, memory_session from .loop import AgentEvent, AgentLimits, AgentRun, run_agent from .protocol import ( AgentError, @@ -11,6 +39,16 @@ parse_action, tool_catalog_hash, ) +from .receipts import EffectReceipt, ReceiptStore +from .reconcile import ( + ReconciliationResult, + SafeCheckpoint, + largest_safe_checkpoint, + reconcile_effect, + reconcile_interrupted_effects, +) +from .status import AgentStateCard, StatusQuery, StatusQueryResult, build_state_card +from .workspace import ToolPolicy, Workspace __all__ = [ @@ -18,14 +56,53 @@ "AgentEvent", "AgentLimits", "AgentRun", + "AgentStateCard", + "AgentInterrupted", + "CancellationToken", + "BranchStats", + "EffectReceipt", + "EquivalenceReport", + "CacheManifest", + "CompactionError", + "CompactionResult", "FinalAction", + "GenerationSession", "GenerationStats", + "ManifestError", + "PlaneResult", + "ReceiptStore", + "ReconciliationResult", + "RewindError", + "RunSnapshot", + "SafeCheckpoint", + "SequentialBranch", + "StatusQuery", + "StatusQueryResult", + "SteeringHandle", + "StatusQueryResult", + "SessionEvent", + "SessionLog", + "SessionStore", "TOOL_CATALOG_HASH", "ToolAction", + "ToolPolicy", + "Workspace", + "build_state_card", + "compare_runs", "build_system_prompt", "generate_response", + "export_cache_manifest", + "compact_tool_results", + "expand_receipt_range", "initial_messages", + "largest_safe_checkpoint", + "memory_session", "parse_action", + "reconcile_effect", + "reconcile_interrupted_effects", + "reexpand_receipt_message", "run_agent", + "snapshot_run", + "validate_resume", "tool_catalog_hash", ] diff --git a/src/tiny_llm_ref/agent/branch.py b/src/tiny_llm_ref/agent/branch.py new file mode 100644 index 00000000..72545bac --- /dev/null +++ b/src/tiny_llm_ref/agent/branch.py @@ -0,0 +1,130 @@ +# WARNING: Under review - generated by LLM. + +"""Sequential checkpoint, rewind, and do-again for derived KV state. + +Branching is sequential, not copy-on-write. A learner can checkpoint the +cache at a stable boundary, try one action, rewind to the boundary, and try +another action. The honest KV lesson is unchanged: the unchanged prefix is +reused, only the divergent suffix is recomputed, and a rewind never rolls +back external world state. + +This module deliberately does NOT implement concurrent forks, shared prefix +pages, refcounts, or a radix registry. Those belong to a later extension +(COW/radix cache); the sequential model covers every learner scenario — +"what are you doing?" side queries, what-if exploration, undo/redo — with +``GenerationSession`` longest-prefix reuse plus an explicit rewind. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .generation import GenerationSession + + +class RewindError(RuntimeError): + """A sequential branch operation is invalid or unsafe.""" + + +@dataclass(frozen=True) +class BranchStats: + """Honest token/page accounting for one sequential branch attempt.""" + + checkpoint_tokens: int + reused_tokens: int + rewound_tokens: int + prefilled_tokens: int + output_tokens: int + cold_start: bool = False + + def to_dict(self) -> dict[str, int]: + return { + "checkpoint_tokens": self.checkpoint_tokens, + "reused_tokens": self.reused_tokens, + "rewound_tokens": self.rewound_tokens, + "prefilled_tokens": self.prefilled_tokens, + "output_tokens": self.output_tokens, + "cold_start": int(self.cold_start), + } + + +class SequentialBranch: + """Checkpoint -> act -> rewind -> do again, over one generation session. + + The branch owns a single ``GenerationSession``. ``checkpoint()`` records + the exact cached token prefix at a stable boundary; ``act()`` runs one + generation turn from the current cache state; ``rewind()`` discards the + cache tail back to the boundary (reusing the unchanged prefix); + ``do_again()`` rewinds and runs a different turn from the same boundary. + Closing the branch releases the session. + """ + + def __init__(self, session: "GenerationSession"): + self._session = session + self._checkpoint_tokens: tuple[int, ...] | None = None + self._closed = False + + @property + def session(self) -> "GenerationSession": + return self._session + + @property + def checkpoint_tokens(self) -> tuple[int, ...] | None: + return self._checkpoint_tokens + + def checkpoint(self) -> tuple[int, ...]: + """Record the exact cached token prefix as the branch boundary.""" + + if self._closed: + raise RewindError("cannot checkpoint a closed branch") + tokens = self._session.cached_token_ids + if not tokens: + raise RewindError("cannot checkpoint a cold cache") + self._checkpoint_tokens = tokens + return tokens + + def act(self, messages: list[dict[str, str]]) -> tuple[str, BranchStats]: + """Run one generation turn, then report the honest token accounting.""" + + if self._closed: + raise RewindError("cannot act on a closed branch") + if self._checkpoint_tokens is None: + raise RewindError("checkpoint() must precede act()") + output = self._session(messages) + stats = self._session.last_stats + return output, BranchStats( + checkpoint_tokens=len(self._checkpoint_tokens), + reused_tokens=stats.reused_tokens, + rewound_tokens=stats.rewound_tokens, + prefilled_tokens=stats.prefilled_tokens, + output_tokens=stats.output_tokens, + cold_start=stats.cold_start, + ) + + def rewind(self) -> int: + """Restore the cache exactly to the checkpoint boundary. + + Only the divergent suffix is discarded: the unchanged prefix stays + cached and is reused by the next ``act``. Rewinding the cache never + rewinds the world; callers must reconcile external state separately. + """ + + if self._closed: + raise RewindError("cannot rewind a closed branch") + if self._checkpoint_tokens is None: + raise RewindError("checkpoint() must precede rewind()") + return self._session.rewind_to(len(self._checkpoint_tokens)) + + def do_again(self, messages: list[dict[str, str]]) -> tuple[str, BranchStats]: + """Rewind to the checkpoint and run a different turn.""" + + self.rewind() + return self.act(messages) + + def close(self) -> None: + if self._closed: + return + self._session.close() + self._closed = True diff --git a/src/tiny_llm_ref/agent/checkpoint.py b/src/tiny_llm_ref/agent/checkpoint.py new file mode 100644 index 00000000..ed72107a --- /dev/null +++ b/src/tiny_llm_ref/agent/checkpoint.py @@ -0,0 +1,272 @@ +# WARNING: Under review - generated by LLM. + +"""Content-addressed KV checkpoint export and import. + +A cache checkpoint is a durable manifest that binds the exact token prefix to +the model, tokenizer, layer geometry, and the checkpoint validity rule: the +workspace fingerprint and enabled tool catalog under which it was produced. +Importing a manifest fails closed on any mismatch; a failed or partial import +releases everything and falls back cold. + +The manifest is content-addressed: its digest covers every binding field, so +the checkpoint file itself can be verified end to end. Warm resume reuses +only the exact token prefix; everything after the first changed causal token +must be recomputed. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .generation import GenerationSession + + +class ManifestError(ValueError): + """A checkpoint manifest is invalid, mismatched, or cannot be imported.""" + + +def _json_copy(value: Any) -> Any: + try: + return json.loads( + json.dumps(value, ensure_ascii=True, sort_keys=True, separators=(",", ":")) + ) + except (TypeError, ValueError) as error: + raise ManifestError("manifest payload must be JSON serializable") from error + + +def _canonical(value: Any) -> bytes: + return json.dumps( + value, ensure_ascii=True, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + + +def _valid_hash(value: Any) -> bool: + return ( + isinstance(value, str) + and len(value) == 64 + and all(character in "0123456789abcdef" for character in value) + ) + + +@dataclass(frozen=True) +class CacheManifest: + """Durable, content-addressed description of one KV checkpoint. + + Fields + ------ + model_hash: + Content hash of the model weights/revision that produced the cache. + tokenizer_hash: + Content hash of the tokenizer/chat-template identity. + layers: + Number of layer caches; every layer must agree on the token offset. + position: + Exact token offset represented by the cache at every layer. + prefix_hash: + Content hash of the exact token IDs cached at this checkpoint. + tool_catalog_hash: + Hash of the enabled tool schema catalog at export time. + workspace_fingerprint: + Fingerprint of the workspace-visible state at export time (observed + file digests plus a mutation counter). A changed workspace makes the + checkpoint stale even when the token prefix is identical. + """ + + model_hash: str + tokenizer_hash: str + layers: int + position: int + prefix_hash: str + tool_catalog_hash: str = "" + workspace_fingerprint: str = "" + + def __post_init__(self) -> None: + for name, value in ( + ("model_hash", self.model_hash), + ("tokenizer_hash", self.tokenizer_hash), + ("prefix_hash", self.prefix_hash), + ): + if not _valid_hash(value): + raise ManifestError(f"invalid {name}: must be a sha256 hex digest") + if isinstance(self.layers, bool) or not isinstance(self.layers, int): + raise ManifestError("invalid layer count") + if self.layers <= 0: + raise ManifestError("layer count must be positive") + if isinstance(self.position, bool) or not isinstance(self.position, int): + raise ManifestError("invalid cache position") + if self.position < 0: + raise ManifestError("cache position must be non-negative") + if not isinstance(self.tool_catalog_hash, str): + raise ManifestError("invalid tool catalog hash") + if not isinstance(self.workspace_fingerprint, str): + raise ManifestError("invalid workspace fingerprint") + + def digest(self) -> str: + """Content address of this manifest over every binding field.""" + + return hashlib.sha256(self._payload()).hexdigest() + + def _payload(self) -> bytes: + return _canonical( + { + "model_hash": self.model_hash, + "tokenizer_hash": self.tokenizer_hash, + "layers": self.layers, + "position": self.position, + "prefix_hash": self.prefix_hash, + "tool_catalog_hash": self.tool_catalog_hash, + "workspace_fingerprint": self.workspace_fingerprint, + } + ) + + def to_dict(self) -> dict[str, Any]: + return _json_copy( + { + "digest": self.digest(), + "model_hash": self.model_hash, + "tokenizer_hash": self.tokenizer_hash, + "layers": self.layers, + "position": self.position, + "prefix_hash": self.prefix_hash, + "tool_catalog_hash": self.tool_catalog_hash, + "workspace_fingerprint": self.workspace_fingerprint, + } + ) + + @classmethod + def from_dict(cls, value: Any) -> "CacheManifest": + """Reconstruct and verify a manifest from its durable JSON form.""" + + if not isinstance(value, dict): + raise ManifestError("manifest must be a JSON object") + required = { + "digest", + "model_hash", + "tokenizer_hash", + "layers", + "position", + "prefix_hash", + "tool_catalog_hash", + "workspace_fingerprint", + } + if set(value) != required: + raise ManifestError("invalid manifest fields") + manifest = cls( + model_hash=value["model_hash"], + tokenizer_hash=value["tokenizer_hash"], + layers=value["layers"], + position=value["position"], + prefix_hash=value["prefix_hash"], + tool_catalog_hash=value["tool_catalog_hash"], + workspace_fingerprint=value["workspace_fingerprint"], + ) + if manifest.digest() != value["digest"]: + raise ManifestError("manifest digest does not match its payload") + return manifest + + def require_current( + self, + *, + model_hash: str, + tokenizer_hash: str, + layers: int, + tool_catalog_hash: str, + workspace_fingerprint: str, + ) -> None: + """Fail closed unless this checkpoint is current and identity-safe. + + KV equality never implies world equality. The cache is reusable only + when the exact model, tokenizer, layer geometry, tool catalog, and + workspace fingerprint all still match; any difference forces a cold + fallback or a reconciliation from a safe boundary. + """ + + if self.model_hash != model_hash: + raise ManifestError("checkpoint model does not match the live model") + if self.tokenizer_hash != tokenizer_hash: + raise ManifestError("checkpoint tokenizer does not match the live one") + if self.layers != layers: + raise ManifestError("checkpoint layer geometry does not match") + if self.tool_catalog_hash != tool_catalog_hash: + raise ManifestError("tool catalog changed since this checkpoint") + if self.workspace_fingerprint != workspace_fingerprint: + raise ManifestError("workspace changed since this checkpoint") + + +def export_cache_manifest( + session: "GenerationSession", + token_ids: tuple[int, ...], + *, + tool_catalog_hash: str = "", + workspace_fingerprint: str = "", +) -> CacheManifest: + """Export a durable content-addressed manifest for a live cache state. + + ``token_ids`` is the exact token prefix represented by the session's + layer caches (``session.cached_token_ids``). The manifest binds it to the + session's model/tokenizer identity and the current validity rule inputs. + """ + + if not isinstance(token_ids, tuple) or any( + isinstance(token, bool) or not isinstance(token, int) or token < 0 + for token in token_ids + ): + raise ManifestError("token prefix must be a tuple of non-negative ints") + if session.layer_count == 0: + raise ManifestError("cannot export a manifest from a cold cache") + prefix_hash = hashlib.sha256( + json.dumps(list(token_ids), ensure_ascii=True, separators=(",", ":")).encode( + "utf-8" + ) + ).hexdigest() + return CacheManifest( + model_hash=session.model_identity, + tokenizer_hash=session.tokenizer_identity, + layers=session.layer_count, + position=len(token_ids), + prefix_hash=prefix_hash, + tool_catalog_hash=tool_catalog_hash, + workspace_fingerprint=workspace_fingerprint, + ) + + +def validate_resume( + manifest: CacheManifest, + session: "GenerationSession", + token_ids: tuple[int, ...], + *, + tool_catalog_hash: str, + workspace_fingerprint: str, +) -> None: + """Fail closed before resuming from a manifest. + + This is the import gate: the live session must present the exact token + prefix the manifest describes (same model, tokenizer, layer count, + position, content hash, tool catalog, and workspace fingerprint). Any + mismatch rejects the resume; callers then go cold or reconcile from a + safe boundary. + """ + + manifest.require_current( + model_hash=session.model_identity, + tokenizer_hash=session.tokenizer_identity, + layers=session.layer_count, + tool_catalog_hash=tool_catalog_hash, + workspace_fingerprint=workspace_fingerprint, + ) + if len(token_ids) != manifest.position: + raise ManifestError( + f"checkpoint position {manifest.position} does not match " + f"the presented prefix {len(token_ids)}" + ) + prefix_hash = hashlib.sha256( + json.dumps(list(token_ids), ensure_ascii=True, separators=(",", ":")).encode( + "utf-8" + ) + ).hexdigest() + if prefix_hash != manifest.prefix_hash: + raise ManifestError("checkpoint token prefix does not match its hash") diff --git a/src/tiny_llm_ref/agent/compaction.py b/src/tiny_llm_ref/agent/compaction.py new file mode 100644 index 00000000..0b85ae07 --- /dev/null +++ b/src/tiny_llm_ref/agent/compaction.py @@ -0,0 +1,148 @@ +# WARNING: Under review - generated by LLM. + +"""Receipt-backed model-visible compaction with verified re-expansion. + +Compaction never deletes canonical evidence. A large tool result stays in +the durable receipt store byte-for-byte; only its model-visible rendering is +replaced by a compact receipt handle. Re-expansion is an explicit, verified +observation: the receipt is looked up by its content address, its digest must +match, and the requested byte range is returned from the canonical result. + +The durable event trace is never mutated by compaction. The model sees a +bounded working view; the harness retains evidence and can re-expand any +omitted range on demand. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from .receipts import ReceiptStore + +if TYPE_CHECKING: + from .session import SessionLog + +Message = dict[str, str] + + +class CompactionError(ValueError): + """A compaction or re-expansion operation is invalid or unsafe.""" + + +@dataclass(frozen=True) +class CompactionResult: + """The outcome of one receipt-backed compaction pass.""" + + compacted: int + visible_bytes_before: int + visible_bytes_after: int + messages: tuple[Message, ...] + + @property + def saved_bytes(self) -> int: + return max(0, self.visible_bytes_before - self.visible_bytes_after) + + +def _payload_bytes(value: str) -> int: + return len(value.encode("utf-8")) + + +def compact_tool_results( + session: "SessionLog", + store: ReceiptStore, + *, + max_result_bytes: int, + max_handle_chars: int = 240, +) -> CompactionResult: + """Replace oversized tool-result renderings with verified receipt handles. + + Every ``tool_result`` event whose content exceeds ``max_result_bytes`` is + looked up in the receipt store by its tool-call ID. The model-visible + copy is replaced by the receipt's compact rendering; the canonical bytes + stay in the store. Results without a receipt are left untouched (never + guessed). The durable session events themselves are not modified. + """ + + if isinstance(max_result_bytes, bool) or not isinstance(max_result_bytes, int): + raise CompactionError("max_result_bytes must be an integer") + if max_result_bytes <= 0: + raise CompactionError("max_result_bytes must be positive") + + messages: list[Message] = [] + compacted = 0 + visible_before = 0 + visible_after = 0 + for event in session.events: + if event.type != "tool_result": + continue + content = event.data.get("content") + if not isinstance(content, str): + continue + visible_before += _payload_bytes(content) + if _payload_bytes(content) <= max_result_bytes: + visible_after += _payload_bytes(content) + messages.append({"role": "user", "content": f"Tool result:\n{content}"}) + continue + tool_call_id = event.data.get("tool_call_id") + receipt = ( + store.by_tool_call(tool_call_id) if isinstance(tool_call_id, str) else None + ) + if receipt is None: + visible_after += _payload_bytes(content) + messages.append({"role": "user", "content": f"Tool result:\n{content}"}) + continue + rendering = receipt.compact_rendering(max_handle_chars) + visible_after += _payload_bytes(rendering) + compacted += 1 + messages.append( + {"role": "user", "content": f"Tool result (compacted):\n{rendering}"} + ) + return CompactionResult( + compacted, + visible_before, + visible_after, + tuple(messages), + ) + + +def expand_receipt_range( + store: ReceiptStore, + receipt_id: str, + *, + start: int = 0, + end: int | None = None, +) -> str: + """Verify and re-expand a bounded byte range of one receipt's result. + + Re-expansion fails closed on digest or range mismatch; the canonical + evidence is never altered. + """ + + return store.expand(receipt_id, start=start, end=end) + + +def reexpand_receipt_message( + store: ReceiptStore, + rendering: str, + *, + start: int = 0, + end: int | None = None, +) -> str: + """Re-expand the receipt referenced by one compact rendering. + + Parses the ``expand via receipt `` handle from a compact rendering + produced by ``EffectReceipt.compact_rendering`` and returns the verified + canonical bytes for the requested range. + """ + + marker = "expand via receipt " + index = rendering.rfind(marker) + if index < 0: + raise CompactionError("rendering does not reference a receipt") + receipt_id = rendering[index + len(marker) :].strip() + if len(receipt_id) != 64 or any( + character not in "0123456789abcdef" for character in receipt_id + ): + raise CompactionError("rendering references an invalid receipt ID") + return store.expand(receipt_id, start=start, end=end) diff --git a/src/tiny_llm_ref/agent/control.py b/src/tiny_llm_ref/agent/control.py new file mode 100644 index 00000000..efcde035 --- /dev/null +++ b/src/tiny_llm_ref/agent/control.py @@ -0,0 +1,83 @@ +# WARNING: Under review - generated by LLM. + +from __future__ import annotations + +from collections.abc import Callable +from threading import Lock +from typing import TYPE_CHECKING, TypeVar + +if TYPE_CHECKING: + from .session import SessionEvent, SessionLog + + +_T = TypeVar("_T") + + +class AgentInterrupted(RuntimeError): + """A cooperative stop observed at a named agent boundary.""" + + def __init__(self, reason: str, phase: str): + if not isinstance(reason, str) or not reason.strip(): + raise ValueError("interruption reason must not be blank") + if not isinstance(phase, str) or not phase.strip(): + raise ValueError("interruption phase must not be blank") + self.reason = reason + self.phase = phase + super().__init__(f"agent interrupted during {phase}: {reason}") + + +class CancellationToken: + """A thread-safe, first-writer-wins cooperative cancellation signal.""" + + def __init__(self): + self._lock = Lock() + self._reason: str | None = None + + @property + def cancelled(self) -> bool: + with self._lock: + return self._reason is not None + + @property + def reason(self) -> str | None: + with self._lock: + return self._reason + + def cancel(self, reason: str = "user_interrupt") -> bool: + """Publish the first cancellation reason and report whether it won.""" + + if not isinstance(reason, str) or not reason.strip(): + raise ValueError("cancellation reason must not be blank") + with self._lock: + if self._reason is not None: + return False + self._reason = reason + return True + + def raise_if_cancelled(self, phase: str) -> None: + """Raise with the stable cancellation reason when cancellation won.""" + + with self._lock: + reason = self._reason + if reason is not None: + raise AgentInterrupted(reason, phase) + + def run_if_active(self, phase: str, operation: Callable[[], _T]) -> _T: + """Linearize one terminal operation against cancellation publication.""" + + with self._lock: + if self._reason is not None: + raise AgentInterrupted(self._reason, phase) + return operation() + + +class SteeringHandle: + """Queue operator steering in the durable session transcript.""" + + def __init__(self, session: SessionLog): + self.session = session + + def submit(self, message: str) -> SessionEvent: + """Durably queue one non-blank message for the next safe turn boundary.""" + + return self.session.queue_steering(message) diff --git a/src/tiny_llm_ref/agent/generation.py b/src/tiny_llm_ref/agent/generation.py index fd4a4882..7dd63f8a 100644 --- a/src/tiny_llm_ref/agent/generation.py +++ b/src/tiny_llm_ref/agent/generation.py @@ -1,7 +1,9 @@ # WARNING: Under review - generated by LLM. -from collections.abc import Callable, Sequence +import hashlib +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass +from time import perf_counter from typing import Any @@ -33,6 +35,298 @@ def initial_messages(task: str, system_prompt: str) -> list[Message]: ] +class GenerationSession: + """Generate responses while reusing compatible model KV-cache state.""" + + last_stats: GenerationStats + + def __init__( + self, + model, + tokenizer, + cache_factory: Callable[[], Iterable[Any]], + max_tokens: int, + enable_thinking: bool = False, + cancellation: Any | None = None, + model_hash: str | None = None, + tokenizer_hash: str | None = None, + ): + if max_tokens <= 0: + raise ValueError("max_tokens must be positive") + self._model = model + self._tokenizer = tokenizer + self._cache_factory = cache_factory + self._max_tokens = max_tokens + self._enable_thinking = enable_thinking + self._cancellation = cancellation + self._model_hash = model_hash or self._default_identity(model) + self._tokenizer_hash = tokenizer_hash or self._default_identity(tokenizer) + self._caches: list[Any] = [] + self._cached_token_ids: list[int] = [] + self._closed = False + self.last_stats = GenerationStats() + + @staticmethod + def _default_identity(value: Any) -> str: + """Derive a stable-enough identity hash when the harness does not pin one.""" + + name = type(value).__module__ + "." + type(value).__qualname__ + return hashlib.sha256(name.encode("utf-8")).hexdigest() + + @property + def model_identity(self) -> str: + """Content-addressable model identity used by checkpoint manifests.""" + + return self._model_hash + + @property + def tokenizer_identity(self) -> str: + """Content-addressable tokenizer identity used by checkpoint manifests.""" + + return self._tokenizer_hash + + @property + def layer_count(self) -> int: + """Number of layer caches currently materialized (0 when cold).""" + + return len(self._caches) + + @property + def cached_token_ids(self) -> tuple[int, ...]: + """Return the exact token prefix currently represented by the caches.""" + + return tuple(self._cached_token_ids) + + def encode_messages(self, messages: list[Message]) -> tuple[int, ...]: + """Render semantic messages and return the model's prompt token IDs.""" + + prompt = self._tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=self._enable_thinking, + ) + return tuple( + int(token) + for token in self._tokenizer.encode(prompt, add_special_tokens=False) + ) + + @staticmethod + def _release(caches: Sequence[Any]) -> None: + for cache in caches: + try: + release = getattr(cache, "release", None) + if release is not None: + release() + except BaseException: + # Continue releasing the remaining layers. The discarded cache + # objects are never reused after cleanup starts. + pass + + @staticmethod + def _cache_offset(caches: Sequence[Any]) -> int: + if not caches: + raise ValueError("cache_factory must create at least one layer cache") + offsets: list[int] = [] + for cache in caches: + offset = getattr(cache, "offset", None) + if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0: + raise ValueError( + "every layer cache must expose a non-negative integer offset" + ) + offsets.append(offset) + if any(offset != offsets[0] for offset in offsets[1:]): + raise ValueError("all layer cache offsets must agree") + return offsets[0] + + def _drop_caches(self) -> None: + caches, self._caches = self._caches, [] + self._cached_token_ids = [] + self._release(caches) + + def _create_caches(self) -> None: + caches: list[Any] = [] + try: + created = self._cache_factory() + iterator = iter(created) + while True: + try: + cache = next(iterator) + except StopIteration: + break + caches.append(cache) + if self._cache_offset(caches) != 0: + raise ValueError("new layer caches must start at offset zero") + except BaseException: + self._release(caches) + raise + self._caches = caches + self._cached_token_ids = [] + + def _reset_caches(self) -> None: + self._drop_caches() + self._create_caches() + + def _validated_offset(self, expected: int) -> int: + offset = self._cache_offset(self._caches) + if offset != expected: + raise ValueError( + f"layer cache offset {offset} does not match token prefix {expected}" + ) + return offset + + @staticmethod + def _common_prefix(left: Sequence[int], right: Sequence[int]) -> int: + length = 0 + for left_token, right_token in zip(left, right): + if left_token != right_token: + break + length += 1 + return length + + def _rewind(self, count: int, expected_offset: int) -> None: + if count == 0: + return + # Validate the complete layer set before mutating any cache. A failure + # partway through is handled by discarding the whole set. + self._validated_offset(expected_offset) + for cache in self._caches: + rewind = getattr(cache, "rewind", None) + if rewind is None: + raise ValueError("every layer cache must support rewind") + rewind(count) + self._validated_offset(expected_offset - count) + + def __call__(self, messages: list[Message]) -> str: + """Render and decode one response, reusing the safe token prefix.""" + + if self._closed: + raise RuntimeError("generation session is closed") + if self._cancellation is not None: + self._cancellation.raise_if_cancelled("model_request") + + import mlx.core as mx + + started = perf_counter() + desired_tokens = self.encode_messages(messages) + if not desired_tokens: + raise ValueError("rendered prompt must contain at least one token") + + cold_start = not self._caches + rewound_tokens = 0 + if cold_start: + self._create_caches() + common_prefix = 0 + else: + try: + cached_length = self._validated_offset(len(self._cached_token_ids)) + common_prefix = self._common_prefix( + self._cached_token_ids, desired_tokens + ) + rewind_count = cached_length - common_prefix + self._rewind(rewind_count, cached_length) + rewound_tokens = rewind_count + except Exception: + # Offset disagreement, unsupported rewind, or a partially + # failed rewind makes every layer unsafe to reuse. + self._reset_caches() + cold_start = True + common_prefix = 0 + rewound_tokens = 0 + except BaseException: + self._drop_caches() + raise + + # A model call needs at least one token to produce next-token logits. + # If the desired prompt is already represented, replay its final token. + if common_prefix == len(desired_tokens): + try: + self._rewind(1, common_prefix) + common_prefix -= 1 + rewound_tokens += 1 + except Exception: + self._reset_caches() + cold_start = True + common_prefix = 0 + rewound_tokens = 0 + except BaseException: + self._drop_caches() + raise + + reused_tokens = common_prefix + prompt_suffix = list(desired_tokens[common_prefix:]) + represented_tokens = list(desired_tokens[:common_prefix]) + output: list[int] = [] + input_ids = prompt_suffix + offset = common_prefix + + try: + for _ in range(self._max_tokens): + if self._cancellation is not None: + self._cancellation.raise_if_cancelled("model_decode") + tokens = mx.array(input_ids) + logits = self._model(tokens[None], offset, self._caches)[:, -1, :] + offset += len(input_ids) + self._validated_offset(offset) + represented_tokens.extend(input_ids) + + token = int(mx.argmax(logits, axis=-1).item()) + if token == self._tokenizer.eos_token_id: + break + output.append(token) + input_ids = [token] + if self._cancellation is not None: + self._cancellation.raise_if_cancelled("model_decode") + except BaseException: + self._drop_caches() + raise + + self._cached_token_ids = represented_tokens + self.last_stats = GenerationStats( + input_tokens=len(desired_tokens), + reused_tokens=reused_tokens, + rewound_tokens=rewound_tokens, + prefilled_tokens=len(prompt_suffix), + output_tokens=len(output), + cold_start=cold_start, + latency_seconds=perf_counter() - started, + ) + return self._tokenizer.decode(output) + + def rewind_to(self, token_count: int) -> int: + """Discard the cache tail so exactly ``token_count`` tokens remain. + + The unchanged prefix stays cached and is reused by the next call; + only the divergent suffix is dropped. Rewinding the cache never + rewinds the world: external effects must be reconciled separately. + """ + + if self._closed: + raise RuntimeError("generation session is closed") + if isinstance(token_count, bool) or not isinstance(token_count, int): + raise ValueError("token_count must be an integer") + if token_count < 0: + raise ValueError("token_count must be non-negative") + if not self._caches: + raise ValueError("cannot rewind a cold cache") + current = self._validated_offset(len(self._cached_token_ids)) + if token_count > current: + raise ValueError("cannot rewind past the cached prefix") + if token_count == current: + return current + self._rewind(current - token_count, current) + self._cached_token_ids = list(self._cached_token_ids[:token_count]) + return token_count + + def close(self) -> None: + """Release every layer cache. Closing more than once is safe.""" + + if self._closed: + return + self._closed = True + self._drop_caches() + + def generate_response( model, tokenizer, diff --git a/src/tiny_llm_ref/agent/harness.py b/src/tiny_llm_ref/agent/harness.py new file mode 100644 index 00000000..6af6d293 --- /dev/null +++ b/src/tiny_llm_ref/agent/harness.py @@ -0,0 +1,188 @@ +# WARNING: Under review - generated by LLM. + +"""End-to-end semantic/evidence/policy equivalence harness. + +Token, page, and latency gains count only when the optimized run preserves +the accepted outcome, the retained evidence, and the policy decisions of the +baseline run. This harness compares pairs of runs — warm/cold, fork/cold, +compact/full, crash/resume — and reports a structured verdict over three +planes: + +- semantic: the final action/outcome and accepted result; +- evidence: durable receipts and tool observations retained; +- policy: approvals, epochs, and effect decisions agree. + +Any mismatch fails the pair; the harness never reports a speedup as correct +when the behavior differs. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .receipts import EffectReceipt + from .session import SessionLog + + +@dataclass(frozen=True) +class RunSnapshot: + """The observable contract of one run, independent of its mechanism.""" + + final: str | None + reason: str + completed: bool + tool_actions: tuple[tuple[str, str], ...] + receipts: tuple[str, ...] + observations: tuple[str, ...] + approval_state: str = "unknown" + world_state: str = "unknown" + token_accounting: dict[str, int] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PlaneResult: + """PASS/FAIL verdict for one equivalence plane.""" + + plane: str + ok: bool + detail: str = "" + + +@dataclass(frozen=True) +class EquivalenceReport: + """The full verdict of comparing a baseline run with an optimized run.""" + + baseline: RunSnapshot + optimized: RunSnapshot + planes: tuple[PlaneResult, ...] + + @property + def ok(self) -> bool: + return all(plane.ok for plane in self.planes) + + def render(self) -> str: + lines = ["semantic/evidence/policy equivalence: PASS" if self.ok else "FAIL"] + for plane in self.planes: + status = "PASS" if plane.ok else "FAIL" + lines.append(f"- {plane.plane}: {status} ({plane.detail})") + return "\n".join(lines) + + +def _receipt_ids(receipts: tuple[EffectReceipt, ...]) -> tuple[str, ...]: + return tuple(sorted(receipt.receipt_id for receipt in receipts)) + + +def snapshot_run( + *, + final: str | None, + reason: str, + completed: bool, + session: "SessionLog | None" = None, + receipts: tuple[EffectReceipt, ...] = (), + approval_state: str = "unknown", + world_state: str = "unknown", + token_accounting: dict[str, int] | None = None, +) -> RunSnapshot: + """Capture the observable contract of one run. + + ``session`` supplies the durable event evidence; ``receipts`` supplies the + canonical effect evidence. Everything the harness compares is public and + mechanism-independent. + """ + + tool_actions: list[tuple[str, str]] = [] + observations: list[str] = [] + if session is not None: + for event in session.events: + if event.type == "tool_call": + tool = event.data.get("tool", "?") + arguments = json.dumps(event.data.get("arguments", {}), sort_keys=True) + tool_actions.append((tool, arguments)) + elif event.type == "tool_result": + content = event.data.get("content", "") + if isinstance(content, str): + observations.append(content) + return RunSnapshot( + final=final, + reason=reason, + completed=completed, + tool_actions=tuple(tool_actions), + receipts=_receipt_ids(receipts), + observations=tuple(observations), + approval_state=approval_state, + world_state=world_state, + token_accounting=dict(token_accounting or {}), + ) + + +def compare_runs( + baseline: RunSnapshot, + optimized: RunSnapshot, +) -> EquivalenceReport: + """Compare the observable contracts of a baseline and an optimized run. + + Planes compared: + - semantic: completed/final/reason and the exact tool-action sequence; + - evidence: durable receipt IDs and tool observations; + - policy: approval and world state at the end of each run. + + The token/page/latency accounting in each snapshot is reported for + context but never used to pass a run whose behavior differs. + """ + + semantic_ok = ( + baseline.completed == optimized.completed + and baseline.final == optimized.final + and baseline.reason == optimized.reason + and baseline.tool_actions == optimized.tool_actions + ) + semantic = PlaneResult( + "semantic", + semantic_ok, + detail=( + "final/actions match" + if semantic_ok + else ( + f"baseline final={baseline.final!r} actions={baseline.tool_actions}; " + f"optimized final={optimized.final!r} actions={optimized.tool_actions}" + ) + ), + ) + evidence_ok = ( + baseline.receipts == optimized.receipts + and baseline.observations == optimized.observations + ) + evidence = PlaneResult( + "evidence", + evidence_ok, + detail=( + "receipts and observations match" + if evidence_ok + else ( + f"baseline receipts={baseline.receipts} observations=" + f"{baseline.observations}; optimized receipts={optimized.receipts} " + f"observations={optimized.observations}" + ) + ), + ) + policy_ok = ( + baseline.approval_state == optimized.approval_state + and baseline.world_state == optimized.world_state + ) + policy = PlaneResult( + "policy", + policy_ok, + detail=( + "state matches" + if policy_ok + else ( + f"baseline approval={baseline.approval_state} " + f"world={baseline.world_state}; optimized " + f"approval={optimized.approval_state} world={optimized.world_state}" + ) + ), + ) + return EquivalenceReport(baseline, optimized, (semantic, evidence, policy)) diff --git a/src/tiny_llm_ref/agent/receipts.py b/src/tiny_llm_ref/agent/receipts.py new file mode 100644 index 00000000..670e435d --- /dev/null +++ b/src/tiny_llm_ref/agent/receipts.py @@ -0,0 +1,303 @@ +# WARNING: Under review - generated by LLM. + +"""Immutable effect receipts and their durable store. + +An ``EffectReceipt`` is the canonical evidence record of one executed tool +effect: the validated action, its normalized inputs, the exit/result state, +and the changed artifacts. A receipt is immutable and content-addressed: its +ID is the SHA-256 of its canonical payload, so a later compaction, replay, or +audit can verify that the bytes it points at have not been altered. + +The receipt store is append-only and durable. Compaction never deletes +canonical evidence; it only replaces the model-visible rendering. Re-expansion +fails closed on digest mismatch, and unmatched call/result pairs are never +split. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +def _json_copy(value: Any) -> Any: + """Validate and detach a JSON-compatible receipt payload.""" + + try: + return json.loads( + json.dumps(value, ensure_ascii=True, sort_keys=True, separators=(",", ":")) + ) + except (TypeError, ValueError) as error: + raise ValueError("receipt payload must be JSON serializable") from error + + +def _canonical_payload(value: Any) -> bytes: + """Serialize a receipt payload into stable canonical bytes.""" + + return json.dumps( + value, ensure_ascii=True, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + + +@dataclass(frozen=True) +class EffectReceipt: + """Immutable evidence of one executed tool effect.""" + + tool_call_id: str + tool: str + arguments: dict[str, Any] + exit_state: str # "ok" | "error" | "uncertain" + result: str + changed_artifacts: tuple[str, ...] + # Optional bounded head/tail rendering used by the compaction feature. + head: str = "" + tail: str = "" + + def __post_init__(self) -> None: + if not isinstance(self.tool_call_id, str) or not self.tool_call_id: + raise ValueError("tool call ID must be a non-empty string") + if not isinstance(self.tool, str) or not self.tool: + raise ValueError("tool name must be a non-empty string") + if self.exit_state not in {"ok", "error", "uncertain"}: + raise ValueError("exit state must be ok, error, or uncertain") + if not isinstance(self.result, str): + raise ValueError("result must be a string") + if not isinstance(self.changed_artifacts, tuple) or not all( + isinstance(path, str) and path for path in self.changed_artifacts + ): + raise ValueError("changed artifacts must be a tuple of paths") + object.__setattr__(self, "arguments", _json_copy(self.arguments)) + object.__setattr__( + self, "changed_artifacts", tuple(sorted(self.changed_artifacts)) + ) + if not isinstance(self.head, str) or not isinstance(self.tail, str): + raise ValueError("head and tail must be strings") + + @property + def receipt_id(self) -> str: + """Content address: SHA-256 of the canonical payload.""" + + return hashlib.sha256(self._payload()).hexdigest() + + def _payload(self) -> bytes: + return _canonical_payload( + { + "tool_call_id": self.tool_call_id, + "tool": self.tool, + "arguments": self.arguments, + "exit_state": self.exit_state, + "result": self.result, + "changed_artifacts": self.changed_artifacts, + "head": self.head, + "tail": self.tail, + } + ) + + def to_dict(self) -> dict[str, Any]: + return _json_copy( + { + "receipt_id": self.receipt_id, + "tool_call_id": self.tool_call_id, + "tool": self.tool, + "arguments": self.arguments, + "exit_state": self.exit_state, + "result": self.result, + "changed_artifacts": self.changed_artifacts, + "head": self.head, + "tail": self.tail, + } + ) + + @classmethod + def from_dict(cls, value: Any) -> "EffectReceipt": + """Reconstruct and verify a receipt from its durable JSON form.""" + + if not isinstance(value, dict): + raise ValueError("receipt must be a JSON object") + required = { + "receipt_id", + "tool_call_id", + "tool", + "arguments", + "exit_state", + "result", + "changed_artifacts", + "head", + "tail", + } + if set(value) != required: + raise ValueError("invalid receipt fields") + receipt = cls( + tool_call_id=value["tool_call_id"], + tool=value["tool"], + arguments=value["arguments"], + exit_state=value["exit_state"], + result=value["result"], + changed_artifacts=tuple(value["changed_artifacts"]), + head=value["head"], + tail=value["tail"], + ) + if receipt.receipt_id != value["receipt_id"]: + raise ValueError("receipt digest does not match its payload") + return receipt + + def compact_rendering(self, max_chars: int = 240) -> str: + """Build the bounded model-visible rendering of this receipt. + + Compaction never deletes canonical evidence: the receipt keeps the + full result and changed artifacts, while the model sees a compact + handle with the receipt ID, tool, exit state, and bounded head/tail. + Re-expansion is an explicit verified observation. + """ + + if isinstance(max_chars, bool) or not isinstance(max_chars, int): + raise ValueError("max_chars must be an integer") + if max_chars < 80: + raise ValueError("max_chars must leave room for the receipt handle") + head = self.head or self.result[:80] + tail = self.tail or (self.result[-80:] if len(self.result) > 160 else "") + handle = f"expand via receipt {self.receipt_id}" + header = ( + f"[tool result {self.tool_call_id[:8]} " + f"{self.exit_state} sha256:{self.receipt_id[:16]} " + f"bytes:{len(self.result.encode('utf-8'))}]" + ) + head_line = f"head: {head[:40]!r}" + tail_line = f"tail: {tail[:40]!r}" + rendered = f"{header}\n{head_line}\n{tail_line}\n{handle}" + encoded = rendered.encode("utf-8") + if len(encoded) <= max_chars: + return rendered + # The receipt handle must stay whole so re-expansion can verify the + # exact content address; only the head/tail previews are shortened. + fixed = f"{header}\n{handle}" + budget = max_chars - len(fixed.encode("utf-8")) + if budget < 0: + raise ValueError("max_chars is too small to keep the receipt handle") + preview_budget = budget // 2 + head_preview = head[: max(0, preview_budget - 7)] + tail_preview = tail[: max(0, preview_budget - 7)] + return f"{header}\nhead: {head_preview!r}\ntail: {tail_preview!r}\n{handle}" + + +class ReceiptStore: + """Append-only durable store of immutable effect receipts. + + Each receipt is persisted as one canonical JSON line with an fsync at the + store level. The store never mutates or deletes a published receipt; + ``put`` is idempotent by content address, and ``get`` verifies the digest. + """ + + def __init__(self, path: Path | None = None): + self._path = path + self._receipts: dict[str, EffectReceipt] = {} + self._by_tool_call: dict[str, str] = {} + self._order: list[str] = [] + if path is not None: + self._load(path) + + def _load(self, path: Path) -> None: + path = path.resolve() + if path.exists(): + with path.open("r", encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if not line: + continue + receipt = EffectReceipt.from_dict(json.loads(line)) + if receipt.receipt_id not in self._receipts: + self._receipts[receipt.receipt_id] = receipt + self._order.append(receipt.receipt_id) + self._by_tool_call.setdefault( + receipt.tool_call_id, receipt.receipt_id + ) + + def put(self, receipt: EffectReceipt) -> str: + """Durably publish one receipt; idempotent by content address.""" + + receipt_id = receipt.receipt_id + if receipt_id in self._receipts: + return receipt_id + if self._path is not None: + self._path = self._path.resolve() + self._path.parent.mkdir(parents=True, exist_ok=True) + line = ( + json.dumps( + receipt.to_dict(), + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ) + + "\n" + ) + with self._path.open("a", encoding="utf-8") as handle: + handle.write(line) + handle.flush() + os.fsync(handle.fileno()) + self._receipts[receipt_id] = receipt + self._order.append(receipt_id) + self._by_tool_call.setdefault(receipt.tool_call_id, receipt_id) + return receipt_id + + def get(self, receipt_id: str) -> EffectReceipt | None: + """Return the verified receipt or None when absent.""" + + return self._receipts.get(receipt_id) + + def require(self, receipt_id: str) -> EffectReceipt: + receipt = self.get(receipt_id) + if receipt is None: + raise KeyError(f"receipt not found: {receipt_id}") + return receipt + + def by_tool_call(self, tool_call_id: str) -> EffectReceipt | None: + """Return the receipt recorded for one tool-call event, if any. + + Receipts are content-addressed, so this is an explicit secondary index + keyed by the session tool-call ID — the link an exactly-once + reconciler needs after a crash. + """ + + receipt_id = self._by_tool_call.get(tool_call_id) + if receipt_id is None: + return None + return self._receipts.get(receipt_id) + + def expand(self, receipt_id: str, *, start: int = 0, end: int | None = None) -> str: + """Verify and re-expand a bounded byte range of one receipt's result. + + Re-expansion is an explicit, verified observation: the stored receipt + is reconstructed from its canonical payload, its content address must + match, and the requested range must be within the result. Any digest + or range mismatch fails closed instead of returning partial bytes. + """ + + receipt = self.require(receipt_id) + result = receipt.result + if isinstance(start, bool) or not isinstance(start, int) or start < 0: + raise ValueError("start must be a non-negative integer") + if end is None: + end = len(result) + if isinstance(end, bool) or not isinstance(end, int): + raise ValueError("end must be an integer") + if end < start or end > len(result): + raise ValueError("receipt byte range is outside the stored result") + return result[start:end] + + def __iter__(self): + for receipt_id in self._order: + yield self._receipts[receipt_id] + + def __len__(self) -> int: + return len(self._order) + + def close(self) -> None: + """Drop the in-memory index; the durable file is already fsynced.""" + + self._receipts.clear() + self._by_tool_call.clear() + self._order = [] diff --git a/src/tiny_llm_ref/agent/reconcile.py b/src/tiny_llm_ref/agent/reconcile.py new file mode 100644 index 00000000..7f6a436f --- /dev/null +++ b/src/tiny_llm_ref/agent/reconcile.py @@ -0,0 +1,170 @@ +# WARNING: Under review - generated by LLM. + +"""Exactly-once crash/effect reconciliation and safe-checkpoint recovery. + +A process can crash after a modifying tool commits but before its observation +is appended. On restart the harness must reconcile the effect receipt, append +exactly one observation, and resume from the largest safe checkpoint — never +guessing, never blindly re-running the effect just to reconstruct model state. + +Rules: + +- every effect ID reaches one terminal observation; +- an orphaned intent is reconciled, never guessed or blindly rerun; +- replay is deterministic and side-effect-free; effectful tools are never + re-executed to rebuild context; +- ``largest_safe_checkpoint`` returns the most advanced manifest that still + passes fail-closed identity checks, or a cold fallback reason. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from .checkpoint import CacheManifest, ManifestError +from .receipts import ReceiptStore + +if TYPE_CHECKING: + from .session import SessionLog + + +@dataclass(frozen=True) +class ReconciliationResult: + """Outcome of reconciling one interrupted tool call.""" + + tool_call_id: str + status: str # "observation_appended" | "no_receipt" | "already_closed" + receipt_id: str | None = None + content: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "tool_call_id": self.tool_call_id, + "status": self.status, + "receipt_id": self.receipt_id, + "content": self.content, + } + + +def reconcile_effect( + session: "SessionLog", + store: ReceiptStore, + tool_call_id: str, +) -> ReconciliationResult: + """Close one interrupted effect exactly once. + + If the session already has a ``tool_result`` for the call, it is left + untouched (exactly-once). Otherwise the durable receipt is looked up by + tool-call ID and its verified result is appended as the single terminal + observation. A call without a receipt is reported as ``no_receipt`` and + never re-executed. + """ + + events = session.events + if any( + event.type == "tool_result" and event.data.get("tool_call_id") == tool_call_id + for event in events + ): + return ReconciliationResult(tool_call_id, "already_closed") + + receipt = store.by_tool_call(tool_call_id) + if receipt is None: + return ReconciliationResult(tool_call_id, "no_receipt") + + session.append( + "tool_result", + tool_call_id=tool_call_id, + tool=receipt.tool, + arguments=receipt.arguments, + is_error=receipt.exit_state == "error", + content=receipt.result, + reconciled=True, + receipt_id=receipt.receipt_id, + ) + return ReconciliationResult( + tool_call_id, + "observation_appended", + receipt_id=receipt.receipt_id, + content=receipt.result, + ) + + +def reconcile_interrupted_effects( + session: "SessionLog", + store: ReceiptStore, +) -> tuple[ReconciliationResult, ...]: + """Reconcile every unclosed tool call in one pass, in event order. + + Deterministic and side-effect-free: only the durable receipt store is + consulted; no tool is re-executed and no receipt is guessed. + """ + + events = session.events + closed = { + event.data.get("tool_call_id") + for event in events + if event.type == "tool_result" and event.data.get("tool_call_id") is not None + } + results = [] + for event in events: + if event.type != "tool_call": + continue + if event.id in closed: + continue + results.append(reconcile_effect(session, store, event.id)) + closed.add(event.id) + return tuple(results) + + +@dataclass(frozen=True) +class SafeCheckpoint: + """The largest manifest that may be resumed, or a cold fallback reason.""" + + manifest: CacheManifest | None + reason: str # "resume" | "cold_no_checkpoint" | "cold_" + reason + + @property + def can_resume(self) -> bool: + return self.manifest is not None + + +def largest_safe_checkpoint( + manifests: tuple[CacheManifest, ...], + *, + model_hash: str, + tokenizer_hash: str, + layers: int, + tool_catalog_hash: str, + workspace_fingerprint: str, +) -> SafeCheckpoint: + """Return the largest checkpoint that passes every fail-closed check. + + ``largest`` means the highest ``position``; ties resolve to the first in + input order. Any manifest that mismatches the live model, tokenizer, + layer geometry, tool catalog, or workspace fingerprint is rejected, and if + none survives the result is cold with the first rejection reason recorded. + """ + + if not manifests: + return SafeCheckpoint(None, "cold_no_checkpoint") + valid: list[CacheManifest] = [] + first_reason = "cold_unknown" + for manifest in manifests: + try: + manifest.require_current( + model_hash=model_hash, + tokenizer_hash=tokenizer_hash, + layers=layers, + tool_catalog_hash=tool_catalog_hash, + workspace_fingerprint=workspace_fingerprint, + ) + except ManifestError as error: + if first_reason == "cold_unknown": + first_reason = f"cold_{str(error).replace(' ', '_')}" + continue + valid.append(manifest) + if not valid: + return SafeCheckpoint(None, first_reason) + best = max(valid, key=lambda manifest: (manifest.position,)) + return SafeCheckpoint(best, "resume") diff --git a/src/tiny_llm_ref/agent/session.py b/src/tiny_llm_ref/agent/session.py new file mode 100644 index 00000000..367c3524 --- /dev/null +++ b/src/tiny_llm_ref/agent/session.py @@ -0,0 +1,1340 @@ +# WARNING: Under review - generated by LLM. + +import fcntl +import hashlib +import json +import os +import re +import stat +import uuid +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from threading import Lock +from typing import Any, TypeVar + +from .generation import Message + + +_SESSION_ID = re.compile(r"[a-f0-9]{32}") +_EVENT_TYPE = re.compile(r"[a-z][a-z0-9_]{0,63}") +_MAX_EVENT_BYTES = 1024 * 1024 +_MAX_LOG_BYTES = 16 * 1024 * 1024 +_MAX_EVENTS = 100_000 +_MAX_INSTRUCTIONS_BYTES = 64 * 1024 +_T = TypeVar("_T") + + +def _json_copy(value: Any) -> Any: + """Validate and detach the JSON value used by a durable event.""" + + try: + return json.loads( + json.dumps(value, ensure_ascii=True, sort_keys=True, separators=(",", ":")) + ) + except (TypeError, ValueError) as error: + raise ValueError("session event data must be JSON serializable") from error + + +def _read_regular_file( + path: Path, + limit: int, + label: str, + *, + dir_fd: int | None = None, + private: bool = False, +) -> tuple[bytes, tuple[int, int]]: + """Read a bounded regular file without following a final-component symlink.""" + + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + descriptor = os.open(path, flags, dir_fd=dir_fd) + except OSError as error: + raise ValueError(f"{label} does not exist or is unsafe") from error + try: + status = os.fstat(descriptor) + if ( + not stat.S_ISREG(status.st_mode) + or status.st_nlink != 1 + or status.st_size > limit + or (private and stat.S_IMODE(status.st_mode) & 0o077 != 0) + ): + raise ValueError(f"{label} is invalid or too large") + with os.fdopen(descriptor, "rb", closefd=False) as file: + content = file.read(limit + 1) + if len(content) > limit: + raise ValueError(f"{label} is too large") + return content, (status.st_dev, status.st_ino) + finally: + os.close(descriptor) + + +@dataclass(frozen=True) +class SessionEvent: + """One immutable event in an append-only agent transcript.""" + + id: str + timestamp: str + type: str + parent_id: str | None + data: dict[str, Any] + + def to_dict(self) -> dict[str, Any]: + return _json_copy( + { + "id": self.id, + "timestamp": self.timestamp, + "type": self.type, + "parent_id": self.parent_id, + "data": self.data, + } + ) + + @classmethod + def from_dict(cls, value: Any) -> "SessionEvent": + if not isinstance(value, dict) or set(value) != { + "id", + "timestamp", + "type", + "parent_id", + "data", + }: + raise ValueError("invalid session event fields") + event_id = value["id"] + event_type = value["type"] + parent_id = value["parent_id"] + if not isinstance(event_id, str) or not _SESSION_ID.fullmatch(event_id): + raise ValueError("invalid session event ID") + if not isinstance(event_type, str) or not _EVENT_TYPE.fullmatch(event_type): + raise ValueError("invalid session event type") + if parent_id is not None and ( + not isinstance(parent_id, str) or not _SESSION_ID.fullmatch(parent_id) + ): + raise ValueError("invalid parent event ID") + if not isinstance(value["timestamp"], str) or not value["timestamp"]: + raise ValueError("invalid session event timestamp") + if not isinstance(value["data"], dict): + raise ValueError("invalid session event data") + return cls( + event_id, + value["timestamp"], + event_type, + parent_id, + _json_copy(value["data"]), + ) + + +class SessionLog: + """A memory-backed or fsync-backed append-only session event stream.""" + + def __init__( + self, + session_id: str, + workspace: Path, + model: str, + *, + path: Path | None = None, + events: tuple[SessionEvent, ...] = (), + identity: tuple[int, int] | None = None, + persisted_size: int | None = None, + ): + if not _SESSION_ID.fullmatch(session_id): + raise ValueError("invalid session ID") + self.session_id = session_id + self.workspace = workspace.resolve() + self.model = model + self.path = path + self._identity = identity + self._events = [SessionEvent.from_dict(event.to_dict()) for event in events] + self._persisted_size = ( + None if path is None else (0 if persisted_size is None else persisted_size) + ) + self._lock = Lock() + self._steering_delivery_lock = Lock() + self._write_failed = False + + @property + def events(self) -> tuple[SessionEvent, ...]: + with self._lock: + return tuple( + SessionEvent.from_dict(event.to_dict()) for event in self._events + ) + + @property + def instructions(self) -> tuple[dict[str, str], ...]: + """Return the newest recorded project-instruction snapshot.""" + + snapshot: Any = () + for event in self.events: + if event.type == "session_started": + snapshot = event.data.get("instructions", ()) + elif event.type == "instructions_changed": + snapshot = event.data.get("current", ()) + if not isinstance(snapshot, (list, tuple)): + return () + return tuple(item for item in snapshot if isinstance(item, dict)) + + def active_path_events(self) -> tuple[SessionEvent, ...]: + """Deterministically reconstruct the model-visible active path. + + The append-only ``id``/``parentId`` tree keeps every branch's full + history, so the active path is exactly this session's own event + sequence: events after ``branch_completed`` are the leaf's private + tail, and the copied prefix before it is the inherited path. This is + the Pi-inspired deterministic reconstruction: the same session always + yields the same path, and a branch never rewrites its ancestors. + """ + + events = self.events + completed = [event for event in events if event.type == "branch_completed"] + if len(completed) > 1: + raise ValueError("session contains multiple branch completion markers") + return events + + def append( + self, event_type: str, *, parent_id: str | None = None, **data: Any + ) -> SessionEvent: + """Append, flush, and fsync before making an event visible in memory.""" + + if not _EVENT_TYPE.fullmatch(event_type): + raise ValueError("invalid session event type") + with self._lock: + if self._write_failed: + raise ValueError("session log is unavailable after a failed append") + if len(self._events) >= _MAX_EVENTS: + raise ValueError("session event limit exceeded") + if not self._events and event_type != "session_started": + raise ValueError("session_started must be the first event") + if self._events and event_type == "session_started": + raise ValueError("session metadata appears more than once") + if parent_id is None and self._events: + parent_id = self._events[-1].id + elif self._events: + if parent_id != self._events[-1].id: + raise ValueError("session event chain is invalid") + elif not self._events and parent_id is not None: + if ( + event_type != "session_started" + or data.get("branch_parent_event_id") != parent_id + or not isinstance(data.get("branch_parent_session_id"), str) + ): + raise ValueError("session event chain is invalid") + detached_data = _json_copy(data) + if event_type == "user_message" and detached_data.get("kind") == "steering": + source_id = detached_data.get("source_event_id") + queued = { + item.id: item.data.get("content") + for item in self._events + if item.type == "steering_queued" + } + delivered = { + item.data.get("source_event_id") + for item in self._events + if item.type == "user_message" + and item.data.get("kind") == "steering" + } + if ( + source_id not in queued + or source_id in delivered + or detached_data.get("content") != queued[source_id] + ): + raise ValueError( + "steering delivery source is invalid or duplicated" + ) + if event_type == "tool_result": + call_id = detached_data.get("tool_call_id") + if call_id is not None: + calls = { + item.id for item in self._events if item.type == "tool_call" + } + completed = { + item.data.get("tool_call_id") + for item in self._events + if item.type == "tool_result" + } + if call_id not in calls or call_id in completed: + raise ValueError("tool result references an invalid call") + assistant_id = detached_data.get("assistant_event_id") + if assistant_id is not None: + assistants = { + item.id + for item in self._events + if item.type == "assistant_message" + } + dispositions = { + item.data.get("assistant_event_id") + for item in self._events + if item.type == "tool_result" + } + if ( + call_id is not None + or assistant_id not in assistants + or assistant_id in dispositions + ): + raise ValueError( + "assistant disposition reference is invalid or duplicated" + ) + if event_type == "command_started": + command_id = detached_data.get("command_id") + argv = detached_data.get("argv") + started_commands = { + item.data.get("command_id") + for item in self._events + if item.type == "command_started" + } + if ( + not isinstance(command_id, str) + or _SESSION_ID.fullmatch(command_id) is None + or command_id in started_commands + or not isinstance(argv, list) + or not argv + or any( + not isinstance(argument, str) or not argument + for argument in argv + ) + ): + raise ValueError("command start event is invalid") + if event_type == "command_finished": + command_id = detached_data.get("command_id") + started_commands = { + item.data.get("command_id") + for item in self._events + if item.type == "command_started" + } + finished_commands = { + item.data.get("command_id") + for item in self._events + if item.type == "command_finished" + } + returncode = detached_data.get("returncode") + launched = detached_data.get("launched") + cleanup_incomplete = detached_data.get("cleanup_incomplete") + if ( + command_id not in started_commands + or command_id in finished_commands + or isinstance(returncode, bool) + or not isinstance(returncode, int) + or not isinstance(launched, bool) + or not isinstance(cleanup_incomplete, bool) + ): + raise ValueError("command finish event is invalid") + event = SessionEvent( + uuid.uuid4().hex, + datetime.now(timezone.utc).isoformat(), + event_type, + parent_id, + detached_data, + ) + encoded = ( + json.dumps( + event.to_dict(), + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ) + + "\n" + ).encode("utf-8") + if len(encoded) > _MAX_EVENT_BYTES: + raise ValueError("session event is too large") + if self.path is not None: + flags = os.O_RDWR | os.O_APPEND | getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(self.path, flags) + original_size = None + try: + fcntl.flock(descriptor, fcntl.LOCK_EX) + status = os.fstat(descriptor) + if ( + not stat.S_ISREG(status.st_mode) + or status.st_nlink != 1 + or stat.S_IMODE(status.st_mode) & 0o077 != 0 + or self._identity != (status.st_dev, status.st_ino) + or self._persisted_size != status.st_size + or status.st_size + len(encoded) > _MAX_LOG_BYTES + ): + raise ValueError("session file changed or is invalid") + self._validate_persisted_head(descriptor, status.st_size) + original_size = status.st_size + written = 0 + while written < len(encoded): + count = os.write(descriptor, encoded[written:]) + if count <= 0: + raise OSError("session append made no progress") + written += count + os.fsync(descriptor) + self._events.append(event) + self._persisted_size = original_size + len(encoded) + except BaseException: + event_published = bool( + self._events and self._events[-1].id == event.id + ) + if event_published: + if original_size is not None: + self._persisted_size = original_size + len(encoded) + else: + try: + if original_size is not None: + os.ftruncate(descriptor, original_size) + os.fsync(descriptor) + self._persisted_size = original_size + except BaseException: + self._write_failed = True + raise + finally: + try: + fcntl.flock(descriptor, fcntl.LOCK_UN) + except OSError: + pass + os.close(descriptor) + else: + self._events.append(event) + return SessionEvent.from_dict(event.to_dict()) + + def _validate_persisted_head(self, descriptor: int, size: int) -> None: + """Reject a stale in-memory head before appending another child event.""" + + expected_id = self._events[-1].id if self._events else None + if size == 0: + if expected_id is not None: + raise ValueError("session persisted head is missing") + return + window = min(size, _MAX_EVENT_BYTES) + tail = os.pread(descriptor, window, size - window) + if len(tail) != window or not tail.endswith(b"\n"): + raise ValueError("session persisted tail is invalid") + try: + persisted = SessionEvent.from_dict(json.loads(tail.splitlines()[-1])) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError("session persisted tail is invalid") from error + if persisted.id != expected_id: + raise ValueError("session persisted head changed") + + def messages(self, system_prompt: str) -> list[Message]: + """Rebuild semantic model messages without exposing audit-only events.""" + + blocks = [item.get("content", "") for item in self.instructions] + blocks = [block for block in blocks if isinstance(block, str) and block] + if blocks: + system_prompt += "\n\nProject instructions:\n" + "\n\n".join(blocks) + messages: list[Message] = [{"role": "system", "content": system_prompt}] + for event in self.events: + content = event.data.get("content") + if event.type == "user_message" and isinstance(content, str): + messages.append({"role": "user", "content": content}) + elif event.type == "assistant_message" and isinstance(content, str): + messages.append({"role": "assistant", "content": content}) + elif event.type == "tool_result" and isinstance(content, str): + messages.append({"role": "user", "content": f"Tool result:\n{content}"}) + return messages + + def pending_steering(self) -> tuple[SessionEvent, ...]: + """Return queued corrections not yet materialized at a turn boundary.""" + + events = self.events + queued = { + event.id: event for event in events if event.type == "steering_queued" + } + delivered: set[str] = set() + for event in events: + if event.type != "user_message" or event.data.get("kind") != "steering": + continue + source_id = event.data.get("source_event_id") + if not isinstance(source_id, str) or source_id not in queued: + raise ValueError("delivered steering references an invalid queue event") + if source_id in delivered: + raise ValueError("steering queue event was delivered more than once") + delivered.add(source_id) + return tuple( + event for event_id, event in queued.items() if event_id not in delivered + ) + + def queue_steering(self, message: str) -> SessionEvent: + """Linearize a durable steering submission with terminal acceptance.""" + + if not isinstance(message, str) or not message.strip(): + raise ValueError("steering message must not be blank") + with self._steering_delivery_lock: + return self.append("steering_queued", content=message) + + def run_if_no_pending_steering( + self, operation: Callable[[], _T] + ) -> tuple[bool, _T | None]: + """Run one terminal operation only if queued steering did not win.""" + + with self._steering_delivery_lock: + if self.pending_steering(): + return False, None + return True, operation() + + def deliver_pending_steering(self) -> tuple[SessionEvent, ...]: + """Append queued corrections as semantic messages at a safe boundary.""" + + with self._steering_delivery_lock: + delivered = [] + for event in self.pending_steering(): + content = event.data.get("content") + if not isinstance(content, str) or not content.strip(): + raise ValueError("queued steering message is invalid") + delivered.append( + self.append( + "user_message", + content=content, + kind="steering", + source_event_id=event.id, + ) + ) + return tuple(delivered) + + def recover_unmatched_tool_calls(self) -> tuple[SessionEvent, ...]: + """Record interruption observations without repeating an uncertain call.""" + + events = self.events + matched = { + event.data.get("tool_call_id") + for event in events + if event.type == "tool_result" + } + recovered = [] + for event in events: + if event.type == "tool_call" and event.id not in matched: + recovered.append( + self.append( + "tool_result", + tool_call_id=event.id, + tool=event.data.get("tool"), + is_error=True, + content=( + "error: the prior process stopped before this tool call " + "recorded a result; it was not repeated" + ), + ) + ) + return tuple(recovered) + + def recover_incomplete_turns(self) -> tuple[SessionEvent, ...]: + """Close crash gaps after a durable response without interpreting it.""" + + recovered = list(self.recover_unmatched_tool_calls()) + events = self.events + protocol_events = { + "assistant_message", + "tool_call", + "tool_result", + "run_finished", + "run_started", + "user_message", + } + for index, event in enumerate(events): + if event.type != "assistant_message": + continue + if any( + candidate.type == "tool_result" + and candidate.data.get("assistant_event_id") == event.id + for candidate in events[index + 1 :] + ): + continue + following = next( + ( + candidate + for candidate in events[index + 1 :] + if candidate.type in protocol_events + ), + None, + ) + if ( + following is None + or following.type + in { + "assistant_message", + "run_started", + "user_message", + } + or ( + following.type == "run_finished" + and following.data.get("completed") is not True + ) + ): + recovered.append( + self.append( + "tool_result", + tool_call_id=None, + tool=None, + is_error=True, + assistant_event_id=event.id, + content=( + "error: the prior process stopped after recording a " + "model response; no action from it was repeated" + ), + ) + ) + return tuple(recovered) + + +class SessionStore: + """Create and resume workspace-bound local session transcripts.""" + + def __init__( + self, + workspace: Path, + model: str, + *, + expected_workspace_identity: tuple[int, int] | None = None, + ): + self.workspace = workspace.resolve() + self.model = model + self.directory = self.workspace / ".tiny-llm" / "sessions" + self._expected_workspace_identity = expected_workspace_identity + + def _validate_workspace_identity(self, descriptor: int) -> None: + """Reject a workspace pathname rebound before a session-side effect.""" + + if self._expected_workspace_identity is None: + return + status = os.fstat(descriptor) + if (status.st_dev, status.st_ino) != self._expected_workspace_identity: + raise ValueError("session workspace identity changed") + + def _ensure_directory(self) -> None: + descriptor = self._open_directory(create=True) + os.close(descriptor) + + def _open_directory(self, *, create: bool) -> int: + flags = ( + os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0) + ) + root = os.open(self.workspace, flags) + parent = None + try: + self._validate_workspace_identity(root) + parent_created = False + if create: + try: + os.mkdir(".tiny-llm", 0o700, dir_fd=root) + parent_created = True + except FileExistsError: + pass + if parent_created: + os.fsync(root) + parent = os.open(".tiny-llm", flags, dir_fd=root) + if not stat.S_ISDIR(os.fstat(parent).st_mode): + raise ValueError("session path is unsafe") + os.fchmod(parent, 0o700) + directory_created = False + if create: + try: + os.mkdir("sessions", 0o700, dir_fd=parent) + directory_created = True + except FileExistsError: + pass + if directory_created: + os.fsync(parent) + directory = os.open("sessions", flags, dir_fd=parent) + if not stat.S_ISDIR(os.fstat(directory).st_mode): + os.close(directory) + raise ValueError("session path is unsafe") + os.fchmod(directory, 0o700) + return directory + except OSError as error: + raise ValueError("session path is unsafe") from error + finally: + if parent is not None: + os.close(parent) + os.close(root) + + def _fsync_directory(self, descriptor: int | None = None) -> None: + owned = descriptor is None + if descriptor is None: + descriptor = self._open_directory(create=False) + try: + os.fsync(descriptor) + finally: + if owned: + os.close(descriptor) + + def _instructions(self) -> tuple[dict[str, str], ...]: + flags = ( + os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0) + ) + root = os.open(self.workspace, flags) + try: + self._validate_workspace_identity(root) + try: + content, _ = _read_regular_file( + Path("AGENTS.md"), + _MAX_INSTRUCTIONS_BYTES, + "AGENTS.md", + dir_fd=root, + ) + except ValueError as error: + try: + os.stat("AGENTS.md", dir_fd=root, follow_symlinks=False) + except FileNotFoundError: + return () + raise error + finally: + os.close(root) + text = content.decode("utf-8") + return ( + { + "path": "AGENTS.md", + "sha256": hashlib.sha256(content).hexdigest(), + "content": text, + }, + ) + + def create(self, session_id: str | None = None) -> SessionLog: + return self._create_session(session_id) + + def _create_session( + self, + session_id: str | None, + *, + branch_parent_session_id: str | None = None, + branch_parent_event_id: str | None = None, + ) -> SessionLog: + instructions = self._instructions() + session_id = session_id or uuid.uuid4().hex + if not _SESSION_ID.fullmatch(session_id): + raise ValueError("invalid session ID") + directory = self._open_directory(create=True) + path = self.directory / f"{session_id}.jsonl" + temporary = self.directory / f".{session_id}-{uuid.uuid4().hex}.tmp" + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0) + try: + descriptor = os.open(temporary.name, flags, 0o600, dir_fd=directory) + except BaseException: + os.close(directory) + raise + status = os.fstat(descriptor) + os.close(descriptor) + log = SessionLog( + session_id, + self.workspace, + self.model, + path=temporary, + identity=(status.st_dev, status.st_ino), + ) + linked = False + try: + metadata: dict[str, Any] = { + "session_id": session_id, + "workspace": str(self.workspace), + "model": self.model, + "instructions": instructions, + } + if branch_parent_event_id is not None: + metadata.update( + branch_id=session_id, + branch_parent_session_id=branch_parent_session_id, + branch_parent_event_id=branch_parent_event_id, + ) + log.append( + "session_started", + parent_id=branch_parent_event_id, + **metadata, + ) + os.link( + temporary.name, + path.name, + src_dir_fd=directory, + dst_dir_fd=directory, + follow_symlinks=False, + ) + linked = True + os.unlink(temporary.name, dir_fd=directory) + log.path = path + self._fsync_directory(directory) + except BaseException: + try: + os.unlink(temporary.name, dir_fd=directory) + except FileNotFoundError: + pass + if linked: + try: + os.unlink(path.name, dir_fd=directory) + except FileNotFoundError: + pass + self._fsync_directory(directory) + raise + finally: + os.close(directory) + return log + + @staticmethod + def _validate_events(events: list[SessionEvent]) -> None: + if not events or events[0].type != "session_started": + raise ValueError("session metadata is missing") + external_parent = events[0].data.get("branch_parent_event_id") + external_session = events[0].data.get("branch_parent_session_id") + if events[0].parent_id is None and ( + external_parent is not None or external_session is not None + ): + raise ValueError("session branch metadata is invalid") + if events[0].parent_id is not None and ( + events[0].parent_id != external_parent + or not isinstance(external_session, str) + or not _SESSION_ID.fullmatch(external_session) + or not isinstance(external_parent, str) + or not _SESSION_ID.fullmatch(external_parent) + ): + raise ValueError("session event chain is invalid") + calls: set[str] = set() + results: set[str] = set() + queued_steering: dict[str, object] = {} + delivered_steering: set[str] = set() + assistants: set[str] = set() + assistant_dispositions: set[str] = set() + commands: set[str] = set() + command_results: set[str] = set() + for index, event in enumerate(events): + if index and event.type == "session_started": + raise ValueError("session metadata appears more than once") + if index and event.parent_id != events[index - 1].id: + raise ValueError("session event chain is invalid") + if event.type == "tool_call": + calls.add(event.id) + elif event.type == "steering_queued": + queued_steering[event.id] = event.data.get("content") + elif event.type == "user_message" and event.data.get("kind") == "steering": + source_id = event.data.get("source_event_id") + if ( + source_id not in queued_steering + or source_id in delivered_steering + or event.data.get("content") != queued_steering[source_id] + ): + raise ValueError( + "steering delivery source is invalid or duplicated" + ) + delivered_steering.add(source_id) + elif event.type == "assistant_message": + assistants.add(event.id) + if event.type == "tool_result": + call_id = event.data.get("tool_call_id") + if call_id is not None: + if call_id not in calls or call_id in results: + raise ValueError("tool result references an invalid call") + results.add(call_id) + assistant_id = event.data.get("assistant_event_id") + if assistant_id is not None: + if ( + call_id is not None + or assistant_id not in assistants + or assistant_id in assistant_dispositions + ): + raise ValueError( + "assistant disposition reference is invalid or duplicated" + ) + assistant_dispositions.add(assistant_id) + if event.type == "command_started": + command_id = event.data.get("command_id") + argv = event.data.get("argv") + if ( + not isinstance(command_id, str) + or _SESSION_ID.fullmatch(command_id) is None + or command_id in commands + or not isinstance(argv, list) + or not argv + or any( + not isinstance(argument, str) or not argument + for argument in argv + ) + ): + raise ValueError("command start event is invalid") + commands.add(command_id) + elif event.type == "command_finished": + command_id = event.data.get("command_id") + returncode = event.data.get("returncode") + launched = event.data.get("launched") + cleanup_incomplete = event.data.get("cleanup_incomplete") + if ( + command_id not in commands + or command_id in command_results + or isinstance(returncode, bool) + or not isinstance(returncode, int) + or not isinstance(launched, bool) + or not isinstance(cleanup_incomplete, bool) + ): + raise ValueError("command finish event is invalid") + command_results.add(command_id) + + def _recover_publication(self, directory: int, name: str) -> None: + """Finish the link publication if a crash left its private name.""" + + try: + status = os.stat(name, dir_fd=directory, follow_symlinks=False) + except FileNotFoundError as error: + raise ValueError("session file does not exist or is unsafe") from error + if status.st_nlink == 1: + return + prefix = f".{Path(name).stem}-" + matches = [] + for candidate in os.listdir(directory): + if not candidate.startswith(prefix) or not candidate.endswith(".tmp"): + continue + candidate_status = os.stat( + candidate, dir_fd=directory, follow_symlinks=False + ) + if (candidate_status.st_dev, candidate_status.st_ino) == ( + status.st_dev, + status.st_ino, + ): + matches.append(candidate) + if status.st_nlink != 2 or len(matches) != 1: + raise ValueError("session file has an unsafe link count") + os.unlink(matches[0], dir_fd=directory) + os.fsync(directory) + + def _truncate_partial_tail( + self, + path: Path, + expected_content: bytes, + complete_size: int, + identity: tuple[int, int], + ) -> None: + """Discard only a verified incomplete final append from a private log.""" + + if complete_size <= 0 or complete_size >= len(expected_content): + raise ValueError("session partial-tail boundary is invalid") + directory = self._open_directory(create=False) + descriptor = None + try: + descriptor = os.open( + path.name, + os.O_RDWR | getattr(os, "O_NOFOLLOW", 0), + dir_fd=directory, + ) + fcntl.flock(descriptor, fcntl.LOCK_EX) + status = os.fstat(descriptor) + if ( + not stat.S_ISREG(status.st_mode) + or status.st_nlink != 1 + or stat.S_IMODE(status.st_mode) & 0o077 != 0 + or (status.st_dev, status.st_ino) != identity + or status.st_size != len(expected_content) + ): + raise ValueError("session changed during partial-tail recovery") + os.lseek(descriptor, 0, os.SEEK_SET) + current = bytearray() + while len(current) < len(expected_content): + chunk = os.read(descriptor, len(expected_content) - len(current)) + if not chunk: + break + current.extend(chunk) + if bytes(current) != expected_content: + raise ValueError("session changed during partial-tail recovery") + os.ftruncate(descriptor, complete_size) + os.fsync(descriptor) + os.fsync(directory) + finally: + if descriptor is not None: + try: + fcntl.flock(descriptor, fcntl.LOCK_UN) + except OSError: + pass + os.close(descriptor) + os.close(directory) + + def load(self, session_id: str, *, recover: bool = True) -> SessionLog: + if not _SESSION_ID.fullmatch(session_id): + raise ValueError("invalid session ID") + path = self.directory / f"{session_id}.jsonl" + directory = self._open_directory(create=False) + try: + self._recover_publication(directory, path.name) + content, identity = _read_regular_file( + Path(path.name), + _MAX_LOG_BYTES, + "session file", + dir_fd=directory, + private=True, + ) + finally: + os.close(directory) + partial_tail = bool(content and not content.endswith(b"\n")) + if partial_tail and not recover: + raise ValueError("session file ends with a partial event") + complete_size = len(content) + parse_content = content + if partial_tail: + complete_size = content.rfind(b"\n") + 1 + if complete_size <= 0: + raise ValueError("session metadata ends with a partial event") + parse_content = content[:complete_size] + events = [] + seen = set() + for number, line in enumerate(parse_content.splitlines(), 1): + if number > _MAX_EVENTS or len(line) > _MAX_EVENT_BYTES: + raise ValueError("session event limit exceeded") + try: + event = SessionEvent.from_dict(json.loads(line)) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError(f"invalid session JSONL line {number}") from error + if event.id in seen: + raise ValueError("duplicate session event ID") + seen.add(event.id) + events.append(event) + self._validate_events(events) + metadata = events[0].data + if metadata.get("branch_parent_event_id") is not None: + completed = [event for event in events if event.type == "branch_completed"] + if len(completed) != 1: + raise ValueError("session branch is incomplete") + if metadata.get("session_id") != session_id: + raise ValueError("session metadata ID does not match its filename") + if metadata.get("workspace") != str(self.workspace): + raise ValueError("session belongs to a different workspace") + if metadata.get("model") != self.model: + raise ValueError("session belongs to a different model") + if partial_tail: + self._truncate_partial_tail(path, content, complete_size, identity) + log = SessionLog( + session_id, + self.workspace, + self.model, + path=path, + events=tuple(events), + identity=identity, + persisted_size=complete_size, + ) + if recover: + if partial_tail: + log.append( + "session_tail_recovered", + discarded_bytes=len(content) - complete_size, + ) + current = self._instructions() + if current != log.instructions: + log.append( + "instructions_changed", previous=log.instructions, current=current + ) + log.recover_incomplete_turns() + return log + + def branch( + self, + session_id: str, + at_event_id: str, + new_session_id: str | None = None, + ) -> SessionLog: + """Create a private lineage with model context ending at one ancestor.""" + + parent = self.load(session_id, recover=False) + parent_events = parent.events + self._validate_branch_side_effect_state(parent_events) + selected_index = next( + ( + index + for index, event in enumerate(parent_events) + if event.id == at_event_id + ), + None, + ) + if selected_index is None or selected_index == 0: + raise ValueError("branch ancestor is not present in the parent session") + prefix = parent_events[: selected_index + 1] + calls = {event.id for event in prefix if event.type == "tool_call"} + results = { + event.data.get("tool_call_id") + for event in prefix + if event.type == "tool_result" + } + if calls - results or prefix[-1].type in { + "assistant_message", + "mutation_intent", + "summary_attempt", + "tool_call", + }: + raise ValueError("branch ancestor does not close a safe event boundary") + + branch = self._create_session( + new_session_id, + branch_parent_session_id=parent.session_id, + branch_parent_event_id=at_event_id, + ) + branch.append( + "branch_created", + parent_session_id=parent.session_id, + at_event_id=at_event_id, + source_event_count=len(prefix), + ) + copied_types = { + "assistant_message", + "run_finished", + "run_started", + "steering_queued", + "tool_call", + "tool_result", + "user_message", + } + copied_ids: dict[str, str] = {} + for event in prefix[1:]: + if event.type not in copied_types: + continue + data = _json_copy(event.data) + if event.type == "tool_result": + source_call_id = data.get("tool_call_id") + if source_call_id is not None: + if source_call_id not in copied_ids: + raise ValueError( + "branch prefix contains an invalid tool result" + ) + data["tool_call_id"] = copied_ids[source_call_id] + source_assistant_id = data.get("assistant_event_id") + if source_assistant_id is not None: + if source_assistant_id not in copied_ids: + raise ValueError( + "branch prefix contains an invalid assistant disposition" + ) + data["assistant_event_id"] = copied_ids[source_assistant_id] + if event.type == "user_message" and data.get("kind") == "steering": + source_steering_id = data.get("source_event_id") + if source_steering_id not in copied_ids: + raise ValueError("branch prefix contains invalid steering") + data["source_parent_steering_event_id"] = source_steering_id + data["source_event_id"] = copied_ids[source_steering_id] + data["source_parent_event_id"] = event.id + copied = branch.append(event.type, **data) + copied_ids[event.id] = copied.id + branch.append( + "branch_completed", + parent_session_id=parent.session_id, + at_event_id=at_event_id, + copied_event_count=len(copied_ids), + ) + return branch + + def parent_of(self, session_id: str) -> tuple[str, str] | None: + """Return ``(parent_session_id, at_event_id)`` for one branch session. + + The root session of the tree returns ``None``. The answer is derived + only from the durable ``session_started`` metadata, so reconstruction + never depends on in-memory state. + """ + + log = self.load(session_id, recover=False) + first = log.events[0] + parent_session = first.data.get("branch_parent_session_id") + parent_event = first.data.get("branch_parent_event_id") + if parent_session is None or parent_event is None: + return None + if ( + not isinstance(parent_session, str) + or not _SESSION_ID.fullmatch(parent_session) + or not isinstance(parent_event, str) + or not _SESSION_ID.fullmatch(parent_event) + ): + raise ValueError("session branch metadata is invalid") + return (parent_session, parent_event) + + def active_path(self, session_id: str) -> tuple[str, ...]: + """Reconstruct the deterministic root-to-leaf session path. + + Walk ``parent_of`` from the given leaf back to the root, then return + the path root-first. Any cycle or dangling parent fails closed; this + is the Pi-inspired ``id``/``parentId`` tree navigation used by + deterministic active-path reconstruction. + """ + + path: list[str] = [] + seen: set[str] = set() + current: str | None = session_id + while current is not None: + if not _SESSION_ID.fullmatch(current): + raise ValueError("invalid session ID in active path") + if current in seen: + raise ValueError("session tree contains a cycle") + seen.add(current) + path.append(current) + parent = self.parent_of(current) + current = parent[0] if parent is not None else None + return tuple(reversed(path)) + + @staticmethod + def _validate_branch_side_effect_state( + events: tuple[SessionEvent, ...], + ) -> None: + """Reject a branch while shared-workspace side effects are unresolved.""" + + calls: set[str] = set() + call_results: set[str] = set() + mutations: set[str] = set() + mutation_results: set[str] = set() + undo_changes: set[str] = set() + undo_results: set[str] = set() + commands: set[str] = set() + command_results: set[str] = set() + conflicts = False + malformed = False + + def valid_id(value: object) -> bool: + return isinstance(value, str) and _SESSION_ID.fullmatch(value) is not None + + for event in events: + data = event.data + if event.type == "tool_call": + calls.add(event.id) + elif event.type == "tool_result": + call_id = data.get("tool_call_id") + if isinstance(call_id, str): + call_results.add(call_id) + elif event.type == "mutation_intent": + intent_id = data.get("intent_id") + if valid_id(intent_id) and intent_id not in mutations: + mutations.add(intent_id) + else: + malformed = True + elif event.type in {"mutation_committed", "mutation_recovered"}: + intent_id = data.get("intent_id") + if ( + valid_id(intent_id) + and intent_id in mutations + and intent_id not in mutation_results + ): + mutation_results.add(intent_id) + else: + malformed = True + if event.type == "mutation_recovered": + if data.get("status") == "conflict": + conflicts = True + elif data.get("status") != "not_applied": + malformed = True + elif event.type == "undo_change_started": + change_id = data.get("undo_change_id") + if valid_id(change_id) and change_id not in undo_changes: + undo_changes.add(change_id) + else: + malformed = True + elif event.type in {"undo_change_finished", "undo_change_recovered"}: + change_id = data.get("undo_change_id") + if ( + valid_id(change_id) + and change_id in undo_changes + and change_id not in undo_results + ): + undo_results.add(change_id) + else: + malformed = True + if event.type == "undo_change_recovered": + if data.get("status") == "conflict": + conflicts = True + elif data.get("status") != "not_applied": + malformed = True + elif event.type == "undo_completed": + recorded_conflicts = data.get("conflicts") + if not isinstance(recorded_conflicts, list): + malformed = True + elif recorded_conflicts: + conflicts = True + elif event.type == "command_started": + command_id = data.get("command_id") + if valid_id(command_id) and command_id not in commands: + commands.add(command_id) + else: + malformed = True + elif event.type == "command_finished": + command_id = data.get("command_id") + if ( + valid_id(command_id) + and command_id in commands + and command_id not in command_results + ): + command_results.add(command_id) + else: + malformed = True + if not isinstance(data.get("launched"), bool) or not isinstance( + data.get("cleanup_incomplete"), bool + ): + malformed = True + elif data["cleanup_incomplete"]: + conflicts = True + if ( + malformed + or conflicts + or calls - call_results + or mutations - mutation_results + or mutation_results - mutations + or undo_changes - undo_results + or undo_results - undo_changes + or commands - command_results + or command_results - commands + ): + raise ValueError( + "parent session has unresolved or conflicting workspace side effects" + ) + + def latest(self) -> SessionLog: + directory = self._open_directory(create=False) + try: + candidates = [] + for name in os.listdir(directory): + path = Path(name) + if path.suffix != ".jsonl" or not _SESSION_ID.fullmatch(path.stem): + continue + status = os.stat(name, dir_fd=directory, follow_symlinks=False) + if stat.S_ISREG(status.st_mode): + candidates.append((status.st_mtime_ns, name)) + candidates.sort(reverse=True) + if not candidates: + raise ValueError("no saved session exists") + for _, name in candidates: + self._recover_publication(directory, name) + content, _ = _read_regular_file( + Path(name), + _MAX_LOG_BYTES, + "session file", + dir_fd=directory, + private=True, + ) + first_line = content.splitlines()[:1] + if not first_line: + raise ValueError("session metadata is missing") + try: + metadata = SessionEvent.from_dict(json.loads(first_line[0])) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError("invalid session metadata") from error + if metadata.type != "session_started": + raise ValueError("session metadata is missing") + if metadata.data.get("session_id") != Path(name).stem: + raise ValueError("session metadata ID does not match its filename") + if metadata.data.get("workspace") != str(self.workspace): + raise ValueError("session belongs to a different workspace") + if metadata.data.get("model") == self.model: + if metadata.data.get("branch_parent_event_id") is not None: + if content and not content.endswith(b"\n"): + continue + try: + branch_complete = any( + SessionEvent.from_dict(json.loads(line)).type + == "branch_completed" + for line in content.splitlines()[1:] + ) + except ( + UnicodeDecodeError, + json.JSONDecodeError, + ValueError, + ): + continue + if not branch_complete: + continue + selected = Path(name).stem + break + else: + raise ValueError("no saved session exists for this model") + finally: + os.close(directory) + return self.load(selected) + + +def memory_session(workspace: Path, model: str) -> SessionLog: + """Create an event log that never writes a transcript to disk.""" + + store = SessionStore(workspace, model) + log = SessionLog(uuid.uuid4().hex, store.workspace, model) + log.append( + "session_started", + session_id=log.session_id, + workspace=str(store.workspace), + model=model, + instructions=store._instructions(), + ) + return log diff --git a/src/tiny_llm_ref/agent/status.py b/src/tiny_llm_ref/agent/status.py new file mode 100644 index 00000000..a7ac9909 --- /dev/null +++ b/src/tiny_llm_ref/agent/status.py @@ -0,0 +1,167 @@ +# WARNING: Under review - generated by LLM. + +"""Public status cards derived from the durable event trace. + +An operator may ask "what are you doing?" while an agent is running. The +answer is a deterministic ``AgentStateCard`` derived only from the durable +event trace — never from hidden chain-of-thought. The card contains only +ledger-derived public state: goal, current action, approval state, last +evidence, and next safe step. + +A status question is answered on the *main* run via sequential snapshot and +rewind: the harness records the current cache prefix (feature 5), decodes a +short status paraphrase, then rewinds to the prefix so the main run continues +unchanged. No COW fork is needed; the main trace and workspace are untouched. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from .branch import BranchStats, SequentialBranch + +if TYPE_CHECKING: + from .generation import GenerationSession + from .session import SessionLog + + +@dataclass(frozen=True) +class AgentStateCard: + """Deterministic public state derived only from the durable event trace.""" + + goal: str + current_action: str + approval_state: str + last_evidence: str + next_safe_step: str + phase: str = "running" + + def to_dict(self) -> dict[str, Any]: + return { + "goal": self.goal, + "current_action": self.current_action, + "approval_state": self.approval_state, + "last_evidence": self.last_evidence, + "next_safe_step": self.next_safe_step, + "phase": self.phase, + } + + def render(self) -> str: + """Render the card as plain text for display or a status query.""" + + lines = [ + f"goal: {self.goal}", + f"current action: {self.current_action}", + f"approval: {self.approval_state}", + f"last evidence: {self.last_evidence}", + f"next safe step: {self.next_safe_step}", + f"phase: {self.phase}", + ] + return "\n".join(lines) + + +def build_state_card( + session: "SessionLog", + *, + approval_state: str = "unknown", + goal: str = "", +) -> AgentStateCard: + """Derive a public state card from the durable event trace only. + + Every field is a function of recorded events and operator-provided + approval state; hidden reasoning is never requested or exposed. + """ + + current_action = "idle" + last_evidence = "" + next_safe_step = "awaiting instruction" + phase = "running" + for event in session.events: + if event.type == "tool_call": + tool = event.data.get("tool", "?") + arguments = event.data.get("arguments", {}) + current_action = f"{tool} {json.dumps(arguments, sort_keys=True)}" + next_safe_step = "deliver tool result" + elif event.type == "tool_result": + content = event.data.get("content", "") + if isinstance(content, str): + last_evidence = content[:80] + current_action = "observing tool result" + next_safe_step = "continue the loop" + elif event.type == "assistant_message": + current_action = "generating" + next_safe_step = "validate the action" + elif event.type == "run_finished": + phase = "finished" if event.data.get("completed") else "stopped" + current_action = "none" + next_safe_step = "none" + return AgentStateCard( + goal=goal, + current_action=current_action, + approval_state=approval_state, + last_evidence=last_evidence, + next_safe_step=next_safe_step, + phase=phase, + ) + + +@dataclass(frozen=True) +class StatusQueryResult: + """The answer plus honest accounting of the sequential side query.""" + + response: str + stats: BranchStats + + +class StatusQuery: + """Answer one status question on the main run via snapshot + rewind. + + The query pauses the main decode at the current cache prefix, decodes a + short tool-disabled paraphrase of the public card, then rewinds to the + prefix. The main token sequence, cache, trace, and workspace are + byte-identical to a run without the query. + """ + + def __init__(self, session: "GenerationSession"): + self._branch = SequentialBranch(session) + + def ask(self, card: AgentStateCard, generate: Any) -> StatusQueryResult: + """Run one status paraphrase from the current prefix, then rewind. + + ``generate`` is a callable that maps a message list to a short answer. + The branch checkpoints the current prefix, asks, and rewinds so the + main run continues unchanged. + """ + + messages = [ + { + "role": "system", + "content": ( + "Answer with a short public status summary. " + "Do not reveal hidden reasoning. " + "You have no tools." + ), + }, + { + "role": "user", + "content": f"Current public state:\n{card.render()}", + }, + ] + if self._branch.checkpoint_tokens is None: + self._branch.checkpoint() + output = generate(messages) + rewind = self._branch.rewind() + stats = BranchStats( + checkpoint_tokens=rewind, + reused_tokens=0, + rewound_tokens=0, + prefilled_tokens=0, + output_tokens=0, + cold_start=False, + ) + return StatusQueryResult(output, stats) + + def close(self) -> None: + self._branch.close() diff --git a/src/tiny_llm_ref/agent/workspace.py b/src/tiny_llm_ref/agent/workspace.py new file mode 100644 index 00000000..6c77cb8d --- /dev/null +++ b/src/tiny_llm_ref/agent/workspace.py @@ -0,0 +1,706 @@ +# WARNING: Under review - generated by LLM. + +"""Week 4, Day 2: authorize effects and record durable receipts. + +The workspace bounds every tool to one explicit root, rejects protected and +secret paths, refuses symlink traversal, requires inspection before +overwrite, and atomically replaces files so a crash never leaves a torn +write. Model-dispatched side effects are default-deny: an operator +``confirm_tool`` callback must approve writes, edits, and commands. + +Every dispatched tool records an immutable ``EffectReceipt`` (Day 2 feature): +the validated action, normalized inputs, exit state, changed artifacts, and a +content-address. Durable crash/effect recovery is taught by Day 6's +exactly-once reconcile; this file deliberately has no mutation journal or +undo machinery. +""" + +import hashlib +import heapq +import os +import signal +import stat +import subprocess +import uuid +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from threading import Event, Thread +from time import monotonic + +from .protocol import AgentError, ToolAction +from .receipts import EffectReceipt, ReceiptStore + + +_PROTECTED_NAMES = frozenset( + { + ".env", + ".aws", + ".azure", + ".docker", + ".git", + ".gnupg", + ".kube", + ".netrc", + ".npmrc", + ".pypirc", + ".ssh", + ".tiny-llm", + "id_rsa", + "id_ed25519", + } +) + + +def _protected_path_reason(parts: tuple[str, ...]) -> str | None: + """Describe why path components cross a protected metadata or secret path.""" + + for part in parts: + lower = part.lower() + if lower == ".git": + return ".git is not accessible" + if ( + lower in _PROTECTED_NAMES + or lower.startswith(".env.") + or lower.startswith(".tiny-llm-agent-") + or lower.startswith(".tiny-llm-undo-") + ): + return "potential secret files are not accessible" + if lower.endswith((".pem", ".key")): + return "potential key files are not accessible" + return None + + +@dataclass(frozen=True) +class _PreparedWrite: + """A validated write and the file state it is allowed to replace.""" + + path: Path + content: bytes + expected_digest: bytes | None + expected_mode: int | None + after_mode: int + parent_identity: tuple[int, int] + + +@dataclass(frozen=True) +class ToolPolicy: + """Filesystem and command boundaries for one workspace.""" + + root: Path + allow_writes: bool = False + allowed_commands: tuple[tuple[str, ...], ...] = () + max_file_bytes: int = 64 * 1024 + max_write_bytes: int = 64 * 1024 + max_list_entries: int = 200 + max_tool_output_chars: int = 16_000 + command_timeout_seconds: float = 30.0 + _root_identity: tuple[int, int] = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + """Normalize the root and reject invalid limits.""" + + root = self.root.resolve() + object.__setattr__(self, "root", root) + numeric_limits = ( + self.max_file_bytes, + self.max_write_bytes, + self.max_list_entries, + self.max_tool_output_chars, + self.command_timeout_seconds, + ) + if any(limit <= 0 for limit in numeric_limits): + raise ValueError("tool policy limits must be positive") + if not root.exists() or not root.is_dir(): + raise ValueError("workspace root must be an existing directory") + root_status = os.stat(root, follow_symlinks=False) + if not stat.S_ISDIR(root_status.st_mode): + raise ValueError("workspace root must be a directory") + object.__setattr__( + self, + "_root_identity", + (root_status.st_dev, root_status.st_ino), + ) + home = Path.home().resolve() + if root.parent == root or root == home or root in home.parents: + raise ValueError("workspace root is too broad") + if reason := _protected_path_reason(root.parts): + raise ValueError(f"workspace root is protected: {reason}") + for command in self.allowed_commands: + if not command or any( + not isinstance(part, str) or not part or "\x00" in part + for part in command + ): + raise ValueError("allowed commands must contain non-empty arguments") + + +@dataclass +class Workspace: + """Bounded tools over one explicit workspace root.""" + + policy: ToolPolicy + confirm_tool: Callable[[ToolAction], bool] | None = None + observed_files: dict[Path, bytes] = field(default_factory=dict, init=False) + modified_files: set[Path] = field(default_factory=set, init=False) + uncertain_modified_files: set[Path] = field(default_factory=set, init=False) + receipt_store: ReceiptStore | None = field(default=None, init=False) + last_receipt: EffectReceipt | None = field(default=None, init=False) + + def bind_receipt_store(self, store: ReceiptStore) -> None: + """Attach a durable receipt store; every effect after this is recorded.""" + + if self.last_receipt is not None or store is None: + raise ValueError("receipt store must be bound before any effect") + self.receipt_store = store + + def _record_receipt( + self, + action: ToolAction, + result: str, + changed_artifacts: tuple[str, ...], + *, + tool_call_id: str | None = None, + ) -> None: + """Persist one immutable effect receipt for a dispatched tool.""" + + exit_state = "uncertain" + if result.startswith("error:"): + exit_state = "error" + elif not self.uncertain_modified_files: + exit_state = "ok" + receipt = EffectReceipt( + tool_call_id=tool_call_id or f"{uuid.uuid4().hex}", + tool=action.tool, + arguments=action.arguments, + exit_state=exit_state, + result=result, + changed_artifacts=changed_artifacts, + ) + self.last_receipt = receipt + if self.receipt_store is not None: + self.receipt_store.put(receipt) + + @property + def available_tools(self) -> frozenset[str]: + """Expose only tools enabled by operator policy.""" + + tools = {"list_files", "read_file"} + if self.policy.allow_writes: + tools.update({"write_file", "edit_file"}) + if self.policy.allowed_commands: + tools.add("run_command") + return frozenset(tools) + + def resolve_path(self, raw: str, *, must_exist: bool = True) -> Path: + """Resolve one tool path inside the workspace, rejecting traversal.""" + + if not isinstance(raw, str) or not raw or "\x00" in raw: + raise AgentError("path must be a non-empty string") + candidate = Path(raw) + if not candidate.is_absolute(): + candidate = self.policy.root / candidate + try: + relative = candidate.resolve().relative_to(self.policy.root) + except ValueError as error: + raise AgentError("path escapes the workspace") from error + self._reject_protected(relative) + path = self.policy.root / relative + if must_exist: + # Reject symlinked components: a resolved path that differs from + # the literal path means a link was followed. + try: + literal = candidate.relative_to(self.policy.root) + except ValueError as error: + raise AgentError("path escapes the workspace") from error + for part in literal.parts: + probe = self.policy.root.joinpath( + *literal.parts[: literal.parts.index(part) + 1] + ) + if probe.is_symlink(): + raise AgentError("symlinks are not accessible") + if not path.exists(): + raise AgentError("path does not exist") + return path + + def _reject_protected(self, relative: Path) -> None: + if reason := _protected_path_reason(relative.parts): + raise AgentError(reason) + + def list_files(self, raw: str = ".") -> str: + """List one directory without following symlinks.""" + + path = self.resolve_path(raw) + if not path.is_dir(): + raise AgentError("list_files path must be a directory") + lines: list[str] = [] + items = heapq.nsmallest( + self.policy.max_list_entries, + path.iterdir(), + key=lambda entry: entry.name, + ) + for item in items: + if item.is_symlink(): + continue + try: + relative = item.relative_to(self.policy.root) + self._reject_protected(relative) + except AgentError: + continue + kind = "dir" if item.is_dir() else "file" + lines.append(f"{kind} {relative}") + if len(lines) == self.policy.max_list_entries: + break + return "\n".join(lines) + + def read_file(self, raw: str) -> str: + """Read one bounded UTF-8 regular file.""" + + path = self.resolve_path(raw) + data = self._read_bounded_file(path, tool="read_file") + content = data.decode("utf-8") + self.observed_files[path] = self._digest(data) + return content + + def write_file(self, raw: str, content: str) -> str: + """Atomically create or replace an inspected file.""" + + prepared = self._prepare_write(raw, content) + return self._commit_write(prepared) + + def _prepare_write(self, raw: str, content: str) -> _PreparedWrite: + """Validate a write completely before asking for operator approval.""" + + if not self.policy.allow_writes: + raise AgentError("writes are disabled; restart with --allow-writes") + encoded = content.encode("utf-8") + if len(encoded) > self.policy.max_write_bytes: + raise AgentError(f"content exceeds {self.policy.max_write_bytes} bytes") + path = self.resolve_path(raw, must_exist=False) + relative = path.relative_to(self.policy.root) + expected_digest = self.observed_files.get(path) + parent, name = self._open_parent_directory(relative) + try: + parent_status = os.fstat(parent) + current = self._read_regular_at(parent, name, tool="write_file") + finally: + os.close(parent) + if current is None: + if expected_digest is not None: + raise AgentError("file changed since it was read; read it again") + expected_mode = None + after_mode = 0o600 + else: + if expected_digest is None: + raise AgentError( + "existing files must be read before they are overwritten" + ) + current_content, current_status = current + if self._digest(current_content) != expected_digest: + raise AgentError("file changed since it was read; read it again") + expected_mode = stat.S_IMODE(current_status.st_mode) + after_mode = expected_mode & 0o777 + return _PreparedWrite( + path, + encoded, + expected_digest, + expected_mode, + after_mode, + (parent_status.st_dev, parent_status.st_ino), + ) + + def _commit_write(self, prepared: _PreparedWrite) -> str: + """Revalidate and commit one previously prepared write.""" + + relative = prepared.path.relative_to(self.policy.root) + path = prepared.path + self._revalidate_prepared_write(prepared) + self.uncertain_modified_files.add(path) + try: + self._atomic_write( + path, + prepared.content, + expected_digest=prepared.expected_digest, + expected_mode=prepared.expected_mode, + after_mode=prepared.after_mode, + parent_identity=prepared.parent_identity, + ) + except BaseException: + self.uncertain_modified_files.discard(path) + raise + self.observed_files[path] = self._digest(prepared.content) + self.modified_files.add(path) + self.uncertain_modified_files.discard(path) + return f"wrote {relative}" + + def edit_file(self, raw: str, old: str, new: str) -> str: + """Make one exact, reviewable replacement in a read file.""" + + prepared = self._prepare_edit(raw, old, new) + return self._commit_write(prepared) + + def _prepare_edit(self, raw: str, old: str, new: str) -> _PreparedWrite: + """Validate and compute an edit without changing the workspace.""" + + if not self.policy.allow_writes: + raise AgentError("writes are disabled; restart with --allow-writes") + path = self.resolve_path(raw) + if path not in self.observed_files: + raise AgentError("files must be read before they are edited") + if not old: + raise AgentError("old text must not be empty") + data = self._read_bounded_file(path, tool="edit_file") + if self._digest(data) != self.observed_files[path]: + raise AgentError("file changed since it was read; read it again") + content = data.decode("utf-8") + matches = content.count(old) + if matches != 1: + raise AgentError(f"old text must match exactly once; found {matches}") + return self._prepare_write(raw, content.replace(old, new, 1)) + + def run_command(self, argv: list[str]) -> str: + """Run one exact allowlisted command with a bounded timeout.""" + + command = self._prepare_command(argv) + return self._run_command(command) + + def _prepare_command(self, argv: list[str]) -> tuple[str, ...]: + if not self.policy.allowed_commands: + raise AgentError("command execution is disabled") + if ( + not isinstance(argv, list) + or not argv + or any(not isinstance(part, str) or not part for part in argv) + ): + raise AgentError("argv must be a non-empty array of non-empty strings") + command = tuple(argv) + if command not in self.policy.allowed_commands: + raise AgentError("command is not explicitly allowed") + return command + + def _run_command(self, argv: tuple[str, ...]) -> tuple[str, tuple[str, ...]]: + """Execute one command with a hard timeout, returning bounded output.""" + + process = subprocess.Popen( + list(argv), + cwd=self.policy.root, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + timeout = self.policy.command_timeout_seconds + deadline = monotonic() + timeout + output: list[bytes] = [] + finished = Event() + warnings: tuple[str, ...] = () + + def reader() -> None: + try: + while True: + chunk = process.stdout.read(1 << 16) + if not chunk: + break + output.append(chunk) + if ( + sum(len(part) for part in output) + > self.policy.max_tool_output_chars + ): + break + finally: + finished.set() + + thread = Thread(target=reader, daemon=True) + thread.start() + try: + # Drain the pipe until the reader finishes or the timeout fires. + while thread.is_alive() and monotonic() < deadline: + thread.join(timeout=0.05) + if thread.is_alive(): + self._kill_and_reap(process) + warnings = ("command timed out and was killed",) + thread.join(timeout=2.0) + else: + process.wait(timeout=2.0) + finally: + if thread.is_alive(): + thread.join(timeout=2.0) + self.uncertain_modified_files.update(self._recent_modified()) + return self._truncate_result( + b"".join(output).decode("utf-8", errors="replace"), + preserve_command_warnings=warnings, + ) + + def _recent_modified(self) -> set[Path]: + """Return files changed inside the workspace since the last snapshot.""" + + changed: set[Path] = set() + baseline = {path: digest for path, digest in self.observed_files.items()} + for path in list(baseline): + try: + data = self._read_bounded_file(path, tool="run_command") + except (OSError, ValueError, AgentError): + changed.add(path) + continue + if self._digest(data) != baseline[path]: + changed.add(path) + return changed + + def _kill_and_reap(self, process) -> bool: + """Terminate a whole process group and wait for it to exit.""" + + try: + os.killpg(process.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + try: + process.kill() + except ProcessLookupError: + return process.poll() is not None + try: + process.wait(timeout=2.0) + except subprocess.TimeoutExpired: + return False + return process.returncode is not None + + def _require_confirmation(self, action: ToolAction) -> None: + """Default-deny one model-requested side effect.""" + + if self.confirm_tool is None or not self.confirm_tool(action): + raise AgentError(f"operator denied {action.tool}") + + def execute(self, action: ToolAction, *, tool_call_id: str | None = None) -> str: + """Dispatch a validated action and return recoverable errors.""" + + before_modified = frozenset(self.modified_files) + before_uncertain = frozenset(self.uncertain_modified_files) + try: + if action.tool == "list_files": + result = self.list_files(action.arguments.get("path", ".")) + elif action.tool == "read_file": + result = self.read_file(action.arguments["path"]) + elif action.tool == "write_file": + prepared = self._prepare_write( + action.arguments["path"], action.arguments["content"] + ) + self._require_confirmation(action) + result = self._commit_write(prepared) + elif action.tool == "edit_file": + prepared = self._prepare_edit( + action.arguments["path"], + action.arguments["old"], + action.arguments["new"], + ) + self._require_confirmation(action) + result = self._commit_write(prepared) + elif action.tool == "run_command": + command = self._prepare_command(action.arguments["argv"]) + self._require_confirmation(action) + result = self._run_command(command) + else: + raise AgentError(f"unknown tool: {action.tool}") + except (KeyError, OSError, subprocess.SubprocessError, ValueError) as error: + result = f"error: {error}" + truncated = self._truncate_result(result) + self._record_receipt( + action, + truncated, + tuple( + sorted( + str(path.relative_to(self.policy.root)) + for path in (self.modified_files | self.uncertain_modified_files) + - (before_modified | before_uncertain) + ) + ), + tool_call_id=tool_call_id, + ) + return truncated + + def _truncate_result( + self, + result: str, + *, + preserve_command_warnings: tuple[str, ...] = (), + ) -> str: + """Bound a tool result to the configured output limit.""" + + if len(result) <= self.policy.max_tool_output_chars: + return result + truncated = ( + result[: self.policy.max_tool_output_chars] + "\n[... truncated ...]" + ) + if preserve_command_warnings: + truncated += "\n" + "\n".join(preserve_command_warnings) + return truncated + + def _open_parent_directory(self, relative: Path) -> tuple[int, str]: + """Open the parent directory of a relative path without following links.""" + + flags = ( + os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0) + ) + directory = os.open(self.policy.root, flags) + try: + for part in relative.parts[:-1]: + if part in {"", ".", ".."}: + raise AgentError("path is unsafe") + if _protected_path_reason((part,)): + raise AgentError(_protected_path_reason((part,))) + try: + directory = os.open(part, flags, dir_fd=directory) + except OSError as error: + raise AgentError("path is unsafe") from error + return directory, relative.name + except BaseException: + os.close(directory) + raise + + def _read_regular_at( + self, parent: int, name: str, *, tool: str + ) -> tuple[bytes, os.stat_result] | None: + """Read one bounded regular file by parent dir fd, rejecting links.""" + + try: + descriptor = os.open( + name, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + dir_fd=parent, + ) + except FileNotFoundError: + return None + except OSError as error: + raise AgentError(f"{tool} path is unsafe") from error + try: + status = os.fstat(descriptor) + if not stat.S_ISREG(status.st_mode) or status.st_nlink != 1: + raise AgentError(f"{tool} path is not a regular file") + if status.st_size > self.policy.max_file_bytes: + raise AgentError( + f"{tool} file exceeds {self.policy.max_file_bytes} bytes" + ) + data = b"" + while True: + chunk = os.read(descriptor, 1 << 16) + if not chunk: + break + data += chunk + if len(data) > self.policy.max_file_bytes: + raise AgentError( + f"{tool} file exceeds {self.policy.max_file_bytes} bytes" + ) + return data, status + finally: + os.close(descriptor) + + def _revalidate_prepared_write(self, prepared: _PreparedWrite) -> None: + """Reject a write whose target changed after preparation.""" + + path = prepared.path + relative = path.relative_to(self.policy.root) + parent, name = self._open_parent_directory(relative) + try: + current = self._read_regular_at(parent, name, tool="write_file") + finally: + os.close(parent) + if current is None: + if prepared.expected_digest is not None: + raise AgentError("file changed since it was read; read it again") + return + current_content, current_status = current + if prepared.expected_digest is None: + raise AgentError("file changed since it was read; read it again") + if self._digest(current_content) != prepared.expected_digest: + raise AgentError("file changed since it was read; read it again") + if stat.S_IMODE(current_status.st_mode) & 0o777 != prepared.expected_mode: + raise AgentError("file mode changed since it was read; read it again") + + def _read_bounded_file(self, path: Path, *, tool: str) -> bytes: + """Read one bounded regular file, rejecting symlinks and special files.""" + + try: + descriptor = os.open( + path, + os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), + ) + except OSError as error: + raise AgentError(f"{tool} path is unsafe or missing") from error + try: + status = os.fstat(descriptor) + if not stat.S_ISREG(status.st_mode) or status.st_nlink != 1: + raise AgentError(f"{tool} path is not a regular file") + if status.st_size > self.policy.max_file_bytes: + raise AgentError( + f"{tool} file exceeds {self.policy.max_file_bytes} bytes" + ) + data = b"" + while True: + chunk = os.read(descriptor, 1 << 16) + if not chunk: + break + data += chunk + if len(data) > self.policy.max_file_bytes: + raise AgentError( + f"{tool} file exceeds {self.policy.max_file_bytes} bytes" + ) + return data + finally: + os.close(descriptor) + + @staticmethod + def _digest(content: bytes) -> bytes: + """Fingerprint observed bytes for stale-write detection.""" + + return hashlib.sha256(content).digest() + + def _atomic_write( + self, + path: Path, + content: bytes, + *, + expected_digest: bytes | None, + expected_mode: int | None, + after_mode: int, + parent_identity: tuple[int, int], + ) -> None: + """Atomically replace a file inside the workspace, fsyncing parents.""" + + relative = path.relative_to(self.policy.root) + parent, name = self._open_parent_directory(relative) + temporary = None + try: + status = os.fstat(parent) + if (status.st_dev, status.st_ino) != parent_identity: + raise AgentError("parent directory changed during the write") + if expected_digest is not None: + current = self._read_regular_at(parent, name, tool="write_file") + if current is None: + raise AgentError("file changed since it was read; read it again") + current_content, current_status = current + if self._digest(current_content) != expected_digest: + raise AgentError("file changed since it was read; read it again") + if stat.S_IMODE(current_status.st_mode) & 0o777 != expected_mode: + raise AgentError( + "file mode changed since it was read; read it again" + ) + temporary = f".{name}.{uuid.uuid4().hex}.tmp" + descriptor = os.open( + temporary, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), + 0o600, + dir_fd=parent, + ) + try: + view = memoryview(content) + while view: + written = os.write(descriptor, view) + view = view[written:] + os.fchmod(descriptor, after_mode) + os.fsync(descriptor) + finally: + os.close(descriptor) + os.rename(temporary, name, src_dir_fd=parent, dst_dir_fd=parent) + os.fsync(parent) + finally: + if temporary is not None: + try: + os.unlink(temporary, dir_fd=parent) + except FileNotFoundError: + pass + os.close(parent) diff --git a/tests_refsol/test_week_4_day_3.py b/tests_refsol/test_week_4_day_3.py new file mode 100644 index 00000000..4209d5c5 --- /dev/null +++ b/tests_refsol/test_week_4_day_3.py @@ -0,0 +1,204 @@ +# WARNING: Under review - generated by LLM. + + +import pytest + +from .tiny_llm_base import AgentError, ToolAction, ToolPolicy, Workspace + + +def test_task_1_lists_and_reads_bounded_workspace_files(tmp_path): + (tmp_path / "README.md").write_text("hello", encoding="utf-8") + (tmp_path / ".env").write_text("SECRET=value", encoding="utf-8") + workspace = Workspace(ToolPolicy(tmp_path)) + + listing = workspace.list_files() + assert "file README.md" in listing + assert ".env" not in listing + assert workspace.read_file("README.md") == "hello" + + +def test_task_1_list_is_sorted_and_bounded(tmp_path): + for name in ("c.txt", "a.txt", "b.txt"): + (tmp_path / name).write_text(name, encoding="utf-8") + workspace = Workspace(ToolPolicy(tmp_path, max_list_entries=2)) + + assert workspace.list_files().splitlines() == ["file a.txt", "file b.txt"] + + +def test_task_1_rejects_oversized_non_utf8_and_non_file_reads(tmp_path): + (tmp_path / "large.txt").write_text("12345", encoding="utf-8") + (tmp_path / "binary.dat").write_bytes(b"\xff") + (tmp_path / "directory").mkdir() + workspace = Workspace(ToolPolicy(tmp_path, max_file_bytes=4)) + + with pytest.raises(AgentError, match="exceeds 4 bytes"): + workspace.read_file("large.txt") + assert workspace.execute( + ToolAction("read_file", {"path": "binary.dat"}) + ).startswith("error:") + + +@pytest.mark.parametrize( + "path", + [ + "../outside.txt", + "/etc/passwd", + "a/../../outside.txt", + ".git/config", + "secret.pem", + ".env.local", + ], +) +def test_task_2_rejects_unsafe_paths(tmp_path, path): + (tmp_path / "outside.txt").write_text("x", encoding="utf-8") + workspace = Workspace(ToolPolicy(tmp_path)) + + with pytest.raises(AgentError): + workspace.read_file(path) + + +def test_task_2_rejects_symlinks(tmp_path): + (tmp_path / "real.txt").write_text("real", encoding="utf-8") + (tmp_path / "link.txt").symlink_to(tmp_path / "real.txt") + workspace = Workspace(ToolPolicy(tmp_path)) + + with pytest.raises(AgentError): + workspace.read_file("link.txt") + result = workspace.execute(ToolAction("read_file", {"path": "link.txt"})) + assert result.startswith("error:") + + +def test_task_3_writes_are_disabled_by_default(tmp_path): + workspace = Workspace(ToolPolicy(tmp_path)) + + with pytest.raises(AgentError, match="disabled"): + workspace.write_file("new.txt", "data") + assert not (tmp_path / "new.txt").exists() + + +def test_task_3_creates_and_atomically_replaces_files(tmp_path): + workspace = Workspace(ToolPolicy(tmp_path, allow_writes=True)) + + assert workspace.write_file("new.txt", "one") == "wrote new.txt" + assert (tmp_path / "new.txt").read_text(encoding="utf-8") == "one" + assert workspace.write_file("new.txt", "two") == "wrote new.txt" + assert (tmp_path / "new.txt").read_text(encoding="utf-8") == "two" + + +def test_task_3_requires_inspection_before_overwriting(tmp_path): + (tmp_path / "existing.txt").write_text("old", encoding="utf-8") + workspace = Workspace(ToolPolicy(tmp_path, allow_writes=True)) + + with pytest.raises(AgentError, match="must be read"): + workspace.write_file("existing.txt", "new") + assert (tmp_path / "existing.txt").read_text(encoding="utf-8") == "old" + + +def test_task_3_rejects_a_stale_observation_before_confirmation(tmp_path): + path = tmp_path / "answer.py" + path.write_text("ANSWER = 41\n", encoding="utf-8") + workspace = Workspace(ToolPolicy(tmp_path, allow_writes=True)) + workspace.read_file("answer.py") + path.write_text("ANSWER = 42\n", encoding="utf-8") + + with pytest.raises(AgentError, match="changed since"): + workspace.write_file("answer.py", "ANSWER = 43\n") + + +def test_task_3_edit_requires_one_exact_match(tmp_path): + path = tmp_path / "example.py" + path.write_text("answer = 41\n", encoding="utf-8") + workspace = Workspace(ToolPolicy(tmp_path, allow_writes=True)) + workspace.read_file("example.py") + + workspace.edit_file("example.py", "41", "42") + + assert path.read_text(encoding="utf-8") == "answer = 42\n" + assert workspace.modified_files == {path} + + +def test_task_3_model_dispatched_writes_are_denied_without_confirmation(tmp_path): + workspace = Workspace(ToolPolicy(tmp_path, allow_writes=True)) + + result = workspace.execute( + ToolAction("write_file", {"path": "new.txt", "content": "data"}) + ) + + assert result == "error: operator denied write_file" + assert not (tmp_path / "new.txt").exists() + + +def test_task_3_model_dispatched_mutations_prompt_exactly_once(tmp_path): + path = tmp_path / "answer.py" + path.write_text("ANSWER = 41\n", encoding="utf-8") + confirmations = [] + workspace = Workspace( + ToolPolicy(tmp_path, allow_writes=True), + confirm_tool=lambda action: confirmations.append(action.tool) or True, + ) + workspace.read_file("answer.py") + + result = workspace.execute( + ToolAction("edit_file", {"path": "answer.py", "old": "41", "new": "42"}) + ) + + assert result == "wrote answer.py" + assert confirmations == ["edit_file"] + assert path.read_text(encoding="utf-8") == "ANSWER = 42\n" + + +def test_task_3_reads_never_request_confirmation(tmp_path): + (tmp_path / "a.txt").write_text("a", encoding="utf-8") + confirmations = [] + workspace = Workspace( + ToolPolicy(tmp_path), confirm_tool=lambda action: confirmations.append(action) + ) + + workspace.execute(ToolAction("read_file", {"path": "a.txt"})) + + assert confirmations == [] + + +def test_task_3_commands_require_an_exact_operator_allowlist(tmp_path): + workspace = Workspace(ToolPolicy(tmp_path, allowed_commands=(("echo", "ok"),))) + + with pytest.raises(AgentError, match="not explicitly allowed"): + workspace.run_command(["some-command"]) + + +def test_task_3_unallowlisted_command_is_rejected_before_confirmation(tmp_path): + confirmations = [] + workspace = Workspace( + ToolPolicy(tmp_path, allowed_commands=(("echo", "ok"),)), + confirm_tool=lambda action: confirmations.append(action) or True, + ) + + result = workspace.execute(ToolAction("run_command", {"argv": ["rm", "-rf", "."]})) + + assert result.startswith("error:") + assert confirmations == [] + + +def test_task_3_denied_command_never_starts_a_process(tmp_path): + confirmations = [] + workspace = Workspace( + ToolPolicy(tmp_path, allowed_commands=(("echo", "ok"),)), + confirm_tool=lambda action: confirmations.append(action.tool) or False, + ) + + result = workspace.execute(ToolAction("run_command", {"argv": ["echo", "ok"]})) + + assert result == "error: operator denied run_command" + assert confirmations == ["run_command"] + + +def test_task_3_allowed_command_runs_with_confirmation(tmp_path): + workspace = Workspace( + ToolPolicy(tmp_path, allowed_commands=(("echo", "ok"),)), + confirm_tool=lambda action: True, + ) + + result = workspace.execute(ToolAction("run_command", {"argv": ["echo", "ok"]})) + + assert "ok" in result + assert not result.startswith("error:") diff --git a/tests_refsol/test_week_4_day_4.py b/tests_refsol/test_week_4_day_4.py new file mode 100644 index 00000000..99eac157 --- /dev/null +++ b/tests_refsol/test_week_4_day_4.py @@ -0,0 +1,147 @@ +# WARNING: Under review - generated by LLM. + +import pytest + +from .tiny_llm_base import ( + SessionStore, + ToolPolicy, + Workspace, + run_agent, +) + + +def responses(*items): + queue = iter(items) + return lambda messages: next(queue) + + +def test_task_1_session_events_round_trip_and_persist(tmp_path): + store = SessionStore(tmp_path, "test-model") + session = store.create() + call = session.append("tool_call", tool="read_file", arguments={"path": "a"}) + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content="data", + ) + + reloaded = store.load(session.session_id, recover=False) + assert [event.type for event in reloaded.events] == [ + "session_started", + "tool_call", + "tool_result", + ] + + +def test_task_2_loop_records_durable_events(tmp_path): + store = SessionStore(tmp_path, "test-model") + session = store.create() + workspace = Workspace(ToolPolicy(tmp_path)) + + result = run_agent( + "inspect the project", + responses('{"tool":"list_files"}', '{"final":"done"}'), + workspace, + session=session, + ) + + assert result.completed + assert result.session_id == session.session_id + types = [event.type for event in session.events] + assert "user_message" in types + assert "run_started" in types + assert "tool_call" in types + assert "tool_result" in types + assert "run_finished" in types + + +def test_task_2_invalid_json_is_recorded_without_execution(tmp_path): + store = SessionStore(tmp_path, "test-model") + session = store.create() + workspace = Workspace(ToolPolicy(tmp_path)) + + result = run_agent( + "inspect the project", + responses("not json", '{"final":"recovered"}'), + workspace, + session=session, + ) + + assert result.completed + tool_results = [event for event in session.events if event.type == "tool_result"] + assert len(tool_results) == 1 + assert tool_results[0].data["is_error"] is True + assert tool_results[0].data["tool_call_id"] is None + + +def test_task_3_active_path_of_root_session_is_just_the_root(tmp_path): + store = SessionStore(tmp_path, "test-model") + root = store.create() + + assert store.active_path(root.session_id) == (root.session_id,) + assert store.parent_of(root.session_id) is None + + +def test_task_3_branch_builds_a_two_level_active_path(tmp_path): + store = SessionStore(tmp_path, "test-model") + root = store.create() + root.append("user_message", content="task one") + branch = store.branch(root.session_id, at_event_id=root.events[-1].id) + + path = store.active_path(branch.session_id) + assert path == (root.session_id, branch.session_id) + assert store.parent_of(branch.session_id) == ( + root.session_id, + root.events[-1].id, + ) + + +def test_task_3_active_path_events_are_deterministic(tmp_path): + store = SessionStore(tmp_path, "test-model") + root = store.create() + root.append("user_message", content="task one") + branch = store.branch(root.session_id, at_event_id=root.events[-1].id) + branch.append("user_message", content="branch follow-up") + + first = branch.active_path_events() + second = branch.active_path_events() + + assert [event.id for event in first] == [event.id for event in second] + assert first[-1].data.get("content") == "branch follow-up" + + +def test_task_3_branch_keeps_the_full_inherited_prefix(tmp_path): + store = SessionStore(tmp_path, "test-model") + root = store.create() + root.append("user_message", content="task one") + branch = store.branch(root.session_id, at_event_id=root.events[-1].id) + + events = branch.active_path_events() + types = [event.type for event in events] + assert "branch_created" in types + assert "branch_completed" in types + assert types.count("user_message") == 1 + + +def test_task_3_grandchild_path_orders_root_first(tmp_path): + store = SessionStore(tmp_path, "test-model") + root = store.create() + root.append("user_message", content="task one") + branch = store.branch(root.session_id, at_event_id=root.events[-1].id) + branch.append("user_message", content="branch follow-up") + child = store.branch(branch.session_id, at_event_id=branch.events[-1].id) + + assert store.active_path(child.session_id) == ( + root.session_id, + branch.session_id, + child.session_id, + ) + + +def test_task_3_active_path_rejects_an_unknown_session(tmp_path): + store = SessionStore(tmp_path, "test-model") + + with pytest.raises(ValueError): + store.active_path("0" * 32) diff --git a/tests_refsol/test_week_4_day_5.py b/tests_refsol/test_week_4_day_5.py new file mode 100644 index 00000000..55145e77 --- /dev/null +++ b/tests_refsol/test_week_4_day_5.py @@ -0,0 +1,296 @@ +# WARNING: Under review - generated by LLM. + +import sys +from types import ModuleType + +import pytest + +from .tiny_llm_base import ( + CacheManifest, + GenerationSession, + ManifestError, + export_cache_manifest, + validate_resume, +) + + +class _FakeCache: + def __init__(self): + self.offset = 0 + self.released = 0 + + def rewind(self, count): + if count < 0 or count > self.offset: + raise ValueError("bad rewind") + self.offset -= count + + def release(self): + self.released += 1 + + +class _FakeArray: + def __init__(self, values): + self.values = list(values) + + def __getitem__(self, _key): + return self + + +class _FakeLogits: + def __init__(self, token): + self.token = token + + def __getitem__(self, _key): + return self + + +class _FakeScalar: + def __init__(self, value): + self.value = value + + def item(self): + return self.value + + +class _FakeTokenizer: + eos_token_id = 0 + + def apply_chat_template(self, messages, **_kwargs): + return "|".join( + f"{message['role']}:{message['content']}" for message in messages + ) + + def encode(self, prompt, **_kwargs): + return [(ord(character) % 251) + 1 for character in prompt] + + def decode(self, tokens): + return ",".join(map(str, tokens)) + + +class _FakeModel: + def __call__(self, tokens, _offset, caches): + for cache in caches: + cache.offset += len(tokens.values) + return _FakeLogits(0 if tokens.values[-1] == 7 else 7) + + +def _install_fake_mlx(monkeypatch): + mlx = ModuleType("mlx") + core = ModuleType("mlx.core") + core.array = _FakeArray + core.argmax = lambda logits, axis: _FakeScalar(logits.token) + mlx.core = core + monkeypatch.setitem(sys.modules, "mlx", mlx) + monkeypatch.setitem(sys.modules, "mlx.core", core) + + +def _make_session( + monkeypatch, *, layers=2, max_tokens=8, model_hash=None, tokenizer_hash=None +): + _install_fake_mlx(monkeypatch) + model = _FakeModel() + tokenizer = _FakeTokenizer() + + def cache_factory(): + return [_FakeCache() for _ in range(layers)] + + return GenerationSession( + model, + tokenizer, + cache_factory, + max_tokens=max_tokens, + model_hash=model_hash, + tokenizer_hash=tokenizer_hash, + ) + + +def test_task_1_checkpoint_exports_content_addressed_manifest(monkeypatch): + session = _make_session(monkeypatch) + session([{"role": "user", "content": "hello"}]) + + manifest = export_cache_manifest( + session, + session.cached_token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + assert manifest.position == len(session.cached_token_ids) + assert manifest.layers == 2 + assert len(manifest.digest()) == 64 + + +def test_task_1_manifest_round_trips_and_verifies_its_digest(monkeypatch): + session = _make_session(monkeypatch) + session([{"role": "user", "content": "hello"}]) + + manifest = export_cache_manifest( + session, + session.cached_token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + restored = CacheManifest.from_dict(manifest.to_dict()) + + assert restored == manifest + assert restored.digest() == manifest.digest() + + +def test_task_1_tampered_manifest_fails_closed(): + manifest = CacheManifest( + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + position=10, + prefix_hash="c" * 64, + ) + payload = manifest.to_dict() + payload["position"] = 999 + + with pytest.raises(ManifestError): + CacheManifest.from_dict(payload) + + +def test_task_1_resume_accepts_an_exact_manifest(monkeypatch): + session = _make_session(monkeypatch) + session([{"role": "user", "content": "hello"}]) + token_ids = session.cached_token_ids + manifest = export_cache_manifest( + session, + token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + validate_resume( + manifest, + session, + token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + +def test_task_1_resume_rejects_a_changed_workspace(monkeypatch): + session = _make_session(monkeypatch) + session([{"role": "user", "content": "hello"}]) + token_ids = session.cached_token_ids + manifest = export_cache_manifest( + session, + token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + with pytest.raises(ManifestError, match="workspace"): + validate_resume( + manifest, + session, + token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="x" * 64, + ) + + +def test_task_1_resume_rejects_a_changed_tool_catalog(monkeypatch): + session = _make_session(monkeypatch) + session([{"role": "user", "content": "hello"}]) + token_ids = session.cached_token_ids + manifest = export_cache_manifest( + session, + token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + with pytest.raises(ManifestError, match="tool catalog"): + validate_resume( + manifest, + session, + token_ids, + tool_catalog_hash="u" * 64, + workspace_fingerprint="w" * 64, + ) + + +def test_task_1_resume_rejects_a_different_prefix(monkeypatch): + session = _make_session(monkeypatch) + session([{"role": "user", "content": "hello"}]) + token_ids = session.cached_token_ids + manifest = export_cache_manifest( + session, + token_ids, + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + with pytest.raises(ManifestError, match="position|prefix"): + validate_resume( + manifest, + session, + token_ids[:-1], + tool_catalog_hash="t" * 64, + workspace_fingerprint="w" * 64, + ) + + +def _make_branch_session(monkeypatch, layers=2, max_tokens=8): + return _make_session(monkeypatch, layers=layers, max_tokens=max_tokens) + + +def test_task_2_sequential_branch_checkpoint_act_rewind_do_again(monkeypatch): + from .tiny_llm_base import SequentialBranch + + session = _make_branch_session(monkeypatch) + branch = SequentialBranch(session) + branch.session([{"role": "user", "content": "setup"}]) + boundary = branch.checkpoint() + + first, first_stats = branch.act([{"role": "user", "content": "try A"}]) + assert first_stats.checkpoint_tokens == len(boundary) + + rewound = branch.rewind() + assert rewound == len(boundary) + + second, second_stats = branch.do_again([{"role": "user", "content": "try B"}]) + assert second_stats.checkpoint_tokens == len(boundary) + branch.close() + + +def test_task_2_branch_requires_checkpoint_before_act(monkeypatch): + from .tiny_llm_base import RewindError, SequentialBranch + + branch = SequentialBranch(_make_branch_session(monkeypatch)) + + with pytest.raises(RewindError, match="checkpoint"): + branch.act([{"role": "user", "content": "try A"}]) + with pytest.raises(RewindError, match="checkpoint"): + branch.rewind() + branch.close() + + +def test_task_2_rewind_never_exceeds_the_checkpoint(monkeypatch): + from .tiny_llm_base import SequentialBranch + + session = _make_branch_session(monkeypatch) + branch = SequentialBranch(session) + branch.session([{"role": "user", "content": "setup"}]) + branch.checkpoint() + branch.rewind() + + assert len(branch.session.cached_token_ids) == len(branch.checkpoint_tokens) + branch.close() + + +def test_task_2_closed_branch_rejects_all_operations(monkeypatch): + from .tiny_llm_base import RewindError, SequentialBranch + + branch = SequentialBranch(_make_branch_session(monkeypatch)) + branch.close() + + with pytest.raises(RewindError, match="closed"): + branch.checkpoint() + with pytest.raises(RewindError, match="closed"): + branch.rewind() + with pytest.raises(RewindError, match="closed"): + branch.act([{"role": "user", "content": "try A"}]) diff --git a/tests_refsol/test_week_4_day_6.py b/tests_refsol/test_week_4_day_6.py new file mode 100644 index 00000000..25927d9d --- /dev/null +++ b/tests_refsol/test_week_4_day_6.py @@ -0,0 +1,117 @@ +# WARNING: Under review - generated by LLM. + +import pytest + +from .tiny_llm_base import ( + CompactionError, + EffectReceipt, + ReceiptStore, + SessionStore, + compact_tool_results, + expand_receipt_range, + reexpand_receipt_message, +) + + +def _session_with_big_result(tmp_path, *, words=500): + store = SessionStore(tmp_path, "test-model") + session = store.create() + session.append("user_message", content="inspect the log") + call = session.append("tool_call", tool="read_file", arguments={"path": "log"}) + body = "line " * words + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content=body, + ) + receipt = EffectReceipt( + tool_call_id=call.id, + tool="read_file", + arguments={"path": "log"}, + exit_state="ok", + result=body, + changed_artifacts=(), + ) + return session, call, receipt + + +def test_task_1_compact_replaces_oversized_tool_results(tmp_path): + session, call, receipt = _session_with_big_result(tmp_path) + store = ReceiptStore(tmp_path / "receipts.jsonl") + store.put(receipt) + + result = compact_tool_results(session, store, max_result_bytes=100) + + assert result.compacted == 1 + assert result.saved_bytes > 0 + assert "compacted" in result.messages[0]["content"] + + +def test_task_1_compact_never_mutates_the_durable_trace(tmp_path): + session, call, receipt = _session_with_big_result(tmp_path) + store = ReceiptStore(tmp_path / "receipts.jsonl") + store.put(receipt) + before = [event.to_dict() for event in session.events] + + compact_tool_results(session, store, max_result_bytes=100) + + assert [event.to_dict() for event in session.events] == before + + +def test_task_1_compact_leaves_results_without_receipts_untouched(tmp_path): + session, call, receipt = _session_with_big_result(tmp_path) + store = ReceiptStore(tmp_path / "receipts.jsonl") + + result = compact_tool_results(session, store, max_result_bytes=100) + + assert result.compacted == 0 + assert "compacted" not in result.messages[0]["content"] + + +def test_task_2_compact_then_expand_recovers_the_omitted_middle(tmp_path): + store = SessionStore(tmp_path, "test-model") + session = store.create() + session.append("user_message", content="inspect the log") + call = session.append("tool_call", tool="read_file", arguments={"path": "log"}) + body = "\n".join(f"fact {number}" for number in range(1, 401)) + "\n" + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content=body, + ) + receipt_store = ReceiptStore(tmp_path / "receipts.jsonl") + receipt = EffectReceipt( + tool_call_id=call.id, + tool="read_file", + arguments={"path": "log"}, + exit_state="ok", + result=body, + changed_artifacts=(), + ) + receipt_store.put(receipt) + + result = compact_tool_results(session, receipt_store, max_result_bytes=100) + rendering = result.messages[0]["content"] + assert result.compacted == 1 + + assert ( + reexpand_receipt_message(receipt_store, rendering, start=0, end=9) == body[:9] + ) + assert reexpand_receipt_message(receipt_store, rendering) == body + assert ( + expand_receipt_range(receipt_store, receipt.receipt_id, start=len(body) - 10) + == body[-10:] + ) + + +def test_task_2_reexpansion_rejects_an_invalid_handle(tmp_path): + store = ReceiptStore(tmp_path / "receipts.jsonl") + + with pytest.raises(CompactionError, match="does not reference"): + reexpand_receipt_message(store, "no handle here") + with pytest.raises(CompactionError, match="invalid receipt ID"): + reexpand_receipt_message(store, "expand via receipt not-a-hash") diff --git a/tests_refsol/test_week_4_day_7.py b/tests_refsol/test_week_4_day_7.py new file mode 100644 index 00000000..795cf2d3 --- /dev/null +++ b/tests_refsol/test_week_4_day_7.py @@ -0,0 +1,92 @@ +# WARNING: Under review - generated by LLM. + +from .tiny_llm_base import ( + build_state_card, + memory_session, +) + + +def test_task_1_state_card_derives_only_public_ledger_state(tmp_path): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="fix the parser") + session.append("tool_call", tool="read_file", arguments={"path": "a.py"}) + + card = build_state_card(session, goal="fix the parser") + + assert card.goal == "fix the parser" + assert "read_file" in card.current_action + assert card.phase == "running" + + +def test_task_1_state_card_tracks_tool_evidence(tmp_path): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="fix the parser") + call = session.append("tool_call", tool="read_file", arguments={"path": "a.py"}) + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content="line 1: def parse()", + ) + + card = build_state_card(session, goal="fix the parser") + + assert "line 1: def parse()" in card.last_evidence + assert card.next_safe_step == "continue the loop" + + +def test_task_1_state_card_finished_phase_is_terminal(tmp_path): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="fix the parser") + session.append("run_finished", completed=True, reason="completed", final="done") + + card = build_state_card(session, goal="fix the parser") + + assert card.phase == "finished" + assert card.current_action == "none" + assert card.next_safe_step == "none" + + +def test_task_1_state_card_never_exposes_hidden_reasoning(tmp_path): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="fix the parser") + session.append("assistant_message", content="I will now reason internally") + + card = build_state_card(session, goal="fix the parser") + + assert "internally" not in card.render() + assert card.current_action == "generating" + + +def test_task_2_steering_handle_queues_durably(tmp_path): + from .tiny_llm_base import SteeringHandle + + session = memory_session(tmp_path, "test-model") + handle = SteeringHandle(session) + event = handle.submit("focus on the parser") + + assert event.type == "steering_queued" + assert session.pending_steering()[0].id == event.id + delivered = session.deliver_pending_steering() + assert delivered[0].data["content"] == "focus on the parser" + assert session.pending_steering() == () + + +def test_task_3_cancellation_is_first_writer_wins(): + import pytest + + from .tiny_llm_base import AgentInterrupted, CancellationToken + + cancellation = CancellationToken() + + assert cancellation.cancel("operator interrupt") is True + assert cancellation.cancel("later timeout") is False + assert cancellation.cancelled is True + assert cancellation.reason == "operator interrupt" + + with pytest.raises(AgentInterrupted) as caught: + cancellation.raise_if_cancelled("model") + + assert caught.value.reason == "operator interrupt" + assert caught.value.phase == "model" diff --git a/tests_refsol/test_week_4_day_8.py b/tests_refsol/test_week_4_day_8.py new file mode 100644 index 00000000..2fe3230d --- /dev/null +++ b/tests_refsol/test_week_4_day_8.py @@ -0,0 +1,168 @@ +# WARNING: Under review - generated by LLM. + +from .tiny_llm_base import ( + CacheManifest, + EffectReceipt, + ReceiptStore, + largest_safe_checkpoint, + memory_session, + reconcile_effect, + reconcile_interrupted_effects, +) + + +def test_task_1_reconcile_appends_exactly_one_observation(tmp_path): + store = ReceiptStore(tmp_path / "receipts.jsonl") + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="edit the file") + call = session.append( + "tool_call", tool="write_file", arguments={"path": "a.txt", "content": "x"} + ) + store.put( + EffectReceipt( + tool_call_id=call.id, + tool="write_file", + arguments={"path": "a.txt", "content": "x"}, + exit_state="ok", + result="wrote a.txt", + changed_artifacts=("a.txt",), + ) + ) + + first = reconcile_effect(session, store, call.id) + second = reconcile_effect(session, store, call.id) + + assert first.status == "observation_appended" + assert second.status == "already_closed" + observations = [ + event + for event in session.events + if event.type == "tool_result" and event.data.get("tool_call_id") == call.id + ] + assert len(observations) == 1 + assert observations[0].data.get("content") == "wrote a.txt" + + +def test_task_1_reconcile_never_guesses_an_orphaned_effect(tmp_path): + store = ReceiptStore(tmp_path / "receipts.jsonl") + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="edit the file") + call = session.append( + "tool_call", tool="write_file", arguments={"path": "a.txt", "content": "x"} + ) + + result = reconcile_effect(session, store, call.id) + + assert result.status == "no_receipt" + assert not any(event.type == "tool_result" for event in session.events) + + +def test_task_1_reconcile_pass_closes_every_interrupted_call(tmp_path): + store = ReceiptStore(tmp_path / "receipts.jsonl") + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="do work") + first = session.append("tool_call", tool="read_file", arguments={"path": "a.py"}) + second = session.append( + "tool_call", tool="write_file", arguments={"path": "b.py", "content": "y"} + ) + store.put( + EffectReceipt( + tool_call_id=first.id, + tool="read_file", + arguments={"path": "a.py"}, + exit_state="ok", + result="def f(): pass", + changed_artifacts=(), + ) + ) + store.put( + EffectReceipt( + tool_call_id=second.id, + tool="write_file", + arguments={"path": "b.py", "content": "y"}, + exit_state="ok", + result="wrote b.py", + changed_artifacts=("b.py",), + ) + ) + + results = reconcile_interrupted_effects(session, store) + + assert [result.status for result in results] == [ + "observation_appended", + "observation_appended", + ] + + +def test_task_2_largest_safe_checkpoint_picks_the_most_advanced(): + smaller = CacheManifest( + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + position=10, + prefix_hash="c" * 64, + ) + larger = CacheManifest( + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + position=40, + prefix_hash="c" * 64, + ) + + result = largest_safe_checkpoint( + (smaller, larger), + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + tool_catalog_hash="", + workspace_fingerprint="", + ) + + assert result.can_resume + assert result.manifest is larger + + +def test_task_2_stale_checkpoint_falls_back_cold(): + stale = CacheManifest( + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + position=40, + prefix_hash="c" * 64, + workspace_fingerprint="w" * 64, + ) + + result = largest_safe_checkpoint( + (stale,), + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + tool_catalog_hash="", + workspace_fingerprint="x" * 64, + ) + + assert not result.can_resume + assert result.reason.startswith("cold_") + + +def test_task_2_model_mismatch_rejects_every_checkpoint(): + other_model = CacheManifest( + model_hash="d" * 64, + tokenizer_hash="b" * 64, + layers=2, + position=40, + prefix_hash="c" * 64, + ) + + result = largest_safe_checkpoint( + (other_model,), + model_hash="a" * 64, + tokenizer_hash="b" * 64, + layers=2, + tool_catalog_hash="", + workspace_fingerprint="", + ) + + assert not result.can_resume + assert "model" in result.reason diff --git a/tests_refsol/test_week_4_day_9.py b/tests_refsol/test_week_4_day_9.py new file mode 100644 index 00000000..077378ab --- /dev/null +++ b/tests_refsol/test_week_4_day_9.py @@ -0,0 +1,162 @@ +# WARNING: Under review - generated by LLM. + +from .tiny_llm_base import ( + EffectReceipt, + compare_runs, + memory_session, + snapshot_run, +) + + +def _receipt(tool_call_id="a" * 32): + return EffectReceipt( + tool_call_id=tool_call_id, + tool="read_file", + arguments={"path": "a"}, + exit_state="ok", + result="data", + changed_artifacts=(), + ) + + +def test_task_1_warm_and_cold_runs_are_equivalent(tmp_path): + def make(): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="inspect") + call = session.append("tool_call", tool="read_file", arguments={"path": "a"}) + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content="data", + ) + return session + + warm = snapshot_run( + final="done", + reason="completed", + completed=True, + session=make(), + token_accounting={"reused_tokens": 100, "prefilled_tokens": 5}, + ) + cold = snapshot_run( + final="done", + reason="completed", + completed=True, + session=make(), + token_accounting={"reused_tokens": 0, "prefilled_tokens": 105}, + ) + + report = compare_runs(cold, warm) + + assert report.ok + assert ( + warm.token_accounting["reused_tokens"] > cold.token_accounting["reused_tokens"] + ) + + +def test_task_1_different_final_actions_fail_semantic_equivalence(tmp_path): + session_a = memory_session(tmp_path, "test-model") + session_a.append("user_message", content="task") + session_b = memory_session(tmp_path, "test-model") + session_b.append("user_message", content="task") + + baseline = snapshot_run( + final="keep", reason="completed", completed=True, session=session_a + ) + optimized = snapshot_run( + final="replace", reason="completed", completed=True, session=session_b + ) + + report = compare_runs(baseline, optimized) + + assert not report.ok + assert not report.planes[0].ok # semantic + assert report.planes[1].ok # evidence still matches + + +def test_task_1_lost_receipt_fails_evidence_equivalence(): + baseline = snapshot_run( + final="done", + reason="completed", + completed=True, + receipts=(_receipt(),), + ) + optimized = snapshot_run( + final="done", reason="completed", completed=True, receipts=() + ) + + report = compare_runs(baseline, optimized) + + assert not report.ok + assert report.planes[1].plane == "evidence" + assert not report.planes[1].ok + + +def test_task_1_policy_state_change_fails_policy_equivalence(): + baseline = snapshot_run( + final="done", reason="completed", completed=True, approval_state="epoch 1" + ) + optimized = snapshot_run( + final="done", reason="completed", completed=True, approval_state="epoch 2" + ) + + report = compare_runs(baseline, optimized) + + assert not report.ok + assert report.planes[2].plane == "policy" + assert not report.planes[2].ok + + +def test_task_2_compacted_and_full_runs_keep_equivalent_evidence(tmp_path): + def make(compact): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="task") + call = session.append("tool_call", tool="read_file", arguments={"path": "a"}) + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content="head [...] tail" if compact else "head middle tail", + ) + return session + + full = snapshot_run( + final="done", reason="completed", completed=True, session=make(False) + ) + compacted = snapshot_run( + final="done", reason="completed", completed=True, session=make(True) + ) + + report = compare_runs(full, compacted) + + # Different visible observation text is a real evidence difference; the + # harness must not pass it silently. + assert not report.planes[1].ok + + +def test_task_2_crash_resume_preserves_evidence(tmp_path): + def make(): + session = memory_session(tmp_path, "test-model") + session.append("user_message", content="task") + call = session.append("tool_call", tool="read_file", arguments={"path": "a"}) + session.append( + "tool_result", + tool_call_id=call.id, + tool="read_file", + is_error=False, + content="data", + reconciled=True, + ) + return session + + baseline = snapshot_run( + final="done", reason="completed", completed=True, session=make() + ) + resumed = snapshot_run( + final="done", reason="completed", completed=True, session=make() + ) + + assert compare_runs(baseline, resumed).ok