diff --git a/migrations/versions/064_turn_slices.py b/migrations/versions/064_turn_slices.py new file mode 100644 index 000000000..eb1386434 --- /dev/null +++ b/migrations/versions/064_turn_slices.py @@ -0,0 +1,73 @@ +"""Migration 064: per-turn and per-subagent-run slices of checkpoint history. + +A slice is stored once when a turn or run settles, so replay reads rows +instead of materializing checkpoint state. It is derived data: deleting rows +loses nothing, the next read rebuilds them. A slice is valid while +``slice_key`` names the code that cut it; ``lines`` is the projected replay +of the turn, wire-ready, valid while ``lines_key`` matches. ``run_claims`` +pins each subagent launch without a ledger stamp to the run it started, which +only a pass over the whole thread can work out. + +Rollback: ``downgrade`` drops both tables; the previous build never reads them. +""" + +from alembic import op + +revision = "064" +down_revision = "063" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Each foreign key takes a lock on the table it references, which live + # turns write; fail fast rather than queue every writer behind it. + op.execute("SET LOCAL lock_timeout = '5s'") + op.execute(""" + CREATE TABLE IF NOT EXISTS turn_slices ( + conversation_response_id UUID PRIMARY KEY + REFERENCES conversation_responses(conversation_response_id) + ON DELETE CASCADE, + conversation_thread_id UUID NOT NULL, + input_checkpoint_id TEXT NOT NULL, + tail_checkpoint_id TEXT NOT NULL, + slice_key TEXT NOT NULL, + slice_codec TEXT NOT NULL, + slice BYTEA NOT NULL, + lines_key TEXT, + lines TEXT, + run_claims JSONB, + built_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """) + # Every read by thread filters on the code that cut the slice too. + op.execute( + "CREATE INDEX IF NOT EXISTS idx_turn_slices_thread " + "ON turn_slices (conversation_thread_id, slice_key)" + ) + op.execute(""" + CREATE TABLE IF NOT EXISTS task_run_slices ( + conversation_thread_id UUID NOT NULL + REFERENCES conversation_threads(conversation_thread_id) + ON DELETE CASCADE, + task_id TEXT NOT NULL, + input_checkpoint_id TEXT NOT NULL, + task_run_id UUID REFERENCES subagent_runs(task_run_id) + ON DELETE CASCADE, + tail_checkpoint_id TEXT NOT NULL, + slice_key TEXT NOT NULL, + slice_codec TEXT NOT NULL, + slice BYTEA NOT NULL, + built_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (conversation_thread_id, task_id, input_checkpoint_id) + ) + """) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_task_run_slices_run " + "ON task_run_slices (task_run_id) WHERE task_run_id IS NOT NULL" + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS task_run_slices") + op.execute("DROP TABLE IF EXISTS turn_slices") diff --git a/migrations/versions/065_replay_facts.py b/migrations/versions/065_replay_facts.py new file mode 100644 index 000000000..df3437a57 --- /dev/null +++ b/migrations/versions/065_replay_facts.py @@ -0,0 +1,39 @@ +"""Migration 065: replay facts on settled runs. + +Some of what a turn's replay shows was only ever in the stream it emitted: +how long each reasoning block took, which sandbox images were captured to +which URLs, and the steering a subagent run handed back unread. ``replay_facts`` +keeps those few values on the run row itself, so replay reads them next to the +checkpoint instead of from the stored event log. NULL means a run settled +before the column existed, and replay keeps reading the stored events for it. + +Rollback: ``downgrade`` drops both columns; the previous build never reads them. +""" + +from alembic import op + +revision = "065" +down_revision = "064" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Each ALTER waits for an ACCESS EXCLUSIVE lock on a table live turns + # write; fail fast rather than queue every writer behind it. + op.execute("SET LOCAL lock_timeout = '5s'") + op.execute( + "ALTER TABLE conversation_responses " + "ADD COLUMN IF NOT EXISTS replay_facts JSONB" + ) + op.execute( + "ALTER TABLE subagent_runs ADD COLUMN IF NOT EXISTS replay_facts JSONB" + ) + + +def downgrade() -> None: + op.execute("SET LOCAL lock_timeout = '5s'") + op.execute("ALTER TABLE subagent_runs DROP COLUMN IF EXISTS replay_facts") + op.execute( + "ALTER TABLE conversation_responses DROP COLUMN IF EXISTS replay_facts" + ) diff --git a/scripts/_errors.py b/scripts/_errors.py new file mode 100644 index 000000000..c01f71bab --- /dev/null +++ b/scripts/_errors.py @@ -0,0 +1,14 @@ +"""How the operator scripts under ``scripts/`` name a failure in their output.""" + +from __future__ import annotations + +import traceback +from pathlib import Path + + +def where(exc: BaseException) -> str: + """An exception's type and innermost frame. Its message can quote message + content (a validation error echoes its input), so it is not printed.""" + frames = traceback.extract_tb(exc.__traceback__) + at = f" at {Path(frames[-1].filename).name}:{frames[-1].lineno}" if frames else "" + return f"{type(exc).__name__}{at}" diff --git a/scripts/ops/_db.py b/scripts/ops/_db.py index d590be5ed..08c72828c 100644 --- a/scripts/ops/_db.py +++ b/scripts/ops/_db.py @@ -10,6 +10,8 @@ import os from urllib.parse import quote_plus +from psycopg.conninfo import conninfo_to_dict + def build_db_uri(prefix: str = "DB_") -> str: """Assemble ``HOST``/``PORT``/``NAME``/``USER``/``PASSWORD`` into a URI. @@ -28,3 +30,20 @@ def build_db_uri(prefix: str = "DB_") -> str: f"postgresql://{quote_plus(user)}:{quote_plus(password)}" f"@{host}:{port}/{name}?sslmode={sslmode}" ) + + +def point_app_pool_at(conninfo: str) -> None: + """Make the app's pool, which builds its target from DB_*, connect where + the script writes. Must run before the app is imported, whose dotenv + load fills only what is unset.""" + params = conninfo_to_dict(conninfo) + for key, param in ( + ("DB_HOST", "host"), + ("DB_PORT", "port"), + ("DB_NAME", "dbname"), + ("DB_USER", "user"), + ("DB_PASSWORD", "password"), + ("DB_SSLMODE", "sslmode"), + ): + if params.get(param): + os.environ[key] = str(params[param]) diff --git a/scripts/ops/_delta_loads.py b/scripts/ops/_delta_loads.py new file mode 100644 index 000000000..52c2d5fc8 --- /dev/null +++ b/scripts/ops/_delta_loads.py @@ -0,0 +1,444 @@ +"""Loads of the ``messages`` DeltaChannel the way the server performs them, +and the verified transaction that rewrites how a thread stores them. + +``thin_message_snapshots`` plans against a thread's checkpoint tree and checks +its own writes by loading through the server's reader and saver inside the +writing transaction, before and after; these are that tree, those loads, +metered, and the transaction around them. +""" + +from __future__ import annotations + +import enum +import functools +import hashlib +from collections import defaultdict +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any, Generic, TypeVar + +import psycopg +from psycopg.conninfo import conninfo_to_dict +from psycopg.rows import dict_row + +from scripts._errors import where + +MESSAGES = "messages" + +# The releases whose storage format and load path these scripts were checked +# against: the stage-1 walk stops at the first checkpoint whose +# ``channel_values`` has a ``messages`` key, a direct load joins the blob by +# channel version, and a flag without a blob seeds empty. +SUPPORTED_RELEASES = { + "langgraph": "1.2.11", + "langgraph-checkpoint": "4.1.1", + "langgraph-checkpoint-postgres": "3.1.1", +} + + +def unsupported_saver() -> str | None: + """Why the installed saver is not one these scripts may edit, or None.""" + from importlib.metadata import PackageNotFoundError, version + + for dist, wanted in SUPPORTED_RELEASES.items(): + try: + installed = version(dist) + except PackageNotFoundError: + return f"{dist} is not installed" + if installed != wanted: + return f"{dist} {installed} is not a verified release ({wanted})" + # Longer walks reach targets past the first stage-1 page, which only the + # server's walk guard rebuilds correctly on these releases. + import src.server.utils.checkpointer # noqa: F401 + from langgraph.checkpoint.postgres.base import BasePostgresSaver + + if not getattr(BasePostgresSaver._try_advance_walks, "_waits_for_target", False): + return "the server's delta walk guard is not installed" + return None + + +class Mode(enum.Enum): + DRY_RUN = "dry run" + REHEARSE = "rehearse" + APPLY = "apply" + + +@dataclass +class Walk: + """What one load cost the saver: its stage-1 scan and the writes it replayed.""" + + pages: int = 0 + rows: int = 0 + chain: int = 0 + writes: int = 0 + seed: bool = False + + def __str__(self) -> str: + if not self.pages: + return "no walk" + return ( + f"pages={self.pages} stage1_rows={self.rows} chain={self.chain} " + f"writes={self.writes} seed={'yes' if self.seed else 'no'}" + ) + + def no_longer_than(self, other: Walk) -> bool: + return ( + self.pages <= other.pages + and self.chain <= other.chain + and self.writes <= other.writes + ) + + +_NODES_SQL = """ +SELECT checkpoint_id, parent_checkpoint_id, + checkpoint -> 'channel_versions' ->> 'messages' AS version, + (checkpoint -> 'channel_values' -> 'messages') IS NOT NULL AS flagged, + xmin::text AS xmin +FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = '' +ORDER BY checkpoint_id +""" + +_BLOB_SIZES_SQL = """ +SELECT version, length(blob) AS size FROM checkpoint_blobs +WHERE thread_id = %s AND checkpoint_ns = '' AND channel = 'messages' + AND type <> 'empty' AND blob IS NOT NULL +""" + + +@dataclass +class Node: + checkpoint_id: str + parent: str | None + version: str | None + flagged: bool + xmin: str + # messages versions along the path from the root: the updates a walk + # from an ancestor replays. + updates: int = 0 + + +class Tree: + """A thread's root-namespace checkpoints, parents before children.""" + + def __init__(self, rows: list[dict[str, Any]]) -> None: + self.nodes: dict[str, Node] = {} + self.order: list[str] = [] + self.children: dict[str, list[str]] = defaultdict(list) + self.by_version: dict[str, list[str]] = defaultdict(list) + self.problem: str | None = None + for row in rows: + cid = row["checkpoint_id"] + self.nodes[cid] = Node( + cid, row["parent_checkpoint_id"], row["version"], row["flagged"], row["xmin"] + ) + self.order.append(cid) + for cid in self.order: + node = self.nodes[cid] + parent = self.nodes.get(node.parent) if node.parent else None + if node.parent is not None and parent is None: + self.problem = "a checkpoint's parent is missing" + elif parent is not None and node.parent >= cid: + # The planners visit parents first, in id order. + self.problem = "a parent sorts after its child" + elif parent is not None: + self.children[node.parent].append(cid) + node.updates = parent.updates + (node.version != parent.version) + if node.version is not None: + self.by_version[node.version].append(cid) + + def leaves(self) -> list[str]: + return [cid for cid in self.order if not self.children.get(cid)] + + def walk(self, checkpoint_id: str, flags: set[str]) -> tuple[str | None, list[str]]: + """The saver's walk for a checkpoint that loads no blob: parents from + its own up to and including the first flagged one, which seeds it.""" + chain: list[str] = [] + cur = self.nodes[checkpoint_id].parent + while cur is not None: + chain.append(cur) + if cur in flags: + return cur, chain + cur = self.nodes[cur].parent + return None, chain + + +async def read_tree(conn: psycopg.AsyncConnection, thread_id: str) -> tuple[Tree, dict[str, int]]: + """The thread's tree, and the size of each stored ``messages`` blob by version.""" + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute(_NODES_SQL, (thread_id,)) + tree = Tree(await cur.fetchall()) + await cur.execute(_BLOB_SIZES_SQL, (thread_id,)) + sizes = {r["version"]: r["size"] for r in await cur.fetchall()} + return tree, sizes + + +@functools.cache +def _metered_saver_class(): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + from langgraph.checkpoint.postgres.base import BasePostgresSaver + + class MeteredSaver(AsyncPostgresSaver): + """The server's saver, counting what each delta rebuild reads. + + Bound to one connection, so a load inside a script's transaction sees + what it has written but not yet committed. + """ + + def __init__(self, conn: psycopg.AsyncConnection) -> None: + super().__init__(conn) + self.walk = Walk() + + def _ingest_stage1_page(self, stage1_rows, *args): # type: ignore[override] + self.walk.pages += 1 + self.walk.rows += len(stage1_rows) + return BasePostgresSaver._ingest_stage1_page(stage1_rows, *args) + + def _build_delta_channels_writes_history(self, **kwargs): + histories = super()._build_delta_channels_writes_history(**kwargs) + history = histories.get(MESSAGES) or {} + self.walk.chain += len(kwargs["chain_by_ch"].get(MESSAGES, ())) + self.walk.writes += len(history.get("writes", ())) + self.walk.seed = self.walk.seed or "seed" in history + return histories + + return MeteredSaver + + +class Loader: + """Reads ``messages`` through the server's reader graph on one connection.""" + + def __init__(self, conn: psycopg.AsyncConnection) -> None: + # Importing the checkpointer module installs the server's walk guard on + # the saver base class; the reader is what the history endpoints use. + import src.server.utils.checkpointer # noqa: F401 + from src.server.services.history.reader import CheckpointHistoryReader + + self.saver = _metered_saver_class()(conn) + self.reader = CheckpointHistoryReader(self.saver) + + @property + def serde(self) -> Any: + return self.saver.serde + + @property + def spec(self) -> Any: + """The ``messages`` channel as the reader graph declares it.""" + return self.reader.graph.channels[MESSAGES] + + async def channel_value(self, thread_id: str, checkpoint_id: str) -> tuple[list, Walk]: + """The channel's value at the checkpoint, before its own pending writes. + + This is what a snapshot must hold: a child replays the checkpoint's + writes on top of it. ``aget_state`` would also apply null-task writes. + """ + from langgraph.pregel._checkpoint import achannels_from_checkpoint + + self.saver.walk = Walk() + config = { + "configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": checkpoint_id} + } + saved = await self.saver.aget_tuple(config) + if saved is None: + raise LookupError(f"checkpoint {checkpoint_id} not found") + graph = self.reader.graph + graph._migrate_checkpoint(saved.checkpoint) # as aget_state does first + channels, _ = await achannels_from_checkpoint( + graph.channels, saved.checkpoint, saver=self.saver, config=saved.config + ) + return list(channels[MESSAGES].get()), self.saver.walk + + async def state_messages(self, thread_id: str, checkpoint_id: str) -> tuple[list, Walk]: + """``messages`` as the history reader serves the checkpoint's state.""" + self.saver.walk = Walk() + snapshot = await self.reader.aget_state(thread_id, checkpoint_id) + return list(snapshot.values.get(MESSAGES) or []), self.saver.walk + + async def digest(self, thread_id: str, checkpoint_id: str) -> tuple[str, Walk]: + messages, walk = await self.state_messages(thread_id, checkpoint_id) + return messages_digest(self.serde, messages), walk + + +def messages_digest(serde: Any, messages: list) -> str: + """A list's serialized bytes, hashed: two lists match when their digests + do, without holding either. Ids, types, content and every field count.""" + tag, data = serde.dumps_typed(list(messages)) + digest = hashlib.blake2b(tag.encode(), digest_size=16) + digest.update(b"\x1e") + digest.update(data) + return digest.hexdigest() + + +def describe(conninfo: str) -> str: + params = conninfo_to_dict(conninfo) + return f"{params.get('host', '?')}:{params.get('port', '5432')}/{params.get('dbname', '?')}" + + +# -------------------------------------------------------------------------- +# One thread's verified rewrite +# -------------------------------------------------------------------------- + + +class Status(enum.Enum): + PLANNED = "would write" + NOTHING = "nothing to do" + REHEARSED = "rehearsed" + WRITTEN = "written" + SKIPPED = "skipped" + MISMATCH = "mismatch" + ERROR = "error" + + +class Skip(Exception): + """Leave the thread as it is, for the reason given: raised by a plan, or by + its write when a row is no longer what the plan read.""" + + +@dataclass(frozen=True) +class Timeouts: + statement_s: float = 120.0 + lock_s: float = 5.0 + + async def set_on(self, conn: psycopg.AsyncConnection) -> None: + statement = f"{int(self.statement_s * 1000)}ms" + # Planning and loads run in Python with the transaction open, so the + # session gets as long idle in it as one statement, and no longer. + await conn.execute( + "SELECT set_config('statement_timeout', %s, false), set_config('lock_timeout', %s, false)," + " set_config('idle_in_transaction_session_timeout', %s, false)", + (statement, f"{int(self.lock_s * 1000)}ms", statement), + ) + + +R = TypeVar("R") + + +@dataclass +class Rewrite(Generic[R]): + """A thread's planned change, which ``rewrite_thread`` carries out. + + ``write`` raises ``Skip`` to refuse; without one there is nothing to do. + ``verify`` gets every check once loaded again and names what it finds + broken, or returns None. + """ + + report: R + checks: list[str] = field(default_factory=list) + committed: str | None = None # the check loaded again once committed + write: Callable[[], Awaitable[None]] | None = None + verify: Callable[[dict[str, Check]], Awaitable[str | None]] | None = None + + +@dataclass +class Check: + checkpoint_id: str + before: Walk + digest: str # the pre-write load, compared against after the write + after: Walk | None = None + match: bool | None = None + + +@dataclass +class Result(Generic[R]): + thread_id: str + status: Status + reason: str = "" + report: R | None = None + checks: list[Check] = field(default_factory=list) + committed: Check | None = None + committed_load: Walk | None = None + committed_match: bool | None = None + + @property + def failed(self) -> bool: + return self.status in (Status.MISMATCH, Status.ERROR) or self.committed_match is False + + +Planner = Callable[[psycopg.AsyncConnection, Loader], Awaitable[Rewrite[R]]] + + +async def rewrite_thread( + conninfo: str, thread_id: str, mode: Mode, timeouts: Timeouts, plan: Planner[R] +) -> Result[R]: + """Plan, write and verify one thread in its own REPEATABLE READ transaction. + + Every read, the write and both rounds of loads share one snapshot, so what + is compared after the write is exactly what was read before it. A dry run + only plans. The transaction commits only in apply mode once every check + loads as it did before and ``verify`` finds nothing; any other ending + rolls it back. + """ + refusal = unsupported_saver() + if refusal: + return Result(thread_id, Status.ERROR, f"refused: {refusal}") + async with await psycopg.AsyncConnection.connect(conninfo, autocommit=True) as conn: + await timeouts.set_on(conn) + try: + async with conn.transaction() as tx: + isolation = "REPEATABLE READ" + (" READ ONLY" if mode is Mode.DRY_RUN else "") + await conn.execute(f"SET TRANSACTION ISOLATION LEVEL {isolation}") + result = await _rewrite(conn, thread_id, mode, plan) + if not (mode is Mode.APPLY and result.status is Status.WRITTEN): + raise psycopg.Rollback(tx) + except psycopg.errors.SerializationFailure: + return Result(thread_id, Status.SKIPPED, "a checkpoint changed concurrently") + except psycopg.errors.LockNotAvailable: + return Result(thread_id, Status.SKIPPED, "a checkpoint row is locked by another writer") + except psycopg.errors.QueryCanceled: + return Result(thread_id, Status.SKIPPED, "statement timeout") + except Exception as exc: # one thread's failure must not stop the run + return Result(thread_id, Status.ERROR, where(exc)) + + if result.status is Status.WRITTEN and result.committed is not None: + # The next load the way the server does it: a later transaction, + # reading only committed rows. + check = result.committed + try: + digest, walk = await Loader(conn).digest(thread_id, check.checkpoint_id) + except Exception as exc: + result.committed_match, result.reason = False, f"committed load failed: {where(exc)}" + else: + result.committed_load = walk + assert check.after is not None + result.committed_match = digest == check.digest and walk.no_longer_than(check.after) + return result + + +async def _rewrite( + conn: psycopg.AsyncConnection, thread_id: str, mode: Mode, plan: Planner[R] +) -> Result[R]: + loader = Loader(conn) + try: + rewrite = await plan(conn, loader) + except Skip as skip: + return Result(thread_id, Status.SKIPPED, str(skip)) + result = Result(thread_id, Status.NOTHING, report=rewrite.report) + if rewrite.write is None: + return result + if mode is Mode.DRY_RUN: + result.status = Status.PLANNED + return result + + ids = list(rewrite.checks) + if rewrite.committed is not None and rewrite.committed not in ids: + ids.insert(0, rewrite.committed) + for cid in ids: + digest, walk = await loader.digest(thread_id, cid) + result.checks.append(Check(cid, walk, digest)) + try: + await rewrite.write() + except Skip as skip: + result.status, result.reason = Status.SKIPPED, str(skip) + return result + + for check in result.checks: + digest, check.after = await loader.digest(thread_id, check.checkpoint_id) + check.match = digest == check.digest + checks = {c.checkpoint_id: c for c in result.checks} + result.committed = checks.get(rewrite.committed) if rewrite.committed else None + broken = await rewrite.verify(checks) if rewrite.verify else None + if broken or not all(c.match for c in result.checks): + result.status = Status.MISMATCH + result.reason = f"rolled back: {broken}" if broken else "rolled back" + return result + result.status = Status.WRITTEN if mode is Mode.APPLY else Status.REHEARSED + return result diff --git a/scripts/ops/_thin_plan.py b/scripts/ops/_thin_plan.py new file mode 100644 index 000000000..a646df7cb --- /dev/null +++ b/scripts/ops/_thin_plan.py @@ -0,0 +1,290 @@ +"""The thinning planner: which ``messages`` blobs a thread keeps, flags and seeds. + +It reads only the thread's tree, its stored blobs and its writes, and +computes loads the way the saver does to propose a plan. The plan is only a +proposal: what commits is decided by real loads before and after the write +(``_delta_loads.rewrite_thread``). +""" + +from __future__ import annotations + +from collections import Counter, OrderedDict +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from scripts.ops._delta_loads import Loader, Skip, Tree, messages_digest + +# A checkpoint's ``messages`` writes as the saver replays them: (task id, +# channel, value), ordered by task id, then index. +Writes = dict[str, list[tuple[str, str, Any]]] + + +@dataclass +class Unit: + """One stored ``messages`` blob and the checkpoints that load it directly + (every checkpoint at its version).""" + + version: str + size: int + loaders: list[str] + flagged: list[str] + + @property + def is_snapshot(self) -> bool: + return bool(self.flagged) + + +@dataclass +class Seed: + """A snapshot to store at a version that holds no blob: the value the + topmost checkpoint of a walk across a newly flagged list loads today.""" + + checkpoint_id: str + version: str + snapshot: Any # a _DeltaSnapshot + typed: tuple[str, bytes] # serialized once, for the report and the write + + @classmethod + def of(cls, checkpoint_id: str, version: str, value: list, serde: Any) -> Seed: + from langgraph.checkpoint.serde.types import _DeltaSnapshot + + snapshot = _DeltaSnapshot(value) + return cls(checkpoint_id, version, snapshot, serde.dumps_typed(snapshot)) + + @property + def size(self) -> int: + return len(self.typed[1]) + + +@dataclass +class Layout: + """Which checkpoints flag a snapshot, and which versions hold a blob: + one stored now, or a seed not yet written.""" + + flags: set[str] + stored: set[str] + pending: dict[str, Seed] = field(default_factory=dict) + + def holds(self, version: str | None) -> bool: + return version is not None and (version in self.stored or version in self.pending) + + def copy(self) -> Layout: + return Layout(set(self.flags), set(self.stored), dict(self.pending)) + + +class Replayer: + """Loads computed the way the saver computes them, from the rows the + planning transaction read, under any layout.""" + + def __init__( + self, + thread_id: str, + tree: Tree, + loader: Loader, + read_blob: Callable[[str], Awaitable[tuple[str, bytes] | None]], + read_writes: Callable[[], Awaitable[Writes]], + ) -> None: + self._thread_id = thread_id + self._tree = tree + self._spec = loader.spec + self.serde = loader.serde + self._read_blob = read_blob + self._read_writes = read_writes + self._writes: Writes | None = None + self._raw: OrderedDict[str, Any] = OrderedDict() + self._stored: dict[str, tuple[list[tuple[Any, Any]], str]] = {} + + async def writes(self) -> Writes: + """Every ``messages`` write of the thread, decoded on first use: a + plan that simulates no load never pays for it.""" + if self._writes is None: + self._writes = await self._read_writes() + return self._writes + + async def _stored_raw(self, version: str) -> Any: + if version in self._raw: + self._raw.move_to_end(version) + return self._raw[version] + typed = await self._read_blob(version) + if typed is None: + raise LookupError(f"messages blob at version {version} is gone") + value = self.serde.loads_typed(typed) + self._raw[version] = value + if len(self._raw) > 8: + self._raw.popitem(last=False) + return value + + async def _blob(self, version: str, layout: Layout) -> Any: + seed = layout.pending.get(version) + return seed.snapshot if seed is not None else await self._stored_raw(version) + + async def load(self, checkpoint_id: str, layout: Layout) -> list: + from langgraph._internal._typing import MISSING + + node = self._tree.nodes[checkpoint_id] + if layout.holds(node.version): + assert node.version is not None + channel = self._spec.from_checkpoint(await self._blob(node.version, layout)) + else: + seed, chain = self._tree.walk(checkpoint_id, layout.flags) + base: Any = MISSING + seed_version = self._tree.nodes[seed].version if seed else None + if seed_version is not None and layout.holds(seed_version): + base = await self._blob(seed_version, layout) + channel = self._spec.from_checkpoint(base) + writes = await self.writes() + channel.replay_writes([w for cid in reversed(chain) for w in writes.get(cid, ())]) + return list(channel.get()) if channel.is_available() else [] + + async def stored(self, version: str) -> tuple[list[tuple[Any, Any]], str]: + """The shape and digest of what loads of a stored blob return.""" + if version not in self._stored: + value = list(self._spec.from_checkpoint(await self._stored_raw(version)).get()) + self._stored[version] = (_shape(value), messages_digest(self.serde, value)) + return self._stored[version] + + async def reproduces(self, checkpoint_id: str, layout: Layout, expected: tuple) -> bool: + value = await self.load(checkpoint_id, layout) + shape, digest = expected + # Ids first: most differences show there, without serializing. + return _shape(value) == shape and messages_digest(self.serde, value) == digest + + +def _shape(messages: list) -> list[tuple[Any, Any]]: + return [(getattr(m, "type", None), getattr(m, "id", None)) for m in messages] + + +@dataclass +class Plan: + before: Layout + after: Layout + drop: list[Unit] = field(default_factory=list) + promote: list[Unit] = field(default_factory=list) + seeds: list[Seed] = field(default_factory=list) + flags: list[str] = field(default_factory=list) + kept: Counter[str] = field(default_factory=Counter) + + @property + def changes(self) -> bool: + return bool(self.drop or self.promote or self.seeds) + + +async def make_plan( + tree: Tree, + units: dict[str, Unit], + replay: Replayer, + protected: dict[str, str], + spacing: int | None, +) -> Plan: + """Decide each blob, oldest first, against the layout the decisions above + it left, so every load is checked against the snapshot it will seed from. + + ``spacing`` None keeps every blob: lists stored before DeltaChannel are + only flagged, and the walks that crossed them seeded. + """ + old = Layout({c for c, n in tree.nodes.items() if n.flagged}, set(units)) + new = old.copy() + plan = Plan(old, new) + + forced: dict[str, str] = {} + for cid, why in protected.items(): + version = tree.nodes[cid].version + if version in units: + forced.setdefault(version, why) + continue + seed, _ = tree.walk(cid, old.flags) + seed_version = tree.nodes[seed].version if seed else None + if seed_version in units: + forced.setdefault(seed_version, why) + + def keep_reason(unit: Unit) -> str | None: + if spacing is None: + return "keep all" + if unit.version in forced: + return forced[unit.version] + seed, _ = tree.walk(unit.loaders[0], new.flags) + since = tree.nodes[unit.loaders[0]].updates - (tree.nodes[seed].updates if seed else 0) + return "spacing" if since >= spacing else None + + async def droppable(unit: Unit) -> bool: + expected = await replay.stored(unit.version) + new.stored.discard(unit.version) + new.flags.difference_update(unit.flagged) + for cid in unit.loaders: + if not await replay.reproduces(cid, new, expected): + new.stored.add(unit.version) + new.flags.update(unit.flagged) + return False + return True + + by_age = sorted(units.values(), key=lambda u: u.loaders[0]) + + # Lists stored before DeltaChannel: no walk stops at them yet. + for unit in (u for u in by_age if not u.is_snapshot): + reason = keep_reason(unit) + if reason is None and await droppable(unit): + plan.drop.append(unit) + continue + plan.kept[reason or "replay differs"] += 1 + plan.promote.append(unit) + new.flags.update(unit.loaders) + + # Walks that today pass a list just flagged above would now stop at it. + # Seed the topmost of each with what it loads today. + added = new.flags - old.flags + crosses: dict[str, bool] = {} + for cid in tree.order: + parent = tree.nodes[cid].parent + crosses[cid] = parent is not None and ( + parent in added or (parent not in old.flags and crosses[parent]) + ) + exposed = {cid for cid in tree.order if crosses[cid] and not old.holds(tree.nodes[cid].version)} + for cid in tree.order: + if cid not in exposed or tree.nodes[cid].parent in exposed: + continue + version = tree.nodes[cid].version + if version is None: + raise Skip("a walk crossing a kept stored list has no messages version") + if not new.holds(version): + value = await replay.load(cid, old) + expected = (_shape(value), messages_digest(replay.serde, value)) + for peer in tree.by_version[version]: + if peer != cid and not await replay.reproduces(peer, old, expected): + raise Skip("a seed would change a load sharing its version") + seed = Seed.of(cid, version, value, replay.serde) + new.pending[version] = seed + plan.seeds.append(seed) + new.flags.add(cid) + plan.flags.append(cid) + + # Flagged snapshots: DeltaChannel's own, and lists an earlier run flagged + # (of this script, or of seed_delta_tips before --keep all replaced it). + for unit in (u for u in by_age if u.is_snapshot): + reason = keep_reason(unit) + if reason is None and await droppable(unit): + plan.drop.append(unit) + else: + plan.kept[reason or "replay differs"] += 1 + + for unit in plan.promote: + plan.flags.extend(cid for cid in unit.loaders if cid not in old.flags) + return plan + + +def longest_walk(tree: Tree, layout: Layout, write_counts: dict[str, int]) -> tuple[int, int]: + """The longest walk any checkpoint's load takes under a layout, as + (checkpoints, writes) replayed.""" + span: dict[str, tuple[int, int]] = {} + longest = (0, 0) + for cid in tree.order: + node = tree.nodes[cid] + if node.parent is None: + up = (0, 0) + else: + above = (0, 0) if node.parent in layout.flags else span[node.parent] + up = (above[0] + 1, above[1] + write_counts.get(node.parent, 0)) + span[cid] = up + if not layout.holds(node.version): + longest = max(longest, up) + return longest diff --git a/scripts/ops/repair_crossed_threads.py b/scripts/ops/repair_crossed_threads.py new file mode 100644 index 000000000..c4703feed --- /dev/null +++ b/scripts/ops/repair_crossed_threads.py @@ -0,0 +1,1117 @@ +"""Trim the window of threads whose crossing into DeltaChannel reordered their oldest messages. + +The sliding window trims a thread's checkpoint only once the server's turn +slices, read in branch order, hold exactly the messages it would drop +(``compaction.window``, ``history.window.runs_held``). A thread stored before +DeltaChannel whose first delta turn ran before ``thin_message_snapshots.py +--keep all`` seeded it "crossed": that turn's load replayed the old writes +from the root, which orders parallel tool results by task id rather than the +order they were applied and gives id-less user inputs new ids, and froze that +order into the checkpoint. The slices, cut from the stored lists, keep the +order the thread was written in. The same messages, in another order, hash to +another digest, so the window refuses the thread at every summary and its +checkpoint keeps growing. + +The model never reads that order. A summary stands in for everything before +its boundary, the turns there are numbered by how many there are, and the +index a summary lists names only the newest twenty. So this does the trim the +window would do, offline, with the window it would keep had the thread never +crossed: the trimmed runs' count, digest and index text taken from the +slices, whose order replay and the transcripts read. The kept messages, and +every carry other readers count over the trimmed ones, come from the +checkpoint as the middleware's own trim takes them. + +A thread is repaired only when all of this holds, and refused otherwise: + +- no run is in progress, the saver is a verified release, and the thread's + newest checkpoint is its committed tip, with no pending write but a failed + task's error (a new turn discards that task; the repaired tip does not + carry it); +- every turn before the cut has a stored slice replay would serve, and those + slices hold the same messages as the tip's head, compared per message by + what it says and the calls it makes or answers, with the same turn count; +- every position where the two orders differ lies within the messages the + crossing replayed: the list of the first delta checkpoint after the + branch's newest pre-DeltaChannel list. Past it they agree message for + message; +- the model is sent the same requests from the repaired tip as from the + old one: the next call, the next call after an idle pause (Tier 1 cuts), + and the next summary with the call after it, ids it mints aside. + +The trim is written as a turn writes it, by a graph node returning the +middleware's own update through the Pregel loop, so the ``Overwrite`` starts +a new ``messages`` snapshot and the new tip loads without replaying the old +list. The graph declares ``MainAgentState`` plus any other field the tip +holds, so no stored value is dropped. The step's writes are then removed +from the old tip, which a DeltaChannel load of any other child of it would +replay (see ``_write``). + +Each thread is one REPEATABLE READ transaction (``_delta_loads.rewrite_thread``): +it loads the tip and the turn anchors and a sample of checkpoints before the +write and again after, checks that the new tip loads the old one's values +with the trim applied, that the turn anchors only end the last turn on it +and that the slices hold the window it keeps, and rolls back on any +difference. Only then is the thread's tip pointer moved onto the new tip by +compare-and-set, so a turn starting meanwhile waits on the row only until +the commit. ``updated_at`` is left alone: this is not activity, and the +thread lists order by it. A thread whose pointer moved since is skipped, and +a turn that starts from the old tip first simply wins, leaving the repair as +a dead branch. Once committed, the old tip is loaded again, and the last +turn, whose slice is keyed by the checkpoint that ends it, is projected +again and compared with what was stored. + +Rerunning is a no-op: a repaired thread has nothing left before its +summary's run to trim. Threads that are not crossed are reported and never +written. Prints ids, counts and booleans, never message content. + +Usage (from the repo root, inside the backend container), after +``thin_message_snapshots.py --keep all`` and ``backfill_turn_slices.py``: + uv run python scripts/ops/repair_crossed_threads.py # detect, plan + uv run python scripts/ops/repair_crossed_threads.py --rehearse # write, verify, roll back + uv run python scripts/ops/repair_crossed_threads.py --apply # write, verify, commit + uv run python scripts/ops/repair_crossed_threads.py --thread ID --verify all +""" + +from __future__ import annotations + +import argparse +import asyncio +import functools +import json +import logging +import random +import re +import sys +import uuid +from collections import Counter +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field, replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import psycopg + +# Run as a script, sys.path[0] is scripts/ops, not the repo root the sibling +# helpers and the app are imported from. +_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(_ROOT)) +sys.path.insert(0, str(_ROOT / "src")) + +from scripts._errors import where # noqa: E402 +from scripts.ops._db import build_db_uri, point_app_pool_at # noqa: E402 +from scripts.ops._delta_loads import ( # noqa: E402 + MESSAGES, + Check, + Loader, + Mode, + Result, + Rewrite, + Skip, + Status, + Timeouts, + Tree, + describe, + messages_digest, + read_tree, + rewrite_thread, + unsupported_saver, +) +from scripts.utils._thread_job import ( # noqa: E402 + add_selection_args, + app_infra, + for_each_thread, + selection_sql, +) + +logger = logging.getLogger("repair_crossed_threads") + +#: The node the window's trim runs in, so the repair's checkpoint reads like +#: one a turn would have written. +NODE = "CompactionMiddleware.before_agent" +#: The tip's only pending write a repair may leave behind it. +_ERROR = "__error__" + +_THREAD_SQL = """ +SELECT latest_checkpoint_id FROM conversation_threads WHERE conversation_thread_id = %s::uuid +""" + +_LIVE_SQL = """ +SELECT EXISTS ( + SELECT 1 FROM conversation_responses + WHERE conversation_thread_id = %(thread_id)s::uuid AND status = 'in_progress') + OR EXISTS ( + SELECT 1 FROM subagent_runs + WHERE thread_id = %(thread_id)s::uuid AND status = 'in_progress') +""" + +# What a turn started without a checkpoint id builds on. +_NEWEST_SQL = """ +SELECT checkpoint_id FROM checkpoints +WHERE thread_id = %s AND checkpoint_ns = '' ORDER BY checkpoint_id DESC LIMIT 1 +""" + +_PENDING_SQL = """ +SELECT DISTINCT channel FROM checkpoint_writes +WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s +""" + +_TIP_WRITES_SQL = """ +SELECT task_id, idx, channel, md5(blob) FROM checkpoint_writes +WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s +ORDER BY task_id, idx +""" + +_CHANNELS_SQL = """ +SELECT jsonb_object_keys(checkpoint -> 'channel_versions') FROM checkpoints +WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s +""" + +_BLOB_SQL = """ +SELECT type, blob FROM checkpoint_blobs +WHERE thread_id = %s AND checkpoint_ns = '' AND channel = 'messages' AND version = %s +""" + +# Compare-and-set, as ``advance_thread_checkpoint_id`` does, without touching +# ``updated_at``. Under REPEATABLE READ a concurrent update of the row fails +# the statement, which skips the thread. +_ADVANCE_SQL = """ +UPDATE conversation_threads SET latest_checkpoint_id = %s +WHERE conversation_thread_id = %s::uuid AND latest_checkpoint_id = %s +""" + +_THREADS_SQL = """ +SELECT t.conversation_thread_id::text FROM conversation_threads t +JOIN workspaces w ON w.workspace_id = t.workspace_id +WHERE {where} +ORDER BY t.updated_at DESC +""" + + +# -------------------------------------------------------------------------- +# One thread +# -------------------------------------------------------------------------- + + +@dataclass +class Options: + verify_all: bool = False + sample: int = 20 + timeouts: Timeouts = field(default_factory=Timeouts) + + +@dataclass +class Report: + """What detection found on one thread, and what the repair did.""" + + #: yes, no, or ? when the slices cannot say. + crossed: str = "?" + detail: str = "" + refused: str = "" + tip: str | None = None + messages: int = 0 + cut: int = 0 + runs: int = 0 + #: Positions before the cut where the slices' order differs from the tip's. + differing: int = 0 + span: tuple[int, int] | None = None + #: Messages the crossing replayed, which the differences must lie within. + replayed: int = 0 + #: Turns whose index text the slices' order changes. + lines: list[int] = field(default_factory=list) + requests: int = 0 + requests_match: bool | None = None + repaired_tip: str | None = None + #: The last turn projected again once committed, and whether its + #: messages came out as stored. + reprojected: int | None = None + reprojected_same: bool | None = None + + def not_crossed(self, detail: str) -> Report: + self.crossed, self.detail = "no", detail + return self + + +@dataclass(frozen=True) +class Settings: + """What the main agent's compaction reads, from the agent config.""" + + keep_messages: int + offload: Any # OffloadSettings + + +def _key(message: Any) -> str: + """A message as the window's digest keys it: what it says and the calls + it makes or answers, never its id.""" + from ptc_agent.agent.transcript.identity import extend + + return extend("", [message]) + + +async def repair_thread( + conninfo: str, thread_id: str, mode: Mode, opts: Options, settings: Settings +) -> Result[Report]: + """Detect, and in a write mode repair, one thread in its own verified + transaction; once committed, project its last turn again.""" + + async def plan(conn: psycopg.AsyncConnection, loader: Loader) -> Rewrite[Report]: + return await _plan(conn, loader, thread_id, opts, settings) + + out = await rewrite_thread(conninfo, thread_id, mode, opts.timeouts, plan) + if out.status is Status.WRITTEN and out.committed_match and out.report is not None: + try: + out.report.reprojected, out.report.reprojected_same = await _reproject(thread_id) + except Exception as exc: # replay projects it on its next read instead + out.reason = f"committed; projecting the last turn failed: {where(exc)}" + return out + + +async def _plan( + conn: psycopg.AsyncConnection, + loader: Loader, + thread_id: str, + opts: Options, + settings: Settings, +) -> Rewrite[Report]: + """The thread's repair, or the report of why there is none. A refusal is + kept in the report rather than raised, so it says what detection found.""" + report = Report() + try: + return await _repair(conn, loader, thread_id, opts, settings, report) + except Skip as skip: + report.refused = str(skip) + return Rewrite(report) + + +async def _repair( + conn: psycopg.AsyncConnection, + loader: Loader, + thread_id: str, + opts: Options, + settings: Settings, + report: Report, +) -> Rewrite[Report]: + from ptc_agent.agent.middleware.compaction.window import window_cut + from ptc_agent.agent.transcript import Window + + try: + uuid.UUID(thread_id) + except ValueError: + raise Skip("not a conversation thread") from None + async with conn.cursor() as cur: + await cur.execute(_THREAD_SQL, (thread_id,)) + row = await cur.fetchone() + if row is None: + raise Skip("no conversation thread row") + tip = row[0] + if not tip: + return Rewrite(report.not_crossed("no turn has finished")) + await cur.execute(_LIVE_SQL, {"thread_id": thread_id}) + if (await cur.fetchone())[0]: + raise Skip("a run is in progress") + report.tip = tip + + tree, _ = await read_tree(conn, thread_id) + if tree.problem or tip not in tree.nodes: + raise Skip(tree.problem or "the tip is not a root checkpoint") + extra = await _extra_fields(conn, thread_id, tip) + values = await _values(_graph(loader.saver, extra), thread_id, tip) + messages = list(values.get(MESSAGES) or []) + report.messages = len(messages) + event = values.get("_summarization_event") + cut = window_cut(messages, event) + report.cut = cut + base = Window.of(values) + if cut <= 0: + report.detail = _no_cut(event, base) + await _cross_whole(loader, thread_id, tip, messages, base, report) + return Rewrite(report) + + # Detection: what the window's coverage check reads, at this tip. + head = await asyncio.to_thread(base.extend, messages[:cut]) + stored = await _stored_head(loader, thread_id, tip, head.messages) + if stored is None: + raise Skip("a turn before the cut has no stored slice replay would serve") + if await asyncio.to_thread(head.holds, stored): + report.detail = "its slices hold the head; the window trims it itself" + await _cross_whole(loader, thread_id, tip, messages, base, report) + return Rewrite(report) + if base.messages and not base.holds(stored[: base.messages]): + raise Skip("its slices do not hold what the window already trimmed") + canon = stored[base.messages :] + tip_keys = await asyncio.to_thread(lambda: [_key(m) for m in messages[:cut]]) + canon_keys = await asyncio.to_thread(lambda: [_key(m) for m in canon]) + if Counter(tip_keys) != Counter(canon_keys): + raise Skip("its slices hold other messages than the tip's head") + report.crossed = "yes" + window = await asyncio.to_thread(base.extend, canon) + report.runs = window.runs + if window.runs != head.runs: + raise Skip("its slices count the trimmed turns differently") + differing = [i for i, (a, b) in enumerate(zip(tip_keys, canon_keys)) if a != b] + report.differing = len(differing) + report.span = (differing[0], differing[-1]) + report.lines = sorted(n for n, line in window.earlier.items() if head.earlier.get(n) != line) + report.replayed = await _replayed(conn, loader, tree, thread_id, tip) - base.messages + if differing[-1] >= report.replayed: + raise Skip("the orders differ past the messages the crossing replayed") + + # What the write itself needs. + async with conn.cursor() as cur: + await cur.execute(_NEWEST_SQL, (thread_id,)) + if (await cur.fetchone())[0] != tip: + raise Skip("the newest checkpoint is not the thread's tip") + await cur.execute(_PENDING_SQL, (thread_id, tip)) + pending = {r[0] for r in await cur.fetchall()} + if pending - {_ERROR}: + raise Skip("the tip has pending writes (an interrupt or an unfinished step)") + + expected = _applied(values, _trim(values, cut, window, thread_id)) + kept_ids = [m.id for m in expected[MESSAGES]] + before = await _requests(values, thread_id, settings) + report.requests = len(before) + + async def same_requests(after: Mapping[str, Any]) -> str | None: + requests = await _requests(after, thread_id, settings) + report.requests_match = requests == before + if report.requests_match: + return None + return f"the model would be sent other requests ({_changed(before, requests)})" + + # Gated before anything is written, so a dry run answers it too; a write + # gates again on what the repaired tip loads. + refusal = await same_requests(expected) + if refusal: + raise Skip(refusal) + anchors, _ = await loader.reader.aget_turn_anchors(thread_id, tip) + + async def write() -> None: + report.repaired_tip = await _write( + conn, loader.saver, extra, thread_id, tip, cut, window, kept_ids + ) + + async def verify(checks: dict[str, Check]) -> str | None: + new = report.repaired_tip + assert new is not None + after = await _values(_graph(loader.saver, extra), thread_id, new) + broken = _compare(loader.serde, expected, after) + if broken: + return f"the repaired tip loads {broken}" + if not await _stores_snapshot(conn, thread_id, new): + return "the repaired tip does not store a messages snapshot" + moved, _ = await loader.reader.aget_turn_anchors(thread_id, new) + if moved != _ended_at(anchors, new): + return "the turn anchors changed beyond the last turn's end" + # What the next trim's coverage check reads: unless the slices hold + # the window the repaired tip keeps, that trim is refused again. + held = await _stored_head(loader, thread_id, new, window.messages) + if held is None or not await asyncio.to_thread(Window.of(after).holds, held): + return "the slices do not hold the repaired tip's window" + refusal = await same_requests(after) + if refusal: + return refusal + # Last, so the thread's row is locked only from here to the commit: a + # turn starting meanwhile waits on it, and the checks above can take + # a while. + async with conn.cursor() as cur: + await cur.execute(_ADVANCE_SQL, (new, thread_id, tip)) + if cur.rowcount != 1: + return "the thread's tip moved" + return None + + return Rewrite(report, _verification_set(tree, anchors, tip, opts), tip, write, verify) + + +def _no_cut(event: Any, base: Any) -> str: + """Why the window has nothing to trim, so a report tells a thread that + waits for its next compaction from one that never compacted.""" + if not isinstance(event, Mapping): + return "no summary to trim at" + if event.get("anchor_message_id") is None: + return "its summary has no anchor; the next compaction writes one" + if base.runs: + return "trimmed to its summary's run" + return "its summary falls in the first run" + + +_REQUESTS = ("next call", "call after a pause", "summary request", "call after the summary") + + +def _changed(before: list[str], after: list[str]) -> str: + if len(before) != len(after): + return f"{len(before)} requests before, {len(after)} after" + changed = [i for i, (a, b) in enumerate(zip(before, after)) if a != b] + return ", ".join(_REQUESTS[i] if i < len(_REQUESTS) else str(i) for i in changed) + + +async def _extra_fields(conn: psycopg.AsyncConnection, thread_id: str, tip: str) -> list[str]: + """The tip's state fields ``MainAgentState`` does not declare, such as one + a turn's input sets and another middleware's schema declares, or one no + code reads any more. Declared too, the repair carries their values onto + its checkpoint instead of dropping them. The graph's own channels (a + node's trigger, ``__start__``) are not state: the repair's graph writes + those.""" + from typing import get_type_hints + + from ptc_agent.agent.main_state import MainAgentState + + declared = set(get_type_hints(MainAgentState)) + async with conn.cursor() as cur: + await cur.execute(_CHANNELS_SQL, (thread_id, tip)) + names = [r[0] for r in await cur.fetchall()] + return sorted(n for n in names if n not in declared and ":" not in n and not n.startswith("__")) + + +def _graph(saver: Any, extra: Sequence[str], node: Any = None) -> Any: + """The one-node graph the repair reads and writes through, over the + thread's own schema: ``MainAgentState``, as every writer of the main + agent's checkpoints declares it, plus ``extra``. Without ``node`` it + only reads.""" + from typing import Any as AnyType + + from langgraph.graph import START, StateGraph + + from ptc_agent.agent.main_state import MainAgentState + + schema: Any = MainAgentState + if extra: + schema = type(MainAgentState)( + "RepairState", + (MainAgentState,), + {"__annotations__": {name: AnyType for name in extra}, "__module__": __name__}, + total=False, + ) + + async def read_only(state: Mapping[str, Any]) -> dict[str, Any]: + raise RuntimeError("the repair's reading graph never runs") + + return ( + StateGraph(schema) + .add_node(NODE, node or read_only) + .add_edge(START, NODE) + .compile(checkpointer=saver) + ) + + +async def _values(graph: Any, thread_id: str, checkpoint_id: str) -> dict[str, Any]: + snapshot = await graph.aget_state( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + "checkpoint_id": checkpoint_id, + } + } + ) + return dict(snapshot.values) + + +async def _cross_whole( + loader: Loader, thread_id: str, tip: str, messages: list[Any], base: Any, report: Report +) -> None: + """Whether a thread with nothing to repair now crossed anyway: its list + against every message its stored slices hold. A crossing the head does + not reach yet is refused at the trim that first reaches it, so the report + names the thread for a rerun after that compaction.""" + from src.server.database.conversation import turn_slices as slices_db + from src.server.services.history import slices + from src.server.services.history import window as window_reads + + anchors, _ = await loader.reader.aget_turn_anchors(thread_id, tip) + stored: list[Any] = [] + for start in range(0, len(anchors), 64): + batch = anchors[start : start + 64] + rows = { + (row.input_checkpoint_id, row.tail_checkpoint_id): row + for row in await slices_db.get_slices_at( + thread_id, slices.SLICE_KEY, [a.input_checkpoint_id for a in batch] + ) + } + turns = await asyncio.to_thread( + window_reads._decoded, loader.reader.serde, rows, batch, 10**12 + ) + if any(turn is None for turn in turns): + report.detail += "; a turn has no stored slice to compare" + return + stored.extend(m for turn in turns for m in turn) + stored = stored[base.messages :] + if len(stored) > len(messages): + report.detail += "; its slices hold more than its list" + return + ours = await asyncio.to_thread(lambda: [_key(m) for m in messages[: len(stored)]]) + theirs = await asyncio.to_thread(lambda: [_key(m) for m in stored]) + if ours == theirs: + report.crossed = "no" + elif Counter(ours) == Counter(theirs): + differing = [i for i, (a, b) in enumerate(zip(ours, theirs)) if a != b] + report.crossed, report.differing = "yes", len(differing) + report.span = (differing[0], differing[-1]) + report.detail = f"nothing to trim yet: {report.detail}" + else: + report.detail += "; its slices hold other messages than its list" + + +async def _stored_head( + loader: Loader, thread_id: str, tip: str, count: int +) -> list[Any] | None: + """The first ``count`` messages the stored slices of the branch ending at + ``tip`` hold, read as the window's coverage check reads them, or None + when a turn among them has no slice replay would serve. The branch is + walked inside the transaction; the slices come through the app's pool.""" + from src.server.services.history import window as window_reads + + anchors, _ = await loader.reader.aget_turn_anchors(thread_id, tip) + return await window_reads._head(thread_id, anchors, count, extract=False) + + +async def _replayed( + conn: psycopg.AsyncConnection, loader: Loader, tree: Tree, thread_id: str, tip: str +) -> int: + """How many messages the crossing replayed from the root: the list of the + first checkpoint after the branch's newest pre-DeltaChannel list, the one + that first loaded by walking past it. 0 when the branch holds no such + list, so no difference is the crossing's.""" + from langgraph.checkpoint.serde.types import _DeltaSnapshot + + path: list[str] = [] + cur_id: str | None = tip + while cur_id is not None: + path.append(cur_id) + cur_id = tree.nodes[cur_id].parent + stored: set[str | None] = set() + async with conn.cursor() as cur: + for i, cid in enumerate(path): + version = tree.nodes[cid].version + if version is None or version in stored: + continue + stored.add(version) + await cur.execute(_BLOB_SQL, (thread_id, version)) + row = await cur.fetchone() + if row is None or row[1] is None or row[0] == "empty": + continue + value = loader.serde.loads_typed((row[0], bytes(row[1]))) + if isinstance(value, _DeltaSnapshot): + continue + crossing = path[i - 1] if i > 0 else cid + messages, _ = await loader.channel_value(thread_id, crossing) + return len(messages) + return 0 + + +def _trim(values: Mapping[str, Any], cut: int, window: Any, thread_id: str) -> dict[str, Any]: + """The middleware's trim of ``values`` before ``cut`` (``window.trim``), + keeping ``window``: the trimmed runs' carries come from the tip's own + messages, as the agent's trim would take them.""" + from ptc_agent.agent.middleware.compaction.notes import notes_window_carry + from ptc_agent.agent.middleware.compaction.window import trim + from ptc_agent.agent.middleware.subagent_switch import subagents_window_carry + + # The carries the main agent's build wires (``PTCAgent.create_agent``). + carries = [notes_window_carry(thread_id), subagents_window_carry] + return trim(values, cut, carries, head=window) + + +def _applied(values: Mapping[str, Any], update: Mapping[str, Any]) -> dict[str, Any]: + """``values`` with ``update`` applied, as the repaired tip should load.""" + from langgraph.types import Overwrite + + out = dict(values) + for key, value in update.items(): + out[key] = list(value.value) if isinstance(value, Overwrite) else value + return out + + +async def _write( + conn: psycopg.AsyncConnection, + saver: Any, + extra: Sequence[str], + thread_id: str, + tip: str, + cut: int, + window: Any, + kept_ids: Sequence[str], +) -> str: + """Run the trim as the tip's next step, leaving the old tip's pending + writes as they were, and return the checkpoint it leaves. Raises ``Skip`` + when that is not the tip's child. + + The loop stores a step's writes under the checkpoint it ran from, and a + DeltaChannel load replays every write under each ancestor, whichever + child it led to. Left there, the trim's ``Overwrite`` would also reach + any other child of the old tip, such as a maintenance write built on it + before the repair committed. The repaired tip does not need them: it + loads ``messages`` from its own snapshot, and its other channels from + their own blobs.""" + from langgraph.types import Command + + async def before_agent(state: Mapping[str, Any]) -> dict[str, Any]: + # The node trims the state the loop hands it, as the middleware's + # trim does; ``verify`` checks it was the state the plan read. A kept + # message stored without an id gets the one the plan's trim gave it, + # not a second fresh one. + for message, planned in zip(state[MESSAGES][cut:], kept_ids): + if message.id is None: + message.id = planned + return _trim(state, cut, window, thread_id) + + # No checkpoint id: the loop builds on the newest checkpoint, the tip (the + # plan checked it in this snapshot), as a step and not as a fork. + config = { + "configurable": {"thread_id": thread_id}, + "metadata": {"maintenance": "repair_crossed_threads"}, + } + graph = _graph(saver, extra, before_agent) + async with conn.cursor() as cur: + await cur.execute(_TIP_WRITES_SQL, (thread_id, tip)) + pending = await cur.fetchall() + await graph.ainvoke(Command(goto=NODE), config, durability="sync") + await cur.execute(_NEWEST_SQL, (thread_id,)) + new = (await cur.fetchone())[0] + await cur.execute( + "SELECT parent_checkpoint_id FROM checkpoints" + " WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s", + (thread_id, new), + ) + if new == tip or (await cur.fetchone())[0] != tip: + raise Skip("the trim did not land as the tip's child") + # The plan allowed no pending write on the tip but a failed task's. + await cur.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = ''" + " AND checkpoint_id = %s AND channel <> %s", + (thread_id, tip, _ERROR), + ) + await cur.execute(_TIP_WRITES_SQL, (thread_id, tip)) + if await cur.fetchall() != pending: + raise Skip("the trim left other writes on the old tip") + return new + + +def _compare(serde: Any, expected: Mapping[str, Any], after: Mapping[str, Any]) -> str | None: + """What of ``after`` differs from ``expected``, named by field, or None.""" + for key in sorted(set(expected) | set(after)): + if key == MESSAGES: + same = messages_digest(serde, list(expected.get(key) or [])) == messages_digest( + serde, list(after.get(key) or []) + ) + else: + same = _typed(serde, expected.get(key)) == _typed(serde, after.get(key)) + if not same: + return f"another {key}" + return None + + +def _typed(serde: Any, value: Any) -> tuple[str, bytes]: + if isinstance(value, set): + value = sorted(value) + return serde.dumps_typed(value) + + +async def _stores_snapshot( + conn: psycopg.AsyncConnection, thread_id: str, checkpoint_id: str +) -> bool: + """Whether ``checkpoint_id`` stores its own ``messages`` snapshot, so its + loads start there and not at the untrimmed list.""" + async with conn.cursor() as cur: + await cur.execute( + "SELECT checkpoint -> 'channel_versions' ->> 'messages'," + " (checkpoint -> 'channel_values' -> 'messages') IS NOT NULL" + " FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s", + (thread_id, checkpoint_id), + ) + version, flagged = await cur.fetchone() + await cur.execute(_BLOB_SQL, (thread_id, version)) + row = await cur.fetchone() + return bool(flagged) and row is not None and row[1] is not None + + +def _ended_at(anchors: list[Any], new: str) -> list[Any]: + """The turn anchors after the repair: the last turn now ends on ``new``.""" + if not anchors: + return anchors + last = replace(anchors[-1], tail_checkpoint_id=new, end_checkpoint_id=new) + return [*anchors[:-1], last] + + +def _verification_set(tree: Tree, anchors: list[Any], tip: str, opts: Options) -> list[str]: + """The checkpoints loaded before and after the write, in id order: the + tip, its parent, every turn anchor and branch tip, and a sample of the + rest. Only the tip gains rows the loads read (the step's writes), and no + load replays them but the repaired tip's, which starts at its own + snapshot.""" + if opts.verify_all: + return list(tree.order) + checks = {tip, *tree.leaves()} + parent = tree.nodes[tip].parent + if parent: + checks.add(parent) + for anchor in anchors: + for cid in (anchor.input_checkpoint_id, anchor.tail_checkpoint_id, anchor.end_checkpoint_id): + if cid in tree.nodes: + checks.add(cid) + rest = [cid for cid in tree.order if cid not in checks] + checks.update(random.Random(tree.order[0]).sample(rest, min(opts.sample, len(rest)))) + return [cid for cid in tree.order if cid in checks] + + +async def _reproject(thread_id: str) -> tuple[int, bool]: + """Project again the turns a read would project now, the last turn + among them (the checkpoint that ends it moved), as + ``backfill_turn_slices`` does; and whether each one was stored again + with the messages it held before.""" + from src.server.database.conversation import turn_slices as slices_db + from src.server.services.history import slices + from src.server.services.history.reader import CheckpointHistoryReader + from src.server.services.history.replay import ( + load_thread_inputs, + project_turns, + unserved_turns, + ) + + loaded = await load_thread_inputs(thread_id) + if loaded is None: + return 0, True + rows, branch_tip = loaded + unserved = await unserved_turns(rows, branch_tip) + if not unserved: + return 0, True + reader = CheckpointHistoryReader.get_instance() + anchors, _ = await reader.aget_turn_anchors(thread_id, branch_tip) + by_input = {a.input_checkpoint_id: a for a in anchors} + + async def spans(*, current: bool) -> dict[str, str]: + out = {} + for row in await slices_db.get_slices_at(thread_id, slices.SLICE_KEY, list(by_input)): + if current and not slices.matches(row, by_input[row.input_checkpoint_id]): + continue + span = slices.decoded(reader.serde, row) + if span is not None: + out[row.input_checkpoint_id] = messages_digest(reader.serde, span.messages) + return out + + held = await spans(current=False) + await project_turns(rows, branch_tip, turn_indexes=unserved) + # Only rows cut for where the branch ends its turns now: a store that + # fails is logged and skipped, which leaves the old row and its digest. + after = await spans(current=True) + return len(unserved), all(after.get(cid) == digest for cid, digest in held.items()) + + +# -------------------------------------------------------------------------- +# The requests the model is sent +# -------------------------------------------------------------------------- + +_ID = re.compile( + r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}|lc_run--[0-9a-f-]+" +) + + +def _known_ids(messages: Iterable[Any]) -> set[str]: + ids: set[str] = set() + for message in messages: + ids.add(str(message.id)) + ids.update(str(call.get("id")) for call in getattr(message, "tool_calls", None) or ()) + if getattr(message, "tool_call_id", None): + ids.add(str(message.tool_call_id)) + return ids + + +def _dump(messages: Sequence[Any], known: set[str], names: dict[str, str]) -> str: + """A request as text, every field the messages carry, with each id minted + while building it (a summary's) renamed by first appearance.""" + text = json.dumps([m.model_dump() for m in messages], sort_keys=True, default=str) + + def name(match: re.Match[str]) -> str: + value = match.group(0) + return value if value in known else names.setdefault(value, f"") + + return _ID.sub(name, text) + + +@functools.cache +def _recorder() -> type: + """A chat model that answers ``answer`` and records every request.""" + from langchain_core.language_models.chat_models import BaseChatModel + from langchain_core.messages import AIMessage + from langchain_core.outputs import ChatGeneration, ChatResult + from pydantic import Field + + class Recorder(BaseChatModel): + answer: str = "ok" + sent: list[list[Any]] = Field(default_factory=list) + + @property + def _llm_type(self) -> str: + return "recorder" + + def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult: + self.sent.append(list(messages)) + return ChatResult(generations=[ChatGeneration(message=AIMessage(self.answer))]) + + return Recorder + + +class _Mount: + """The transcript mount, as far as a request reads it: a save lands.""" + + async def save_transcript(self, target: Any, messages: Any, *, window: Any) -> bool: + return True + + +def _sandbox() -> Any: + mount = _Mount() + + async def settled_livefs(workspace_id: str | None = None) -> Any: + return mount + + async def aupload_files(files: Any) -> Any: + return SimpleNamespace(error=None) + + return SimpleNamespace( + livefs=mount, + settled_livefs=settled_livefs, + aupload_files=aupload_files, + normalize_path=lambda path: path, + ) + + +async def _requests(values: Mapping[str, Any], thread_id: str, settings: Settings) -> list[str]: + """The requests the main agent's compaction builds from ``values``: the + next call; the next call after an idle pause, with its Tier 1 cuts; and + a forced summary's request and the call after it. These read everything + the repair changes besides the head it drops: the window's count, the + turn numbering and the index text. Fakes stand in for the models and the + transcript mount, and the calls run as the thread's, so the paths they + name are its own.""" + from langchain_core.runnables.config import var_child_runnable_config + + from langsmith import tracing_context + + known = _known_ids(values.get(MESSAGES) or ()) + token = var_child_runnable_config.set( + {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + ) + try: + out: list[str] = [] + # The fakes see the whole conversation; a tracer on in the + # environment must not upload it as runs of its own. + with tracing_context(enabled=False): + for threshold, paused in ( + (float("inf"), False), + (float("inf"), True), + (0.0, False), + ): + names: dict[str, str] = {} + sent = await _sent(values, settings, threshold=threshold, paused=paused) + out += [_dump(messages, known, names) for messages in sent] + return out + finally: + var_child_runnable_config.reset(token) + + +async def _sent( + values: Mapping[str, Any], settings: Settings, *, threshold: float, paused: bool +) -> list[list[Any]]: + """What one model call from ``values`` sends: the summary model's request + when the context reaches ``threshold``, then the model's.""" + from langchain.agents.middleware.types import ModelRequest, ModelResponse + from langchain_core.messages import AIMessage + + from ptc_agent.agent.middleware.compaction.compact import Summarizer + from ptc_agent.agent.middleware.compaction.middleware import CompactionMiddleware + + recorder = _recorder() + model, summary = recorder(), recorder(answer="Summary.") + middleware = CompactionMiddleware( + # The whole span, so the comparison reads all of it. + Summarizer(summary, limit=10**12), + token_threshold=threshold, # type: ignore[arg-type] + keep_messages=settings.keep_messages, + offload=replace(settings.offload, idle_seconds=0.0), + backend=_sandbox(), + workspace_id="repair", + ) + state = dict(values) + if paused: + state["_last_model_response_at"] = 1.0 + state.update(await middleware._offload_idle(state)) + sent: list[list[Any]] = [] + + async def handler(request: ModelRequest) -> ModelResponse: + sent.append(list(request.messages)) + return ModelResponse(result=[AIMessage("ok")]) + + request = ModelRequest(model=model, messages=list(state.get(MESSAGES) or []), state=state) + await middleware.awrap_model_call(request, handler) + return [*summary.sent, *sent] + + +# -------------------------------------------------------------------------- +# CLI +# -------------------------------------------------------------------------- + + +def _outcome(out: Result[Report]) -> str: + if out.report is not None and out.report.refused: + return f"refused ({out.report.refused})" + if out.status is Status.NOTHING: + return "nothing to do" + if out.status is Status.SKIPPED: + return f"refused ({out.reason})" + if out.status is Status.PLANNED: + return "would repair" + if out.status is Status.REHEARSED: + return "rehearsed (rolled back)" + if out.status is Status.WRITTEN: + return "repaired" + (f" ({out.reason})" if out.reason else "") + return f"{out.status.value} ({out.reason})" + + +def _report(out: Result[Report]) -> None: + report = out.report + if report is None: + logger.info("thread=%s outcome=%s", out.thread_id, _outcome(out)) + return + span = f"{report.span[0]}-{report.span[1]}" if report.span else "-" + if report.crossed == "yes" and not report.runs: + logger.info( + "thread=%s crossed=yes messages=%d cut=%d differing=%d at %s (%s)", + out.thread_id, report.messages, report.cut, report.differing, span, report.detail, + ) + elif report.crossed == "yes": + logger.info( + "thread=%s crossed=yes messages=%d cut=%d runs=%d differing=%d at %s " + "of %d replayed index_lines=%s", + out.thread_id, report.messages, report.cut, report.runs, report.differing, + span, report.replayed, report.lines or "none", + ) + else: + logger.info( + "thread=%s crossed=%s messages=%d cut=%d%s", + out.thread_id, report.crossed, report.messages, report.cut, + f" ({report.detail})" if report.detail else "", + ) + if report.requests: + logger.info( + " requests=%d %s", report.requests, + {True: "identical", False: "DIFFER", None: "compared in the plan"}[report.requests_match], + ) + if out.checks: + mismatched = [c.checkpoint_id for c in out.checks if c.match is False] + logger.info(" verified=%d mismatched=%d", len(out.checks), len(mismatched)) + for cid in mismatched[:20]: + logger.info(" MISMATCH %s", cid) + if report.repaired_tip: + logger.info(" tip %s -> %s", report.tip, report.repaired_tip) + if out.committed_match is not None: + logger.info(" old tip committed load %s", "match" if out.committed_match else "MISMATCH") + if report.reprojected is not None: + logger.info( + " reprojected=%d turns, messages %s", report.reprojected, + "as stored" if report.reprojected_same else "DIFFER or not stored", + ) + logger.info(" outcome=%s", _outcome(out)) + + +async def _thread_ids(conninfo: str, args: argparse.Namespace) -> list[str]: + clauses, params = selection_sql(args) + clauses.append("t.latest_checkpoint_id IS NOT NULL") + async with await psycopg.AsyncConnection.connect(conninfo, autocommit=True) as conn: + await conn.execute( + "SELECT set_config('statement_timeout', %s, false)", + (f"{int(args.statement_timeout * 1000)}ms",), + ) + cur = await conn.execute(_THREADS_SQL.format(where=" AND ".join(clauses)), params) + return [r[0] for r in await cur.fetchall()] + + +async def run(args: argparse.Namespace, conninfo: str) -> int: + mode = Mode.APPLY if args.apply else Mode.REHEARSE if args.rehearse else Mode.DRY_RUN + opts = Options( + verify_all=args.verify == "all", + sample=args.sample, + timeouts=Timeouts(args.statement_timeout, args.lock_timeout), + ) + logger.info("target=%s mode=%s", describe(conninfo), mode.value) + refusal = unsupported_saver() + if refusal: + logger.error("refusing to run: %s", refusal) + return 2 + + from src.server.database import pool as db_pool + + app_target = describe(db_pool.get_db_connection_string()) + if app_target != describe(conninfo): + logger.error("the app's pool targets %s, not %s", app_target, describe(conninfo)) + return 2 + + statuses: Counter[str] = Counter() + failed = False + # Redis: projecting the last turn again probes its task streams. The + # agent config: what the compaction reads, and the key stored lines are + # kept under. + async with app_infra(redis=True, agent_config=True): + from ptc_agent.agent.middleware.compaction.types import OffloadSettings + from src.server.app import setup + + compaction = setup.agent_config.compaction + settings = Settings(compaction.keep_messages, OffloadSettings.from_config(compaction)) + thread_ids = await _thread_ids(conninfo, args) + logger.info("threads selected: %d", len(thread_ids)) + + async def one(thread_id: str) -> None: + nonlocal failed + out = await repair_thread(conninfo, thread_id, mode, opts, settings) + _report(out) + crossed = out.report.crossed if out.report else "?" + statuses[f"crossed={crossed}: {_outcome(out)}"] += 1 + # Committed, but its last turn did not project again as stored. + unprojected = ( + out.status is Status.WRITTEN + and out.report is not None + and out.report.reprojected_same is not True + ) + failed = failed or out.failed or unprojected + + tally = await for_each_thread( + thread_ids, one, concurrency=args.concurrency, pause=args.pause, + total=len(thread_ids), + ) + logger.info("done mode=%s threads=%d %s", mode.value, tally.done, dict(statuses) or "") + return 1 if failed or tally.failed else 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=(__doc__ or "").split("\n")[0]) + action = parser.add_mutually_exclusive_group() + action.add_argument("--apply", action="store_true", help="write and commit verified repairs") + action.add_argument( + "--rehearse", action="store_true", + help="write and verify each repair, then roll it back", + ) + add_selection_args(parser, concurrency=1, pause=0.2) + parser.add_argument( + "--verify", choices=("affected", "all"), default="affected", + help="load the tip, turn anchors and a sample before and after, or every checkpoint", + ) + parser.add_argument( + "--sample", type=int, default=20, + help="other checkpoints also loaded per thread under --verify affected", + ) + parser.add_argument("--dsn", help="database URI (default: from DB_*)") + parser.add_argument( + "--statement-timeout", type=float, default=120.0, metavar="SECONDS", + help="per-statement timeout for the listing and inside each thread's transaction", + ) + parser.add_argument( + "--lock-timeout", type=float, default=5.0, metavar="SECONDS", + help="how long a write waits on a locked row before skipping the thread", + ) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + conninfo = args.dsn or build_db_uri("DB_") + point_app_pool_at(conninfo) + return asyncio.run(run(args, conninfo)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/ops/thin_message_snapshots.py b/scripts/ops/thin_message_snapshots.py new file mode 100644 index 000000000..b34ed0550 --- /dev/null +++ b/scripts/ops/thin_message_snapshots.py @@ -0,0 +1,689 @@ +"""Reclaim checkpoint storage by thinning old ``messages`` snapshots, keeping every checkpoint. + +``messages`` is a DeltaChannel. A checkpoint loads it one of two ways. If a +``messages`` blob is stored at the checkpoint's channel version, the load reads +it directly: a snapshot, or the full list a checkpoint written before +DeltaChannel stored at every step. Otherwise the load walks parent checkpoints +to the nearest one whose ``channel_values`` flags a snapshot and replays the +``messages`` writes in between. Every blob holds the whole list, so on a long +thread blobs are most of its checkpoint bytes. + +This script deletes most of them. Every checkpoint, write and parent link +stays, so branches, forks and edits are untouched; only where loads start +their replay moves. Kept, per thread: + +- what every branch tip and the thread's current tip load directly or seed + from, so the next turn loads exactly as it does today; +- the same for the anchors of the newest ``--keep-recent-turns`` turns, where + retry, regenerate and edit start; +- a snapshot at least every ``--spacing`` messages updates along a branch, + which bounds what a load of an old checkpoint (an edit into an old turn) + replays; +- every blob whose loads a replay from the next one up would not reproduce + byte for byte. Writes stored before DeltaChannel replay with id-less human + inputs, and parallel tool results replay in a different order than they were + applied, so this keeps more than spacing alone would. + +A kept pre-DeltaChannel list is flagged as a snapshot, which it already loads +as, so walks stop at it instead of replaying from the root. A load that +already walks past such lists to the root today keeps its value: the topmost +checkpoint of each such walk is first seeded with what it loads now. + +``--keep all`` drops nothing and does only that: every stored pre-DeltaChannel +list is flagged and every walk that crossed one seeded. Run it on every thread +before the snapshot cadence is lengthened, or a migrated thread's first delta +turn replays the id-less writes from its root and freezes that list into +state. Nothing gets slower, so it does not wait on stored lines. + +A thread is thinned only while its history no longer reads old states: +every turn on its branch streams from a stored slice and lines under the +current keys, and no run is in progress. Task namespaces (``task:...``) are +not touched. + +Each thread is one REPEATABLE READ transaction. It plans from the stored rows +and loads, through the server's own reader, the checkpoints that load a +changed blob or are newly flagged, their children, every branch tip, every +turn anchor and a random sample; deeper walks rest on the planner's own +byte-for-byte check of each dropped blob's loads. ``--verify all`` loads every +checkpoint instead. It then writes, loads them again, and rolls back on any +byte of difference. The tips' loads may get shorter, never longer. + +The script edits the Postgres saver's storage format, so it refuses to run +against releases it was not verified with; re-verify it with ``--rehearse +--verify all`` on a copy before adding a release. + +Usage (from the repo root, inside the backend container): + uv run python scripts/ops/thin_message_snapshots.py # plan only + uv run python scripts/ops/thin_message_snapshots.py --rehearse # write, verify, roll back + uv run python scripts/ops/thin_message_snapshots.py --apply # write, verify, commit + uv run python scripts/ops/thin_message_snapshots.py --keep all --apply + uv run python scripts/ops/thin_message_snapshots.py --thread ID --verify all + +Connects with DB_*, the variables the server builds its pools from, or with +--dsn; the coverage reads go through the app's pool, which is pointed at the +same database. Idempotent: a second run finds nothing left to do. Safe beside +live turns: a turn only adds rows, and a thread whose rows change under it is +skipped. Prints ids, counts and sizes only, never message content. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import random +import sys +import uuid +from collections import Counter, defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import psycopg + +# Run as a script, sys.path[0] is scripts/ops, not the repo root the sibling +# helpers and the app are imported from. +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from scripts._errors import where # noqa: E402 +from scripts.ops._db import build_db_uri, point_app_pool_at # noqa: E402 +from scripts.ops._delta_loads import ( # noqa: E402 + MESSAGES, + Check, + Loader, + Mode, + Result, + Rewrite, + Skip, + Status, + Timeouts, + Tree, + describe, + read_tree, + rewrite_thread, + unsupported_saver, +) +from scripts.ops._thin_plan import ( # noqa: E402 + Plan, + Replayer, + Unit, + Writes, + longest_walk, + make_plan, +) + +logger = logging.getLogger("thin_message_snapshots") + +_BLOB_SQL = """ +SELECT type, blob FROM checkpoint_blobs +WHERE thread_id = %s AND checkpoint_ns = '' AND channel = 'messages' AND version = %s +""" + +_WRITES_SQL = """ +SELECT checkpoint_id, task_id, idx, type, blob FROM checkpoint_writes +WHERE thread_id = %s AND checkpoint_ns = '' AND channel = 'messages' +ORDER BY checkpoint_id, task_id, idx +""" + +_WRITE_COUNTS_SQL = """ +SELECT checkpoint_id, count(*) FROM checkpoint_writes +WHERE thread_id = %s AND checkpoint_ns = '' AND channel = 'messages' +GROUP BY checkpoint_id +""" + +_THREAD_SQL = """ +SELECT latest_checkpoint_id FROM conversation_threads WHERE conversation_thread_id = %s::uuid +""" + +_LIVE_SQL = """ +SELECT EXISTS ( + SELECT 1 FROM conversation_responses + WHERE conversation_thread_id = %(thread_id)s::uuid AND status = 'in_progress') + OR EXISTS ( + SELECT 1 FROM subagent_runs + WHERE thread_id = %(thread_id)s::uuid AND status = 'in_progress') +""" + +_CANDIDATES_SQL = """ +SELECT thread_id, count(*) AS blobs, sum(length(blob)) AS bytes FROM checkpoint_blobs +WHERE checkpoint_ns = '' AND channel = 'messages' AND type <> 'empty' AND blob IS NOT NULL +GROUP BY thread_id +HAVING count(*) > 1 AND sum(length(blob)) >= %s +ORDER BY bytes DESC +""" + +# Under --keep all only a list no checkpoint flags yet can change anything. +_UNFLAGGED_LISTS_SQL = """ +WITH snapshots AS ( + SELECT DISTINCT thread_id, checkpoint -> 'channel_versions' ->> 'messages' AS version + FROM checkpoints + WHERE checkpoint_ns = '' AND checkpoint -> 'channel_values' -> 'messages' IS NOT NULL +) +SELECT b.thread_id, count(*) AS blobs, sum(length(b.blob)) AS bytes FROM checkpoint_blobs b +LEFT JOIN snapshots s ON s.thread_id = b.thread_id AND s.version = b.version +WHERE b.checkpoint_ns = '' AND b.channel = 'messages' AND b.type <> 'empty' + AND b.blob IS NOT NULL AND s.version IS NULL +GROUP BY b.thread_id +HAVING sum(length(b.blob)) >= %s +ORDER BY bytes DESC +""" + +# Each row must still be the version this transaction read: xmin moves on any +# update, and under REPEATABLE READ a concurrent one fails the statement. +_FLAG_SQL = """ +UPDATE checkpoints +SET checkpoint = jsonb_set(checkpoint, '{channel_values,messages}', 'true'::jsonb) +WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s + AND xmin = %s::xid + AND checkpoint -> 'channel_values' -> 'messages' IS NULL + AND checkpoint -> 'channel_versions' ->> 'messages' = %s +""" + +_UNFLAG_SQL = """ +UPDATE checkpoints +SET checkpoint = checkpoint #- '{channel_values,messages}' +WHERE thread_id = %s AND checkpoint_ns = '' AND checkpoint_id = %s + AND xmin = %s::xid + AND checkpoint -> 'channel_values' -> 'messages' IS NOT NULL + AND checkpoint -> 'channel_versions' ->> 'messages' = %s +""" + +_DELETE_BLOB_SQL = """ +DELETE FROM checkpoint_blobs +WHERE thread_id = %s AND checkpoint_ns = '' AND channel = 'messages' AND version = %s + AND type <> 'empty' AND blob IS NOT NULL +RETURNING length(blob) +""" + +# Same row shape as the saver's own snapshot put. Only an empty placeholder may +# be replaced; a version with a stored value never gets a seed. +_SEED_SQL = """ +INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob) +VALUES (%s, '', 'messages', %s, %s, %s) +ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO UPDATE + SET type = EXCLUDED.type, blob = EXCLUDED.blob + WHERE checkpoint_blobs.type = 'empty' OR checkpoint_blobs.blob IS NULL +""" + + +# -------------------------------------------------------------------------- +# One thread +# -------------------------------------------------------------------------- + + +@dataclass +class Options: + spacing: int = 1000 + keep_recent_turns: int = 2 + sample: int = 20 + verify_all: bool = False + keep_all: bool = False + timeouts: Timeouts = field(default_factory=Timeouts) + + +@dataclass +class Report: + """What the plan does to one thread, and what verifying it found.""" + + checkpoints: int = 0 + blobs: int = 0 + blob_bytes: int = 0 + dropped: int = 0 + dropped_bytes: int = 0 + promoted: int = 0 + seeded: int = 0 + seed_bytes: int = 0 + kept: Counter[str] = field(default_factory=Counter) + walk_before: tuple[int, int] = (0, 0) + walk_after: tuple[int, int] = (0, 0) + longer: list[str] = field(default_factory=list) + anchors_match: bool | None = None + + @property + def reclaimed(self) -> int: + return self.dropped_bytes - self.seed_bytes + + +@dataclass(frozen=True) +class Candidate: + tip: str + + +@dataclass(frozen=True) +class NotCandidate: + reason: str + unserved: int = 0 # turns lacking current stored lines + + +Candidacy = Candidate | NotCandidate + + +async def preconditions( + conn: psycopg.AsyncConnection, thread_id: str, loader: Loader, *, coverage: bool = True +) -> Candidacy: + """Whether the thread may be thinned now, from its current tip. + + Thinning makes old states slower to load, so it waits until nothing + routine loads them: every turn on the branch streams from its stored + slice and lines (``coverage``), and no run is mid-flight. + """ + try: + uuid.UUID(thread_id) + except ValueError: + return NotCandidate("not a conversation thread") + async with conn.cursor() as cur: + await cur.execute(_THREAD_SQL, (thread_id,)) + row = await cur.fetchone() + if row is None: + return NotCandidate("no conversation thread row") + tip = row[0] + if not tip: + return NotCandidate("thread has no current checkpoint") + await cur.execute(_LIVE_SQL, {"thread_id": thread_id}) + if (await cur.fetchone())[0]: + return NotCandidate("a run is in progress") + if not coverage: + return Candidate(tip) + + from src.server.services.history.replay import ( + CheckpointReplayUnavailable, + load_thread_inputs, + unserved_turns, + ) + + # Judged from the tip read above, which is the one the candidacy carries. + loaded = await load_thread_inputs(thread_id) + if loaded is None: + return NotCandidate("no conversation thread row") + try: + unserved = await unserved_turns(loaded[0], tip, reader=loader.reader) + except CheckpointReplayUnavailable: + return NotCandidate("turns do not pair with checkpoints") + if unserved is None: + return NotCandidate("no turns on the branch") + if unserved: + return NotCandidate("turns without current stored lines", len(unserved)) + return Candidate(tip) + + +async def thin_thread( + conninfo: str, thread_id: str, tip_id: str, mode: Mode, opts: Options +) -> Result[Report]: + """Thin one thread in its own verified transaction (``rewrite_thread``). + + ``tip_id`` is the current tip its candidacy was judged from; the thread is + skipped if it has moved on. + """ + + async def plan(conn: psycopg.AsyncConnection, loader: Loader) -> Rewrite[Report]: + return await _plan(conn, loader, thread_id, tip_id, opts) + + return await rewrite_thread(conninfo, thread_id, mode, opts.timeouts, plan) + + +async def _plan( + conn: psycopg.AsyncConnection, loader: Loader, thread_id: str, tip_id: str, opts: Options +) -> Rewrite[Report]: + async with conn.cursor() as cur: + await cur.execute(_THREAD_SQL, (thread_id,)) + row = await cur.fetchone() + if row is None or row[0] != tip_id: + raise Skip("the thread moved on since its candidacy was checked") + await cur.execute(_LIVE_SQL, {"thread_id": thread_id}) + if (await cur.fetchone())[0]: + raise Skip("a run is in progress") + + tree, sizes = await read_tree(conn, thread_id) + if tree.problem: + raise Skip(tree.problem) + if tip_id not in tree.nodes: + raise Skip("the current tip is not a root checkpoint") + if any(n.flagged and n.version not in sizes for n in tree.nodes.values()): + # Such a load reads the flag itself as the value; leave it be. + raise Skip("a flagged checkpoint has no stored snapshot") + units = { + version: Unit( + version, + size, + list(tree.by_version[version]), + [c for c in tree.by_version[version] if tree.nodes[c].flagged], + ) + for version, size in sizes.items() + if tree.by_version.get(version) + } + report = Report(len(tree.nodes), len(units), sum(u.size for u in units.values())) + + async def read_blob(version: str) -> tuple[str, bytes] | None: + async with conn.cursor() as cur: + await cur.execute(_BLOB_SQL, (thread_id, version)) + row = await cur.fetchone() + return None if row is None else (row[0], bytes(row[1])) + + async def read_writes() -> Writes: + return await _read_writes(conn, thread_id, loader.serde) + + replay = Replayer(thread_id, tree, loader, read_blob, read_writes) + anchors, _ = await loader.reader.aget_turn_anchors(thread_id, tip_id) + protected: dict[str, str] = {cid: "branch tip" for cid in tree.leaves()} + protected[tip_id] = "current tip" + for anchor in anchors[-opts.keep_recent_turns :] if opts.keep_recent_turns else []: + for cid in (anchor.input_checkpoint_id, anchor.tail_checkpoint_id, anchor.end_checkpoint_id): + if cid in tree.nodes: + protected.setdefault(cid, "recent turn") + + plan = await make_plan(tree, units, replay, protected, None if opts.keep_all else opts.spacing) + report.kept = plan.kept + report.dropped = len(plan.drop) + report.dropped_bytes = sum(u.size for u in plan.drop) + report.promoted = len(plan.promote) + report.seeded = len(plan.seeds) + report.seed_bytes = sum(s.size for s in plan.seeds) + if not plan.changes: + return Rewrite(report) + counts = await _write_counts(conn, thread_id) + report.walk_before = longest_walk(tree, plan.before, counts) + report.walk_after = longest_walk(tree, plan.after, counts) + + async def write() -> None: + report.dropped_bytes = await _write(conn, thread_id, tree, plan) + + async def verify(checks: dict[str, Check]) -> str | None: + report.longer = [ + cid for cid, check in checks.items() + if cid in protected and check.after and not check.after.no_longer_than(check.before) + ] + anchors_after, _ = await loader.reader.aget_turn_anchors(thread_id, tip_id) + report.anchors_match = anchors_after == anchors + if report.longer: + return "a tip's walk got longer" + return None if report.anchors_match else "the turn anchors changed" + + checks = _verification_set(tree, plan, anchors, protected, tip_id, opts) + return Rewrite(report, checks, tip_id, write, verify) + + +async def _read_writes(conn: psycopg.AsyncConnection, thread_id: str, serde: Any) -> Writes: + """Each checkpoint's ``messages`` writes as the saver replays them: + ordered by task id, then index.""" + writes: Writes = defaultdict(list) + async with conn.cursor() as cur: + await cur.execute(_WRITES_SQL, (thread_id,)) + async for cid, task_id, _idx, tag, blob in cur: + writes[cid].append((task_id, MESSAGES, serde.loads_typed((tag, bytes(blob))))) + return writes + + +async def _write_counts(conn: psycopg.AsyncConnection, thread_id: str) -> dict[str, int]: + async with conn.cursor() as cur: + await cur.execute(_WRITE_COUNTS_SQL, (thread_id,)) + return {cid: n for cid, n in await cur.fetchall()} + + +def _verification_set( + tree: Tree, + plan: Plan, + anchors: list[Any], + protected: dict[str, str], + tip_id: str, + opts: Options, +) -> list[str]: + """The checkpoints loaded before and after the write, in id order: those + loading a changed blob or newly flagged, their children, every tip and + turn anchor, and a sample of the rest.""" + if opts.verify_all: + return list(tree.order) + changed: set[str] = set(plan.flags) + for unit in [*plan.drop, *plan.promote]: + changed.update(unit.loaders) + for seed in plan.seeds: + changed.update(tree.by_version[seed.version]) + checks = set(changed) + for cid in changed: + checks.update(tree.children.get(cid, ())) + checks.update(protected) + checks.add(tip_id) + for anchor in anchors: + for cid in (anchor.input_checkpoint_id, anchor.tail_checkpoint_id, anchor.end_checkpoint_id): + if cid in tree.nodes: + checks.add(cid) + rest = [cid for cid in tree.order if cid not in checks] + checks.update(random.Random(tree.order[0]).sample(rest, min(opts.sample, len(rest)))) + return [cid for cid in tree.order if cid in checks] + + +async def _write(conn: psycopg.AsyncConnection, thread_id: str, tree: Tree, plan: Plan) -> int: + """Store the seeds, flag what stays, unflag and delete what goes. Returns + the bytes deleted; raises ``Skip`` when a row is no longer what was read.""" + async with conn.cursor() as cur: + for seed in plan.seeds: + await cur.execute(_SEED_SQL, (thread_id, seed.version, *seed.typed)) + if cur.rowcount != 1: + raise Skip("a stored messages blob appeared at a seed's version") + for cid in plan.flags: + node = tree.nodes[cid] + await cur.execute(_FLAG_SQL, (thread_id, cid, node.xmin, node.version)) + if cur.rowcount != 1: + raise Skip("a checkpoint changed since it was read") + deleted = 0 + for unit in plan.drop: + for cid in unit.flagged: + node = tree.nodes[cid] + await cur.execute(_UNFLAG_SQL, (thread_id, cid, node.xmin, node.version)) + if cur.rowcount != 1: + raise Skip("a checkpoint changed since it was read") + await cur.execute(_DELETE_BLOB_SQL, (thread_id, unit.version)) + row = await cur.fetchone() + if row is None: + raise Skip("a messages blob changed since it was read") + deleted += row[0] + return deleted + + +# -------------------------------------------------------------------------- +# CLI +# -------------------------------------------------------------------------- + + +def _mb(n: int) -> str: + return f"{n / 1e6:.1f}MB" + + +def _report(out: Result[Report]) -> None: + report = out.report + if report is None: + logger.info("thread=%s", out.thread_id) + else: + logger.info( + "thread=%s checkpoints=%d blobs=%d (%s) dropped=%d (%s) flagged_kept=%d " + "seeded=%d (%s) reclaimed=%s kept=%s", + out.thread_id, report.checkpoints, report.blobs, _mb(report.blob_bytes), + report.dropped, _mb(report.dropped_bytes), report.promoted, report.seeded, + _mb(report.seed_bytes), _mb(report.reclaimed), dict(report.kept) or "{}", + ) + if out.status is not Status.NOTHING: + logger.info( + " longest walk (checkpoints, writes): before=%s after=%s", + report.walk_before, report.walk_after, + ) + tip = out.committed + if report is not None and tip is not None: + mismatched = [c.checkpoint_id for c in out.checks if c.match is False] + logger.info( + " verified=%d mismatched=%d longer_tip_walks=%d anchors=%s", + len(out.checks), len(mismatched), len(report.longer), + "same" if report.anchors_match else "DIFFER", + ) + for cid in mismatched[:20]: + logger.info(" MISMATCH %s", cid) + for cid in report.longer[:20]: + logger.info(" LONGER WALK %s", cid) + logger.info(" current tip load before: %s | after: %s", tip.before, tip.after) + if out.committed_match is not None: + logger.info(" committed tip load %s", "match" if out.committed_match else "MISMATCH") + logger.info(" outcome=%s%s", out.status.value, f" ({out.reason})" if out.reason else "") + + +async def run(args: argparse.Namespace, conninfo: str) -> int: + mode = Mode.APPLY if args.apply else Mode.REHEARSE if args.rehearse else Mode.DRY_RUN + opts = Options( + spacing=args.spacing, + keep_recent_turns=args.keep_recent_turns, + sample=args.sample, + verify_all=args.verify == "all", + keep_all=args.keep == "all", + timeouts=Timeouts(args.statement_timeout, args.lock_timeout), + ) + logger.info("target=%s mode=%s keep=%s", describe(conninfo), mode.value, args.keep) + refusal = unsupported_saver() + if refusal: + logger.error("refusing to run: %s", refusal) + return 2 + + app_pool = None + if not opts.keep_all: + # Coverage reads go through the app's pool. + from src.server.database import pool as db_pool + + app_target = describe(db_pool.get_db_connection_string()) + if app_target != describe(conninfo): + logger.error("the app's pool targets %s, not %s", app_target, describe(conninfo)) + return 2 + app_pool = db_pool.get_or_create_pool() + statuses: Counter[str] = Counter() + totals: Counter[str] = Counter() + failed = False + try: + if app_pool is not None: + await app_pool.open() + # Stored lines are keyed by the config the server reads (the + # compaction threshold), so coverage is judged under it too. + from ptc_agent.config import ConfigContext, load_from_files + from src.server.app import setup + + setup.agent_config = await load_from_files(context=ConfigContext.SDK) + async with await psycopg.AsyncConnection.connect(conninfo, autocommit=True) as conn: + await conn.execute( + "SELECT set_config('statement_timeout', %s, false)", + (f"{int(args.statement_timeout * 1000)}ms",), + ) + if args.thread: + thread_ids = list(args.thread) + else: + listing = _UNFLAGGED_LISTS_SQL if opts.keep_all else _CANDIDATES_SQL + try: + async with conn.cursor() as cur: + await cur.execute(listing, (args.min_bytes,)) + thread_ids = [r[0] for r in await cur.fetchall()] + except psycopg.errors.QueryCanceled: + logger.error( + "listing threads hit the statement timeout (%ss); raise " + "--statement-timeout or pass --thread", args.statement_timeout, + ) + return 1 + logger.info( + "threads with %s: %d", + "a messages list no checkpoint flags" if opts.keep_all + else "more than one messages blob", + len(thread_ids), + ) + + coverage = Loader(conn) + processed = 0 + for thread_id in thread_ids: + if args.limit is not None and processed >= args.limit: + break + try: + candidacy = await preconditions( + conn, thread_id, coverage, coverage=not opts.keep_all + ) + except Exception as exc: + candidacy = NotCandidate(f"precondition check failed: {where(exc)}") + if isinstance(candidacy, NotCandidate): + statuses[f"not a candidate: {candidacy.reason}"] += 1 + if args.thread or args.verbose: + count = f" ({candidacy.unserved})" if candidacy.unserved else "" + logger.info( + "thread=%s not a candidate: %s%s", thread_id, candidacy.reason, count + ) + continue + processed += 1 + out = await thin_thread(conninfo, thread_id, candidacy.tip, mode, opts) + _report(out) + statuses[out.status.value if not out.reason else f"{out.status.value}: {out.reason}"] += 1 + if out.report is not None and out.status in ( + Status.PLANNED, Status.REHEARSED, Status.WRITTEN + ): + totals["dropped_bytes"] += out.report.dropped_bytes + totals["seed_bytes"] += out.report.seed_bytes + totals["threads"] += 1 + failed = failed or out.failed + finally: + if app_pool is not None: + await app_pool.close() + + logger.info( + "done mode=%s threads=%d dropped=%s seeded=%s reclaimed=%s %s", + mode.value, totals["threads"], _mb(totals["dropped_bytes"]), + _mb(totals["seed_bytes"]), _mb(totals["dropped_bytes"] - totals["seed_bytes"]), + dict(statuses) or "", + ) + return 1 if failed else 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=(__doc__ or "").split("\n")[0]) + action = parser.add_mutually_exclusive_group() + action.add_argument("--apply", action="store_true", help="write and commit verified changes") + action.add_argument( + "--rehearse", action="store_true", + help="write and verify each thread, then roll it back", + ) + parser.add_argument( + "--thread", action="append", metavar="ID", + help="limit to this thread (repeatable)", + ) + parser.add_argument("--limit", type=int, help="process at most N candidate threads") + parser.add_argument( + "--keep", choices=("spaced", "all"), default="spaced", + help="spaced: drop what the tips, recent turns and --spacing do not need; " + "all: drop nothing, only flag lists stored before DeltaChannel", + ) + parser.add_argument( + "--min-bytes", type=int, default=0, + help="only threads whose messages blobs hold at least this many bytes", + ) + parser.add_argument( + "--spacing", type=int, default=1000, + help="keep a snapshot at least every N messages updates along a branch", + ) + parser.add_argument( + "--keep-recent-turns", type=int, default=2, + help="keep what the anchors of the newest N turns load", + ) + parser.add_argument( + "--verify", choices=("affected", "all"), default="affected", + help="load the checkpoints the change can reach, or every checkpoint", + ) + parser.add_argument( + "--sample", type=int, default=20, + help="unaffected checkpoints also loaded per thread under --verify affected", + ) + parser.add_argument("--dsn", help="database URI (default: from DB_*)") + parser.add_argument( + "--statement-timeout", type=float, default=120.0, metavar="SECONDS", + help="per-statement timeout for the scan and inside each thread's transaction", + ) + parser.add_argument( + "--lock-timeout", type=float, default=5.0, metavar="SECONDS", + help="how long a write waits on a locked checkpoint row before skipping the thread", + ) + parser.add_argument("-v", "--verbose", action="store_true", help="list every non-candidate") + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + conninfo = args.dsn or build_db_uri("DB_") + point_app_pool_at(conninfo) + return asyncio.run(run(args, conninfo)) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/utils/_thread_job.py b/scripts/utils/_thread_job.py new file mode 100644 index 000000000..5131ddabc --- /dev/null +++ b/scripts/utils/_thread_job.py @@ -0,0 +1,163 @@ +"""What the thread jobs under ``scripts/utils`` share: the app's infrastructure, +opened and always closed; the thread selection flags; and a paced worker loop +that reports a failure without its message. + +The importing script puts the repo root and ``src`` on ``sys.path`` first. +""" + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import os +import time +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable +from dataclasses import dataclass +from typing import Any, TypeVar + +from scripts._errors import where + +T = TypeVar("T") + + +@contextlib.asynccontextmanager +async def app_infra(*, redis: bool, agent_config: bool) -> AsyncIterator[None]: + """The app's pool and checkpointer, plus Redis and the agent config when + the job needs them, closed on the way out however it ends: an open + pool's workers outlive ``main()``, and ``asyncio.run()`` can wait on + them at exit forever. Each close is registered before its open, so a + failed open still closes what came up before it.""" + from src.server.app import setup + from src.server.database import pool as db_pool + from src.server.utils.checkpointer import ( + close_checkpointer_pool, + get_checkpointer, + open_checkpointer_pool, + ) + + async with contextlib.AsyncExitStack() as stack: + pool = db_pool.get_or_create_pool() + stack.push_async_callback(pool.close) + await pool.open() + if redis: + from src.utils.cache.redis_cache import close_cache, init_cache + + stack.push_async_callback(close_cache) + await init_cache() + checkpointer = get_checkpointer( + "postgres", + db_host=os.getenv("DB_HOST", "localhost"), + db_port=int(os.getenv("DB_PORT", "5432")), + db_name=os.getenv("DB_NAME", "postgres"), + db_user=os.getenv("DB_USER", "postgres"), + db_password=os.getenv("DB_PASSWORD", "postgres"), + ) + stack.push_async_callback(close_checkpointer_pool, checkpointer) + await open_checkpointer_pool(checkpointer) + setup.checkpointer = checkpointer + if agent_config: + from ptc_agent.config import ConfigContext, load_from_files + + setup.agent_config = await load_from_files(context=ConfigContext.SDK) + yield + + +# -------------------------------------------------------------------------- +# Selection +# -------------------------------------------------------------------------- + + +def add_selection_args(parser: argparse.ArgumentParser, *, concurrency: int, pause: float) -> None: + parser.add_argument("--days", type=int, help="only threads active in the last N days") + parser.add_argument("--user", help="only this user's threads") + parser.add_argument("--workspace", help="only this workspace's threads") + parser.add_argument("--thread", action="append", help="only these threads") + parser.add_argument("--concurrency", type=int, default=concurrency) + parser.add_argument("--pause", type=float, default=pause, help="seconds between threads") + + +def selection_sql(args: argparse.Namespace) -> tuple[list[str], dict[str, Any]]: + """WHERE clauses and their named parameters for the selection flags, over + ``conversation_threads t JOIN workspaces w``.""" + clauses: list[str] = [] + params: dict[str, Any] = {} + if args.thread: + clauses.append("t.conversation_thread_id = ANY(%(threads)s::uuid[])") + params["threads"] = args.thread + if args.workspace: + clauses.append("t.workspace_id = %(workspace)s::uuid") + params["workspace"] = args.workspace + if args.user: + clauses.append("w.user_id = %(user)s") + params["user"] = args.user + if args.days: + clauses.append("t.updated_at >= NOW() - make_interval(days => %(days)s)") + params["days"] = args.days + return clauses, params + + +# -------------------------------------------------------------------------- +# The worker loop +# -------------------------------------------------------------------------- + + +@dataclass +class Tally: + done: int = 0 + unavailable: int = 0 # replays from stored events instead; nothing to do + failed: int = 0 + + +async def for_each_thread( + items: Iterable[T] | AsyncIterable[T], + fn: Callable[[T], Awaitable[None]], + *, + concurrency: int, + pause: float, + total: int | None = None, + name: Callable[[T], str] = str, +) -> Tally: + """Run ``fn`` on each item, ``concurrency`` at a time with ``pause`` + seconds after each. An async source is drawn from only as workers free + up, so a job can recheck each page of candidates as it comes up.""" + from src.server.services.history.replay import CheckpointReplayUnavailable + + source = _aiter(items) + draw = asyncio.Lock() + tally = Tally() + began = time.monotonic() + + async def worker() -> None: + while True: + async with draw: + try: + item = await anext(source) + except StopAsyncIteration: + return + try: + await fn(item) + except CheckpointReplayUnavailable as exc: + tally.unavailable += 1 + print(f"{name(item)} unavailable: {where(exc)}") + except Exception as exc: + tally.failed += 1 + print(f"{name(item)} failed: {where(exc)}") + tally.done += 1 + if tally.done % 25 == 0: + rate = tally.done / max(time.monotonic() - began, 1e-6) + of = f"/{total}" if total is not None else "" + print(f"{tally.done}{of} threads, {tally.failed} failed, {rate:.2f}/s") + await asyncio.sleep(pause) + + await asyncio.gather(*(worker() for _ in range(max(1, concurrency)))) + return tally + + +async def _aiter(items: Iterable[T] | AsyncIterable[T]) -> AsyncIterator[T]: + if isinstance(items, AsyncIterable): + async for item in items: + yield item + else: + for item in items: + yield item diff --git a/scripts/utils/backfill_replay_facts.py b/scripts/utils/backfill_replay_facts.py new file mode 100644 index 000000000..26a423b7f --- /dev/null +++ b/scripts/utils/backfill_replay_facts.py @@ -0,0 +1,460 @@ +"""Backfill legacy replay facts, so replay stops reading stored events. + +A turn settled while the live stream was archived replays some rows only that +archive kept: widget data, evicted tool results, reasoning durations, signal +rows in their streamed positions, and the main lane of a turn whose images +never resolved (``replay.legacy``). Replay derives them from ``sse_events`` +on every projection of a row without them. This derives them once, into the +row's ``replay_facts.legacy``, after which replay never reads the row's +stored events. It also stamps ``metadata.credit_pause_message`` on credit +pauses that predate the finalize stamping it, which settling an automation +otherwise reads from the stored events. + +Per thread, newest first, with a pause between threads: + +1. every pending turn is projected afresh from its stored events, and what + they gave it is kept, as the row will hold it, with the row version it + came from; those facts are placed into the same projection again, and a + turn whose rows they place differently is reported and left pending + (it keeps replaying from its stored events; never expected); +2. each row's facts are written with a jsonb merge that keeps every other + key, and only while the row is still that version (a row written + meanwhile is left for the next run); +3. the written turns are projected again, from their facts, so their + stored lines are current. + +A turn whose runs still write, or with no checkpoint boundary but a +completed row (the thread then replays from its stored events whole) or +stored events of its own, is left pending; one with neither is marked known +empty. Safe to stop and rerun: only checked facts are written, a +backfilled row is not derived again, and a write to a row's stored events +takes its facts away until the next run. + +Run inside the backend container (needs the app's env and venv); a +projection reads checkpoints a batch of turns at a time, so keep the +concurrency low: + + /app/.venv/bin/python scripts/utils/backfill_replay_facts.py # count only + /app/.venv/bin/python scripts/utils/backfill_replay_facts.py --apply + /app/.venv/bin/python scripts/utils/backfill_replay_facts.py --apply --thread + +Prints ids, counts and sizes only, never content. +""" + +from __future__ import annotations + +import argparse +import asyncio +import copy +import json +import os +import sys +import time +from collections import Counter +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from psycopg.rows import dict_row + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "src"))) + +from scripts.utils._thread_job import ( # noqa: E402 + add_selection_args, + app_infra, + for_each_thread, + selection_sql, +) +from src.server.database.conversation import settled_attempts # noqa: E402 +from src.server.database.replay_facts import LEGACY_KEY, LEGACY_VERSION_SQL # noqa: E402 + + +def _version() -> str: + from src.server.services.history.replay import legacy + + return str(legacy.FACTS_VERSION) + + +# The row each turn replays (``settled_attempts``): a superseded attempt is +# never replayed, so it needs no facts. +_REPLAY_ROWS = settled_attempts( + "conversation_response_id, turn_index, status, replay_facts, " + "sse_events IS NOT NULL AS has_events, " + "COALESCE(pg_column_size(sse_events), 0) AS events_bytes", + "t.conversation_thread_id", +) + + +async def _pending_threads(args: argparse.Namespace) -> list[dict[str, Any]]: + """Threads with a replayed row not yet backfilled, most recently active + first, with the count and stored-event bytes of those rows.""" + from src.server.database.pool import get_db_connection + + where, params = selection_sql(args) + where.append(f"(r.{LEGACY_VERSION_SQL}) IS DISTINCT FROM %(version)s") + params["version"] = _version() + async with get_db_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + "SELECT t.conversation_thread_id, " + " t.latest_checkpoint_id IS NOT NULL AS has_pointer, " + " count(*) AS rows, " + " count(*) FILTER (WHERE r.has_events) AS with_events, " + " sum(r.events_bytes) AS events_bytes " + "FROM conversation_threads t " + "JOIN workspaces w ON w.workspace_id = t.workspace_id " + f"CROSS JOIN LATERAL ({_REPLAY_ROWS}) r " + f"WHERE {' AND '.join(where)} " + "GROUP BY t.conversation_thread_id, t.latest_checkpoint_id, t.updated_at " + "ORDER BY t.updated_at DESC", + params, + ) + return [ + { + "thread_id": str(row["conversation_thread_id"]), + "has_pointer": row["has_pointer"], + "rows": int(row["rows"]), + "with_events": int(row["with_events"]), + "events_bytes": int(row["events_bytes"] or 0), + } + for row in await cur.fetchall() + ] + + +async def _pending_rows(thread_id: str) -> dict[Any, dict[str, Any]]: + """The thread's replayed rows not yet backfilled, by turn.""" + from src.server.database.pool import get_db_connection + + async with get_db_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + f"SELECT r.conversation_response_id, r.turn_index, r.status, " + f" r.xmin::text AS row_xmin, " + f" r.has_events " + f"FROM (SELECT %s::uuid AS conversation_thread_id) t " + f"CROSS JOIN LATERAL ({_REPLAY_ROWS}) r " + f"WHERE (r.{LEGACY_VERSION_SQL}) IS DISTINCT FROM %s", + (thread_id, _version()), + ) + return { + row["turn_index"]: { + "response_id": str(row["conversation_response_id"]), + "status": row["status"], + "xmin": row["row_xmin"], + "has_events": row["has_events"], + } + for row in await cur.fetchall() + } + + +_WRITE_FACTS = f""" + UPDATE conversation_responses + SET replay_facts = (COALESCE(replay_facts, '{{}}'::jsonb) - '{LEGACY_KEY}') + || jsonb_build_object('{LEGACY_KEY}', %(facts)s::jsonb) + WHERE conversation_response_id = %(id)s::uuid + AND xmin = %(xmin)s::text::xid + AND status <> 'in_progress' + AND ({LEGACY_VERSION_SQL}) IS DISTINCT FROM %(version)s +""" + + +async def _write_facts(rows: list[dict[str, Any]]) -> set[str]: + """Write each row's facts while it is still the version they came from. + The merge keeps every key a live writer set (``reasoning_ms``) and adds + ``v`` only where the row had none. Returns the response ids written.""" + from src.server.database.pool import get_db_connection + from src.server.utils.pg_sanitize import SafeJson + + written: set[str] = set() + if not rows: + return written + async with get_db_connection() as conn: + async with conn.transaction(): + for row in rows: + cur = await conn.execute( + _WRITE_FACTS, + { + "id": row["response_id"], + "xmin": str(row["xmin"]), + "facts": SafeJson(row["facts"]), + "version": _version(), + }, + ) + if cur.rowcount: + written.add(row["response_id"]) + return written + + +@dataclass +class DerivedFacts: + """A turn's legacy facts as its stored events gave them, in the form the + row will hold them, and the row version they were derived from.""" + + response_id: str + xmin: Any + facts: dict[str, Any] + # Every run the turn launched had finished writing. + sealed: bool + # The facts in that form place every row as the stored events did. + verified: bool + + +class Derivation: + """The projection's ``Derive`` hook: keeps each turn's facts as the row + will hold them, and checks that, so held, they place every row where the + stored events did. A jsonb bind drops a NUL, for one.""" + + def __init__(self) -> None: + self.turns: dict[Any, DerivedFacts] = {} + + def __call__( + self, + turn_index: Any, + response: dict[str, Any], + facts: dict[str, Any], + unplaced: tuple[list[dict[str, Any]], list[dict[str, Any]]], + sealed: bool, + ) -> Callable[[list[dict[str, Any]]], None]: + from src.server.services.history.replay import legacy + from src.server.utils.pg_sanitize import jsonb_round_trip + + held = jsonb_round_trip(facts) + # The placement below writes into the items, so the check places + # the held facts into a copy taken before it. + apart = copy.deepcopy(unplaced) + + def placed(items: list[dict[str, Any]]) -> None: + self.turns[turn_index] = DerivedFacts( + response_id=str(response.get("conversation_response_id")), + xmin=response.get("xmin"), + facts=held, + sealed=sealed, + verified=_dumps(legacy.apply_turn(*apart, held)) == _dumps(items), + ) + + return placed + + +def _dumps(items: list[dict[str, Any]]) -> str: + return json.dumps(items, sort_keys=True, default=str) + + +def _sizes(facts: dict[str, Any], stats: Counter) -> None: + """Bytes per facts component, for the report.""" + + def size(value: Any) -> int: + return len(json.dumps(value, ensure_ascii=False, separators=(",", ":"))) + + stats["bytes/total"] += size(facts) + merge = facts.get("merge") or {} + for key in ("widgets", "results", "durations"): + if key in merge: + stats[f"bytes/{key}"] += size(merge[key]) + stats[f"count/{key}"] += len(merge[key]) + for signal in merge.get("signals") or (): + kind = "widget" if "widget" in signal else signal["event"] + stats[f"bytes/signal/{kind}"] += size(signal) + stats[f"count/signal/{kind}"] += 1 + if "verbatim" in facts: + stats["bytes/verbatim"] += size(facts["verbatim"]) + stats["count/verbatim_rows"] += len(facts["verbatim"]) + stats["count/verbatim_turns"] += 1 + + +async def _backfill_thread(thread_id: str, stats: Counter, mismatched: list[str]) -> None: + from src.server.services.history.replay import ( + load_thread_inputs, + project_turns, + ) + + pending = await _pending_rows(thread_id) + if not pending: + return + loaded = await load_thread_inputs(thread_id) + if loaded is None: + stats["threads/no_pointer"] += 1 + return + rows, branch_tip = loaded + + derivation = Derivation() + first = await project_turns( + rows, branch_tip, turn_indexes=sorted(pending), derive=derivation + ) + + writes: list[dict[str, Any]] = [] + for turn_index, row in pending.items(): + found = derivation.turns.get(turn_index) + if found is not None: + if not found.sealed: + stats["turns/runs_live"] += 1 + continue + if not found.verified: + stats["turns/mismatch"] += 1 + mismatched.append(f"{thread_id} turn {turn_index}") + continue + _sizes(found.facts, stats) + writes.append( + { + "turn_index": turn_index, + "response_id": found.response_id, + "xmin": found.xmin, + "facts": found.facts, + } + ) + elif turn_index in first: + # Projected from facts after all: backfilled while this ran. + stats["turns/already"] += 1 + elif row["status"] == "completed": + # No boundary: the thread replays from its stored events whole. + stats["turns/unpaired_completed"] += 1 + elif row["has_events"]: + # No boundary yet: replayed as its stub, but what its events add + # is resolved only once a projection covers it. + stats["turns/stub_with_events"] += 1 + else: + # No boundary and nothing stored: known empty. + stats["turns/stub"] += 1 + writes.append( + { + "turn_index": turn_index, + "response_id": row["response_id"], + "xmin": row["xmin"], + "facts": {"v": int(_version())}, + } + ) + written = await _write_facts(writes) + stats["turns/written"] += len(written) + stats["turns/moved"] += len(writes) - len(written) + + # Bring the written turns' stored lines up to date: the write moved + # each row past the version step 1 keyed them to. + refreshed = [ + w["turn_index"] + for w in writes + if w["response_id"] in written and w["turn_index"] in first + ] + if not refreshed: + return + loaded = await load_thread_inputs(thread_id) + if loaded is None: + return + await project_turns(*loaded, turn_indexes=refreshed) + stats["turns/refreshed"] += len(refreshed) + + +# --------------------------------------------------------------- pauses + + +_PAUSE_ROWS = """ + SELECT r.conversation_response_id + FROM conversation_responses r + JOIN conversation_threads t ON t.conversation_thread_id = r.conversation_thread_id + JOIN workspaces w ON w.workspace_id = t.workspace_id + WHERE r.interrupt_reason = %(reason)s + AND r.status <> 'in_progress' + AND NOT (COALESCE(r.metadata, '{}'::jsonb) ? 'credit_pause_message') +""" + +_STAMP_PAUSE = """ + UPDATE conversation_responses + SET metadata = COALESCE(metadata, '{}'::jsonb) + || jsonb_build_object('credit_pause_message', %(message)s::jsonb) + WHERE conversation_response_id = %(id)s::uuid + AND NOT (COALESCE(metadata, '{}'::jsonb) ? 'credit_pause_message') +""" + + +async def _pause_rows(args: argparse.Namespace) -> list[str]: + """Credit pauses whose denial is still only in their stored events.""" + from src.server.contracts.status import INTERRUPT_REASON_CREDIT_PAUSE + from src.server.database.pool import get_db_connection + + where, params = selection_sql(args) + sql = _PAUSE_ROWS + "".join(f" AND {clause}" for clause in where) + async with get_db_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + await cur.execute(sql, {**params, "reason": INTERRUPT_REASON_CREDIT_PAUSE}) + return [str(row["conversation_response_id"]) for row in await cur.fetchall()] + + +async def _stamp_pauses(response_ids: list[str]) -> tuple[int, int]: + """Stamp each credit pause with the denial its stored interrupt carries, + JSON null when it carries none: either way the key is there, so settling + never reads the stored events for it again. Never over a stamp a + finalize wrote. Returns (stamped, of them without a message).""" + from src.server.database.pool import get_db_connection + from src.server.services.automation_settlement import stored_pause_message as _pause_message + from src.server.utils.pg_sanitize import SafeJson + + stamped = empty = 0 + for response_id in response_ids: + async with get_db_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + "SELECT sse_events FROM conversation_responses " + "WHERE conversation_response_id = %s::uuid", + (response_id,), + ) + row = await cur.fetchone() + message = _pause_message(row["sse_events"] if row else None) + await cur.execute( + _STAMP_PAUSE, {"id": response_id, "message": SafeJson(message)} + ) + stamped += cur.rowcount + empty += cur.rowcount if message is None else 0 + return stamped, empty + + +async def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--apply", action="store_true", help="write; without it, count only") + add_selection_args(parser, concurrency=1, pause=0.5) + args = parser.parse_args() + # Task stream probes decide whether a turn's background runs are sealed; + # without Redis no turn that launched one could be backfilled. + async with app_infra(redis=True, agent_config=False): + return await _run(args) + + +async def _run(args: argparse.Namespace) -> int: + threads = await _pending_threads(args) + pauses = await _pause_rows(args) + rows = sum(t["rows"] for t in threads) + with_events = sum(t["with_events"] for t in threads) + events_mb = sum(t["events_bytes"] for t in threads) / 1e6 + no_pointer = sum(1 for t in threads if not t["has_pointer"]) + print( + f"{len(threads)} threads, {rows} turns pending ({with_events} with stored " + f"events, {events_mb:.1f} MB stored), {no_pointer} threads without a " + f"commit pointer; {len(pauses)} credit pauses unstamped" + ) + if not args.apply: + return 0 + + stamped, empty = await _stamp_pauses(pauses) + print(f"credit pauses stamped: {stamped} ({empty} with no message)") + + stats: Counter = Counter() + mismatched: list[str] = [] + todo = [t["thread_id"] for t in threads if t["has_pointer"]] + began = time.monotonic() + + async def backfill(thread_id: str) -> None: + await _backfill_thread(thread_id, stats, mismatched) + + tally = await for_each_thread( + todo, backfill, concurrency=args.concurrency, pause=args.pause, total=len(todo) + ) + for line in mismatched: + print(f"mismatch: {line}") + for key in sorted(stats): + print(f"{key}: {stats[key]}") + left = await _pending_threads(args) + print( + f"{len(todo)} threads, {tally.unavailable} unavailable, {tally.failed} failed; " + f"{sum(t['rows'] for t in left)} turns still pending " + f"({sum(t['with_events'] for t in left)} with stored events) " + f"in {len(left)} threads, {time.monotonic() - began:.1f}s" + ) + return 1 if tally.failed or mismatched else 0 + + +if __name__ == "__main__": + sys.exit(asyncio.run(main())) diff --git a/scripts/utils/backfill_transcripts.py b/scripts/utils/backfill_transcripts.py index 14cba9ba3..4faa60da0 100644 --- a/scripts/utils/backfill_transcripts.py +++ b/scripts/utils/backfill_transcripts.py @@ -25,35 +25,19 @@ import asyncio import os import sys -import time +from collections.abc import AsyncIterator sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "src"))) -_BATCH = 200 - +from scripts.utils._thread_job import ( # noqa: E402 + add_selection_args, + app_infra, + for_each_thread, + selection_sql, +) -async def _open_infra(): - from src.server.app import setup - from src.server.database import pool as db_pool - from src.server.utils.checkpointer import get_checkpointer, open_checkpointer_pool - from src.utils.cache.redis_cache import init_cache - - pool = db_pool.get_or_create_pool() - await pool.open() - # Raises when Redis is down: an export without the lock could land over a - # server's. - await init_cache() - checkpointer = get_checkpointer( - "postgres", - db_host=os.getenv("DB_HOST", "localhost"), - db_port=int(os.getenv("DB_PORT", "5432")), - db_name=os.getenv("DB_NAME", "postgres"), - db_user=os.getenv("DB_USER", "postgres"), - db_password=os.getenv("DB_PASSWORD", "postgres"), - ) - await open_checkpointer_pool(checkpointer) - setup.checkpointer = checkpointer +_BATCH = 200 async def _threads(args: argparse.Namespace) -> list[tuple[str, str, str]]: @@ -61,24 +45,12 @@ async def _threads(args: argparse.Namespace) -> list[tuple[str, str, str]]: from src.server.database.pool import get_db_connection # Only workspaces on a computer: nothing else ever places a transcript. - where = [ + where, params = selection_sql(args) + where += [ "t.latest_checkpoint_id IS NOT NULL", "w.computer_id IS NOT NULL", "w.status <> 'deleted'", ] - params: list = [] - if args.thread: - where.append("t.conversation_thread_id = ANY(%s::uuid[])") - params.append(args.thread) - if args.workspace: - where.append("t.workspace_id = %s::uuid") - params.append(args.workspace) - if args.user: - where.append("w.user_id = %s") - params.append(args.user) - if args.days: - where.append("t.updated_at >= NOW() - make_interval(days => %s)") - params.append(args.days) async with get_db_connection() as conn: cur = await conn.execute( "SELECT t.conversation_thread_id, t.workspace_id, t.latest_checkpoint_id " @@ -101,15 +73,14 @@ async def _stale(batch: list[tuple[str, str, str]]) -> list[tuple[str, str]]: async def main() -> int: parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) parser.add_argument("--apply", action="store_true", help="store; without it, count only") - parser.add_argument("--days", type=int, help="only threads active in the last N days") - parser.add_argument("--user", help="only this user's threads") - parser.add_argument("--workspace", help="only this workspace's threads") - parser.add_argument("--thread", action="append", help="only these threads") - parser.add_argument("--concurrency", type=int, default=2) - parser.add_argument("--pause", type=float, default=0.2, help="seconds between threads") + add_selection_args(parser, concurrency=2, pause=0.2) args = parser.parse_args() + # Redis is required: an export without the lock could land over a server's. + async with app_infra(redis=True, agent_config=False): + return await _run(args) + - await _open_infra() +async def _run(args: argparse.Namespace) -> int: from src.server.services.transcripts import export_thread threads = await _threads(args) @@ -121,37 +92,27 @@ async def main() -> int: return 0 # A page at a time, each checked against the store as it comes up, so a - # thread a turn end brought current meanwhile is skipped and at most - # ``concurrency`` exports are ever in hand. - done = stored = failed = 0 - began = time.monotonic() - - async def worker(queue: asyncio.Queue[tuple[str, str]]) -> None: - nonlocal done, stored, failed - while True: - try: - thread_id, workspace_id = queue.get_nowait() - except asyncio.QueueEmpty: - return - try: - counts = await export_thread(workspace_id, thread_id) - stored += counts["stored"] - failed += counts["failed"] - except Exception as e: - failed += 1 - print(f"{thread_id} failed: {e}") - done += 1 - if done % 25 == 0: - rate = done / max(time.monotonic() - began, 1e-6) - print(f"{done} done, {stored} stored, {failed} failed, {rate:.1f}/s") - await asyncio.sleep(args.pause) - - for start in range(0, len(threads), _BATCH): - queue: asyncio.Queue[tuple[str, str]] = asyncio.Queue() - for item in await _stale(threads[start : start + _BATCH]): - queue.put_nowait(item) - await asyncio.gather(*(worker(queue) for _ in range(max(1, args.concurrency)))) - print(f"{len(threads)} threads, {done} exported, {stored} stored, {failed} failed") + # thread a turn end brought current meanwhile is skipped. + async def behind() -> AsyncIterator[tuple[str, str]]: + for start in range(0, len(threads), _BATCH): + for item in await _stale(threads[start : start + _BATCH]): + yield item + + stored = failed = 0 + + async def export(item: tuple[str, str]) -> None: + nonlocal stored, failed + thread_id, workspace_id = item + counts = await export_thread(workspace_id, thread_id) + stored += counts["stored"] + failed += counts["failed"] + + tally = await for_each_thread( + behind(), export, concurrency=args.concurrency, pause=args.pause, + name=lambda item: item[0], + ) + failed += tally.failed + tally.unavailable + print(f"{len(threads)} threads, {tally.done} exported, {stored} stored, {failed} failed") return 1 if failed else 0 diff --git a/scripts/utils/backfill_turn_slices.py b/scripts/utils/backfill_turn_slices.py new file mode 100644 index 000000000..84fffb667 --- /dev/null +++ b/scripts/utils/backfill_turn_slices.py @@ -0,0 +1,193 @@ +"""Bring every turn's stored replay up to date: the first-time backfill, and +the re-key step to run after a deploy that changes how turns are projected. + +A thread opens by streaming each turn's stored slice and lines, and projects +from checkpoints only the turns a read cannot serve that way. Two kinds of +turn need it: + +- turns stored before the store existed, which have no row; +- every stored turn after a deploy that changes the projection code or the + config it reads (``lines_epoch()``), or how slices are cut (``SLICE_KEY``): + the rows are still there, under keys no read accepts. + +Until something re-projects them, every open of a long thread pays for it, +``thin_message_snapshots.py`` skips the thread, and after a ``SLICE_KEY`` +change the sliding window cannot trim past turns whose slices are not current. + +This finds candidate threads in SQL (a turn without lines, or a row under an +old key), confirms each with the read path's own test (``unserved_turns``), +and projects only the turns a read would project: one thread at a time with a +pause between, newest first. Safe to stop and rerun: a current turn is read +back, not projected again, so a rerun after a complete pass projects nothing. +Prints ids and counts only. + +A bulk UPDATE of response or query rows (rather than a deploy) leaves every +key's epoch current but changes the row versions the keys cover, which the +SQL cannot see; ``--recheck`` confirms every selected thread instead. + +Run inside the backend container (needs the app's env and venv); a projection +reads checkpoints a batch of turns at a time, so keep the concurrency low: + + /app/.venv/bin/python scripts/utils/backfill_turn_slices.py # count only + /app/.venv/bin/python scripts/utils/backfill_turn_slices.py --apply + /app/.venv/bin/python scripts/utils/backfill_turn_slices.py --apply --days 60 + /app/.venv/bin/python scripts/utils/backfill_turn_slices.py --apply --thread + /app/.venv/bin/python scripts/utils/backfill_turn_slices.py --apply --recheck --days 60 +""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from dataclasses import dataclass + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "src"))) + +from scripts.utils._thread_job import ( # noqa: E402 + add_selection_args, + app_infra, + for_each_thread, + selection_sql, +) + + +@dataclass +class Candidate: + thread_id: str + missing: int # turns without stored lines + stale: int # turns whose lines or slice are under an old key + + +async def _threads(args: argparse.Namespace, *, recheck: bool = False) -> list[Candidate]: + """Threads with a turn a read would project, most recently active first. + + The SQL sees what a deploy changes, not a turn's own inputs: a turn whose + rows changed after its lines were stored re-projects on its next read. + ``recheck`` selects every thread, so ``unserved_turns`` decides alone. + A turn counts by the attempt replay reads (``settled_attempts``); a + superseded attempt is never projected, so never stored. + The counts are an upper bound: a turn whose runs are still live stays + unstored until they settle, and a turn with no place on the branch (its + run died before its first checkpoint) is never projected. + """ + from psycopg.rows import dict_row + + from src.server.database.conversation import settled_attempts + from src.server.database.pool import get_db_connection + from src.server.services.history import slices + from src.server.services.history.replay import lines_epoch + + where, params = selection_sql(args) + params.update(slice_key=slices.SLICE_KEY, epoch=f"{lines_epoch()}.") + stale = ( + "(s.lines IS NULL OR s.lines_key IS NULL " + "OR s.slice_key <> %(slice_key)s " + "OR NOT starts_with(s.lines_key, %(epoch)s))" + ) + where.append("t.latest_checkpoint_id IS NOT NULL") + if not recheck: + where.append(stale) + async with get_db_connection() as conn, conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + "SELECT t.conversation_thread_id AS thread_id, " + " count(*) FILTER (WHERE s.lines IS NULL) AS missing, " + f" count(*) FILTER (WHERE s.lines IS NOT NULL AND {stale}) AS stale " + "FROM conversation_threads t " + "JOIN workspaces w ON w.workspace_id = t.workspace_id " + "CROSS JOIN LATERAL (" + f"{settled_attempts('conversation_response_id', 't.conversation_thread_id')}" + ") r " + "LEFT JOIN turn_slices s " + " ON s.conversation_response_id = r.conversation_response_id " + f"WHERE {' AND '.join(where)} " + "GROUP BY t.conversation_thread_id, t.updated_at " + "ORDER BY t.updated_at DESC", + params, + ) + return [ + Candidate(str(r["thread_id"]), r["missing"], r["stale"]) for r in await cur.fetchall() + ] + + +async def _bring_current(thread_id: str) -> int: + """Project the thread's turns a read would project, and say how many.""" + from src.server.services.history.replay import ( + load_thread_inputs, + project_turns, + unserved_turns, + ) + + loaded = await load_thread_inputs(thread_id) + if loaded is None: + return 0 + rows, branch_tip = loaded + unserved = await unserved_turns(rows, branch_tip) + if not unserved: + return 0 + await project_turns(rows, branch_tip, turn_indexes=unserved) + return len(unserved) + + +def _tally(threads: list[Candidate]) -> str: + missing = sum(c.missing for c in threads) + stale = sum(c.stale for c in threads) + return ( + f"{len(threads)} threads: {missing} turns without stored lines, " + f"{stale} under an old key" + ) + + +async def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--apply", action="store_true", help="store; without it, count only") + parser.add_argument( + "--recheck", + action="store_true", + help="confirm every selected thread, not only those the SQL flags " + "(after a bulk UPDATE of response or query rows)", + ) + add_selection_args(parser, concurrency=2, pause=0.2) + args = parser.parse_args() + # Redis: task stream probes decide whether a turn's background runs are + # sealed, and without them no turn that launched one could be stored. + # The agent config: lines are keyed by the config the server reads (the + # compaction threshold); stored under any other they would be stale on + # arrival. + async with app_infra(redis=True, agent_config=True): + return await _run(args) + + +async def _run(args: argparse.Namespace) -> int: + threads = await _threads(args, recheck=args.recheck) + print(_tally(threads)) + if not args.apply: + return 0 + + projected = turns = current = 0 + + async def bring_current(thread_id: str) -> None: + nonlocal projected, turns, current + n = await _bring_current(thread_id) + if n: + projected += 1 + turns += n + else: + current += 1 + + tally = await for_each_thread( + [c.thread_id for c in threads], bring_current, + concurrency=args.concurrency, pause=args.pause, total=len(threads), + ) + print( + f"{len(threads)} candidate threads: {projected} projected ({turns} turns), " + f"{current} already current, {tally.unavailable} unavailable, {tally.failed} failed" + ) + print(f"left: {_tally(await _threads(args))}") + return 1 if tally.failed else 0 + + +if __name__ == "__main__": + sys.exit(asyncio.run(main())) diff --git a/scripts/utils/replay_parity.py b/scripts/utils/replay_parity.py index 7425fcbc0..3c359ffce 100644 --- a/scripts/utils/replay_parity.py +++ b/scripts/utils/replay_parity.py @@ -1,25 +1,41 @@ """Replay parity gate: checkpoint-sourced vs sse-sourced, without the merge layer. -Cutover step 3 of single-source replay: before sse_events writes stop, every -turn must replay UI-equivalently from checkpoints + tables alone. This script -runs ``build_checkpoint_replay_items`` with the stored-event merge disabled -(``_stored_events`` patched to empty) against ``build_sse_replay_items`` for a -thread corpus and reports per-turn diffs — each diff is a turn still dependent -on the dual-write (legacy payloads, unresolved images, historical event shapes). +Before sse_events writes stop, every turn must replay UI-equivalently from +checkpoints, tables and run facts alone. ``--compare`` picks the two sides: + +- ``sse`` (default): ``build_checkpoint_replay_items`` with the stored-event + merge disabled (``_stored_events`` patched to empty) against + ``build_sse_replay_items``. Each diff is a turn still dependent on the + dual-write (legacy payloads, unresolved images, historical event shapes), + except on a stopped or failed turn whose checkpoints hold the start of what + it streamed: it is expected to differ by the output it never committed, and + is reported apart. +- ``sse-merged``: the same with the merge on, a sanity check. +- ``facts``: each turn replayed twice with the merge on, once with its replay + facts and once with them hidden, reporting where the facts do not + reproduce what the stored events gave. The facts are the run facts and a + backfilled row's legacy facts, so after the backfill this checks that each + row's legacy facts replay as its stored events do. + +Every replay runs with ``cache=False``: turns and runs are cut from their +checkpoints, stored slices and lines are neither read nor written, so a +stale or wrong cache row cannot hide a diff, and the script writes nothing. Run inside the backend container (needs DB env + venv): /app/.venv/bin/python scripts/utils/replay_parity.py --all /app/.venv/bin/python scripts/utils/replay_parity.py --thread [--verbose] - /app/.venv/bin/python scripts/utils/replay_parity.py --all --merge # sanity: with merge on + /app/.venv/bin/python scripts/utils/replay_parity.py --all --compare sse-merged + /app/.venv/bin/python scripts/utils/replay_parity.py --all --compare facts -Exit code 0 = all compared turns equivalent; 1 = diffs found. +Exit code 0 = no unexpected diffs; 1 = diffs found. """ from __future__ import annotations import argparse import asyncio +import dataclasses import json import os import sys @@ -29,6 +45,8 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "src"))) +from scripts.utils._thread_job import app_infra # noqa: E402 + # Event types that never replay (mirrors the ledger's live-only set). _IGNORED = { "metadata", @@ -36,7 +54,6 @@ "warning", "retry", "steering_accepted", - "steering_returned", "model_retry", "model_fallback", # ckpt-projected from the ui channel on new turns, but # legacy checkpoints predate the ui record → corpus-wide compare is noise @@ -49,6 +66,29 @@ # Fields that legitimately differ between the sources. _VOLATILE_KEYS = {"timestamp", "record_id", "artifact_id", "threshold", "response_id"} +# A turn that ended here may hold streamed output its checkpoints never +# committed, which replay no longer shows. +_UNCOMMITTED_STATUSES = {"cancelled", "error", "interrupted"} + + +def _streamed_past_commit(key: str, committed: Any, streamed: Any) -> bool: + """Whether a stopped turn's committed ``key`` is the start of what it + streamed, so the two differ only by output after its last checkpoint. A + change to what the checkpoints did commit, or to any other key, is a diff + like any other.""" + if key == "tool_calls": + return (committed or set()) <= (streamed or set()) + committed, streamed = committed or {}, streamed or {} + if key == "text": + return all(streamed.get(k, "").startswith(v) for k, v in committed.items()) + if key == "signals": + return all(v <= streamed.get(k, 0) for k, v in committed.items()) + if key == "elapsed_ms": + return all(streamed.get(k, [])[: len(v)] == v for k, v in committed.items()) + if key == "results": + return all(streamed.get(k) == v for k, v in committed.items()) + return False + def _lane(data: dict) -> str: agent = data.get("agent") @@ -76,6 +116,9 @@ def _normal_form(items: list[dict]) -> dict[Any, dict]: "user": [], "text": defaultdict(str), "signals": defaultdict(int), + "elapsed_ms": defaultdict(list), + "widgets": [], + "steering_returned": [], "tool_calls": set(), "results": {}, "artifacts": [], @@ -98,6 +141,8 @@ def _normal_form(items: list[dict]) -> dict[Any, dict]: content_type = data.get("content_type") if content_type == "reasoning_signal": turn["signals"][_lane(data)] += 1 + if data.get("elapsed_ms") is not None: + turn["elapsed_ms"][_lane(data)].append(data.get("elapsed_ms")) elif content_type in ("text", "reasoning"): turn["text"][(_lane(data), content_type)] += data.get("content") or "" elif event == "tool_calls": @@ -110,6 +155,10 @@ def _normal_form(items: list[dict]) -> dict[Any, dict]: ) elif event == "artifact": payload = data.get("payload") or {} + if data.get("artifact_type") == "html_widget": + turn["widgets"].append( + (payload.get("title"), _canon(payload.get("data"))) + ) turn["artifacts"].append( ( data.get("artifact_type"), @@ -144,6 +193,15 @@ def _normal_form(items: list[dict]) -> dict[Any, dict]: data.get("tool_call_id"), ) ) + elif event == "steering_returned": + turn[event].append( + ( + _lane(data), + data.get("input_id"), + data.get("reason"), + data.get("content"), + ) + ) elif event in ("steering_delivered", "interrupt", "error", "credit_usage"): turn[event].append(_strip(data)) return dict(turns) @@ -153,11 +211,11 @@ def _diff_turn(a: dict, b: dict) -> list[str]: reasons = [] for key in a.keys() | b.keys(): va, vb = a.get(key), b.get(key) - if key in ("text", "signals"): + if key in ("text", "signals", "elapsed_ms"): va, vb = dict(va or {}), dict(vb or {}) if key == "artifacts": va, vb = sorted(map(_canon, va or [])), sorted(map(_canon, vb or [])) - if key == "context": + if key in ("context", "widgets", "steering_returned"): va, vb = sorted(map(_canon, va or [])), sorted(map(_canon, vb or [])) if key in ("steering_delivered", "interrupt", "error", "credit_usage"): va = sorted(map(_canon, va or [])) @@ -167,28 +225,6 @@ def _diff_turn(a: dict, b: dict) -> list[str]: return reasons -async def _open_infra(): - from src.server.app import setup - from src.server.database import pool as db_pool - from src.server.utils.checkpointer import ( - get_checkpointer, - open_checkpointer_pool, - ) - - pool = db_pool.get_or_create_pool() - await pool.open() - checkpointer = get_checkpointer( - "postgres", - db_host=os.getenv("DB_HOST", "localhost"), - db_port=int(os.getenv("DB_PORT", "5432")), - db_name=os.getenv("DB_NAME", "postgres"), - db_user=os.getenv("DB_USER", "postgres"), - db_password=os.getenv("DB_PASSWORD", "postgres"), - ) - await open_checkpointer_pool(checkpointer) - setup.checkpointer = checkpointer - - async def _thread_ids(only: list[str]) -> list[str]: if only: return only @@ -203,78 +239,130 @@ async def _thread_ids(only: list[str]) -> list[str]: return [str(r[0]) for r in await cur.fetchall()] -async def _compare_thread(thread_id: str, merge: bool, verbose: bool) -> tuple[int, int]: - """Returns (turns_compared, turns_diff).""" - from src.server.database.conversation.replay_rows import get_replay_thread_data +async def _checkpoint_items(rows, tip, *, merge=True, facts=True): + """One uncached checkpoint replay of the whole thread, with the + stored-event merge and the run facts each on or off. A row's legacy facts + are the merge's output, so they go with either.""" from src.server.services.history import replay + from src.server.services.history.replay import facts as replay_facts + from src.server.services.history.replay import legacy - _, thread, queries, responses, usages, provenance = await get_replay_thread_data( - thread_id - ) - if not thread or thread.get("latest_checkpoint_id") is None: - print(f"{thread_id} SKIP (no commit pointer)") - return 0, 0 - responses_by_turn = { - r.get("turn_index"): r for r in responses if isinstance(r, dict) - } - - sse_items = replay.build_sse_replay_items(thread_id, queries, responses_by_turn) - - original_stored_events = replay.stored_merge._stored_events + saved_events = replay.stored_merge._stored_events + saved_durations = replay_facts.reasoning_durations + saved_legacy = legacy.facts_of if not merge: replay.stored_merge._stored_events = lambda response: [] + if not merge or not facts: + legacy.facts_of = lambda response: None + if not facts: + replay_facts.reasoning_durations = lambda response: None + rows = dataclasses.replace(rows, run_facts=[]) try: - checkpoint_items = await replay.build_checkpoint_replay_items( - thread_id, - queries, - responses_by_turn, - branch_tip_checkpoint_id=thread.get("latest_checkpoint_id"), - usages=usages, - provenance=provenance, - ) - except replay.CheckpointReplayUnavailable as e: - print(f"{thread_id} FALLBACK ({e})") - return 0, 0 + return await replay.build_checkpoint_replay_items(rows, tip, cache=False) finally: - replay.stored_merge._stored_events = original_stored_events + replay.stored_merge._stored_events = saved_events + replay_facts.reasoning_durations = saved_durations + legacy.facts_of = saved_legacy + + +async def _compare_thread( + thread_id: str, compare: str, verbose: bool +) -> tuple[int, int, int]: + """Returns (turns_compared, turns_diff, turns_expected_diff).""" + from src.server.services.history import replay - ckpt, sse = _normal_form(checkpoint_items), _normal_form(sse_items) - diffs = 0 - for turn_index in sorted(ckpt.keys() | sse.keys(), key=lambda x: (x is None, x)): - reasons = _diff_turn(ckpt.get(turn_index, {}), sse.get(turn_index, {})) - if reasons: + loaded = await replay.load_thread_inputs(thread_id) + if loaded is None: + print(f"{thread_id} SKIP (no commit pointer)") + return 0, 0, 0 + rows, tip = loaded + # Rows arrive without their stored events; both sides need them here, + # including the rows whose legacy facts stand in for them. + responses_by_turn = await replay.with_stored_events( + rows.responses_by_turn, list(rows.responses_by_turn), verbatim=True + ) + rows = dataclasses.replace(rows, responses_by_turn=responses_by_turn) + + try: + if compare == "facts": + left = await _checkpoint_items(rows, tip) + right = await _checkpoint_items(rows, tip, facts=False) + labels = ("facts", "stored") + else: + left = await _checkpoint_items(rows, tip, merge=compare == "sse-merged") + right = replay.build_sse_replay_items( + thread_id, rows.queries, responses_by_turn + ) + labels = ("ckpt", "sse") + except replay.CheckpointReplayUnavailable as e: + print(f"{thread_id} FALLBACK ({e})") + return 0, 0, 0 + + a, b = _normal_form(left), _normal_form(right) + diffs = expected = 0 + for turn_index in sorted(a.keys() | b.keys(), key=lambda x: (x is None, x)): + ckpt, sse = a.get(turn_index, {}), b.get(turn_index, {}) + reasons = _diff_turn(ckpt, sse) + if not reasons: + continue + status = (responses_by_turn.get(turn_index) or {}).get("status") + if ( + compare != "facts" + and status in _UNCOMMITTED_STATUSES + and all( + _streamed_past_commit(key, ckpt.get(key), sse.get(key)) + for key in reasons + ) + ): + expected += 1 + verdict = f"EXPECTED ({status}, uncommitted output)" + else: diffs += 1 - print(f"{thread_id} turn {turn_index} DIFF: {', '.join(sorted(reasons))}") - if verbose: - for key in reasons: - print(f" ckpt {key}: {_canon(ckpt.get(turn_index, {}).get(key))[:400]}") - print(f" sse {key}: {_canon(sse.get(turn_index, {}).get(key))[:400]}") - total = len(ckpt.keys() | sse.keys()) - if not diffs: + verdict = "DIFF" + print(f"{thread_id} turn {turn_index} {verdict}: {', '.join(sorted(reasons))}") + if verbose: + for key in reasons: + print(f" {labels[0]:<6} {key}: {_canon(a.get(turn_index, {}).get(key))[:400]}") + print(f" {labels[1]:<6} {key}: {_canon(b.get(turn_index, {}).get(key))[:400]}") + total = len(a.keys() | b.keys()) + if not diffs and not expected: print(f"{thread_id} OK ({total} turns)") - return total, diffs + return total, diffs, expected async def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--thread", action="append", default=[], help="thread id (repeatable)") parser.add_argument("--all", action="store_true", help="all threads with a commit pointer") - parser.add_argument("--merge", action="store_true", help="keep the stored-event merge on") + parser.add_argument( + "--compare", + choices=("sse", "sse-merged", "facts"), + default="sse", + help="sse: unmerged replay vs stored events; sse-merged: merged replay " + "vs stored events; facts: replay with replay facts vs with them hidden", + ) parser.add_argument("--verbose", action="store_true") args = parser.parse_args() if not args.thread and not args.all: parser.error("pass --thread or --all") + async with app_infra(redis=False, agent_config=False): + return await _run(args) - await _open_infra() - total_turns = total_diffs = threads = 0 + +async def _run(args: argparse.Namespace) -> int: + total_turns = total_diffs = total_expected = threads = 0 for thread_id in await _thread_ids(args.thread): - compared, diffs = await _compare_thread(thread_id, args.merge, args.verbose) + compared, diffs, expected = await _compare_thread( + thread_id, args.compare, args.verbose + ) threads += 1 total_turns += compared total_diffs += diffs + total_expected += expected print( f"\n{threads} threads, {total_turns} turns compared, " - f"{total_diffs} turn diffs ({'merge ON' if args.merge else 'merge OFF'})" + f"{total_diffs} turn diffs, {total_expected} expected on stopped or " + f"failed turns (--compare {args.compare})" ) return 1 if total_diffs else 0 diff --git a/src/config/settings.py b/src/config/settings.py index 4695ad23d..1bbf713f4 100644 --- a/src/config/settings.py +++ b/src/config/settings.py @@ -338,17 +338,6 @@ def get_recovery_scan_interval() -> float: return value if math.isfinite(value) and value > 0 else 30.0 -def get_replay_projection_cache_ttl() -> int: - """Env REPLAY_PROJECTION_CACHE_TTL_SECONDS, default 14 days. 0 disables.""" - default = 14 * 86400 - try: - return max( - 0, int(os.getenv("REPLAY_PROJECTION_CACHE_TTL_SECONDS", str(default))) - ) - except ValueError: - return default - - def get_redis_socket_connect_timeout() -> int: return get_infrastructure_config().redis.socket_connect_timeout diff --git a/src/llms/error_classification.py b/src/llms/error_classification.py index 0e8b7cb46..94ef7c39f 100644 --- a/src/llms/error_classification.py +++ b/src/llms/error_classification.py @@ -3,7 +3,7 @@ Used by the agent's model-resilience middleware to decide whether a failed model call is worth retrying on the same model or should move straight to the next fallback. Mirrors the exception-chain walking in -``src/server/services/runs/sse_producer.py::classify_stream_exception`` — kept +``src/server/services/runs/stream_errors.py::classify_stream_exception`` — kept separate because ptc_agent code must not import from src/server. """ diff --git a/src/ptc_agent/agent/agent.py b/src/ptc_agent/agent/agent.py index 44e25ffbd..bbccc2e1f 100644 --- a/src/ptc_agent/agent/agent.py +++ b/src/ptc_agent/agent/agent.py @@ -19,7 +19,7 @@ from ptc_agent.core.paths import WorkspaceLayout from ptc_agent.core.project_context import ProjectContext from ptc_agent.agent.middleware import SubAgentMiddleware -from ptc_agent.agent.state import DeltaAgentState +from ptc_agent.agent.main_state import MainAgentState from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware from langchain_anthropic.middleware import AnthropicPromptCachingMiddleware @@ -53,7 +53,9 @@ NotesDueMiddleware, NotesOffMiddleware, ThreadScratchpad, + notes_window_carry, ) +from ptc_agent.agent.middleware.compaction.window import WindowCoverage from ptc_agent.agent.middleware.direct_mcp import ( DirectToolSet, direct_tool_middleware, @@ -63,6 +65,7 @@ from ptc_agent.agent.middleware.subagent_switch import ( SubagentSwitchMiddleware, SubagentSwitchReader, + subagents_window_carry, ) from ptc_agent.agent.context_stack import build_context_middleware from ptc_agent.agent.roles import AgentRole @@ -245,6 +248,7 @@ def create_agent( role: AgentRole = "analyst", harness_blocks: Mapping[str, str | None] | None = None, subagent_switch: SubagentSwitchReader | None = None, + window_coverage: WindowCoverage | None = None, ) -> Any: """Create a deepagent with PTC pattern capabilities. @@ -281,6 +285,10 @@ def create_agent( row to read (thread maintenance, a synthetic turn), which leaves subagents as built. It never changes the tools or the prompt: it adds a history row and refuses launches instead. + window_coverage: Asks the server whether turn slices hold the + runs the main agent's window would trim from its checkpoint + (see ``compaction.window``). None for a build with no server + behind it, which never trims. Returns: Configured BackgroundSubagentOrchestrator wrapping the deepagent. @@ -697,6 +705,13 @@ def create_agent( elif compaction is not None: main_compaction = compaction.with_scratchpad_notes(scratchpad.notes_dir) notes_rows = NotesDueMiddleware(scratchpad, main_compaction) + # Each carry writes a field ``MainAgentState`` declares, so a build + # without the carry's reader (a notification turn has no switch) + # trims too. + if main_compaction and window_coverage and thread_id: + main_compaction = main_compaction.with_window( + window_coverage, [notes_window_carry(thread_id), subagents_window_carry] + ) model_resilience = [build_model_resilience_middleware(self.config, turn)] @@ -947,7 +962,7 @@ def create_agent( middleware=deepagent_middleware, checkpointer=checkpointer, store=store, - state_schema=DeltaAgentState, + state_schema=MainAgentState, ).with_config({"recursion_limit": 2000}) return BackgroundSubagentOrchestrator( diff --git a/src/ptc_agent/agent/graph.py b/src/ptc_agent/agent/graph.py index 6e8342813..10e7de2d7 100644 --- a/src/ptc_agent/agent/graph.py +++ b/src/ptc_agent/agent/graph.py @@ -6,6 +6,7 @@ from ptc_agent.agent.agent import AgentRole, PTCAgent from ptc_agent.agent.middleware.runtime_context import TurnContext +from ptc_agent.agent.middleware.compaction.window import WindowCoverage from ptc_agent.agent.middleware.subagent_switch import SubagentSwitchReader from ptc_agent.config import AgentConfig from ptc_agent.core.project_context import ProjectContext @@ -199,6 +200,7 @@ async def build_ptc_graph_with_session( tool_view: Any | None = None, role: AgentRole = "analyst", subagent_switch: SubagentSwitchReader | None = None, + window_coverage: WindowCoverage | None = None, ) -> Any: """Build a BackgroundSubagentOrchestrator from a pre-acquired session (WorkspaceManager path). @@ -215,7 +217,9 @@ async def build_ptc_graph_with_session( own fields belong to whichever project on the machine resolved last. ``subagent_switch`` reads the thread's subagent switch on every call; a - build without one leaves subagents as built. + build without one leaves subagents as built. ``window_coverage`` lets the + main agent trim its checkpoint to its window (see ``compaction.window``); + a build without one never trims. """ mcp_registry = ( tool_view.mcp_registry if tool_view is not None else session.mcp_registry @@ -289,6 +293,7 @@ async def build_ptc_graph_with_session( role=role, harness_blocks=harness_blocks, subagent_switch=subagent_switch, + window_coverage=window_coverage, ) logger.debug( diff --git a/src/ptc_agent/agent/main_state.py b/src/ptc_agent/agent/main_state.py new file mode 100644 index 000000000..2c033520d --- /dev/null +++ b/src/ptc_agent/agent/main_state.py @@ -0,0 +1,30 @@ +"""The main agent's state schema: every field its middleware keep, declared +whichever of them a build includes. + +A graph writes back only the channels it declares, and the Postgres saver +keeps a primitive value (a count, a flag) inside the checkpoint row rather +than beside it, so one write through a graph that lacks such a field erases +it. The main agent's checkpoints are written by builds that differ (a turn's, +and thread maintenance's, which leaves the subagent switch out) and by the +server's history reader appending a ui record, so each declares this schema. +""" + +from __future__ import annotations + +from ptc_agent.agent.middleware.compaction.notes import NotesDueState +from ptc_agent.agent.middleware.compaction.types import CompactionState +from ptc_agent.agent.middleware.skills.middleware import LoadedSkillsState +from ptc_agent.agent.middleware.subagent_switch import SubagentSwitchState +from ptc_agent.agent.state import DeltaAgentState + + +# DeltaAgentState last: a TypedDict takes each field from the last base that +# declares it, and only DeltaAgentState gives ``messages`` its DeltaChannel. +class MainAgentState( + CompactionState, + NotesDueState, + SubagentSwitchState, + LoadedSkillsState, + DeltaAgentState, +): + """The main agent's checkpointed state; see the module docstring.""" diff --git a/src/ptc_agent/agent/middleware/background_subagent/redis_stream.py b/src/ptc_agent/agent/middleware/background_subagent/redis_stream.py index 8b050c014..481b91b1c 100644 --- a/src/ptc_agent/agent/middleware/background_subagent/redis_stream.py +++ b/src/ptc_agent/agent/middleware/background_subagent/redis_stream.py @@ -675,6 +675,39 @@ def parse_steering_payload(raw: Any) -> dict[str, Any] | None: return None +def returned_entry(payload: dict[str, Any], reason: str) -> dict[str, Any]: + """A returned steering input, as its ``steering_returned`` event and the + run's replay fact both name it.""" + return { + "content": payload["content"], + "input_id": payload["input_id"], + "reason": reason, + } + + +async def record_returned( + registry: Any, run_id: str | None, returned: list[dict[str, Any]] +) -> bool: + """Keep returned steering on the run's ledger row for replay, which reads + the run's transcript from its checkpoint, where nothing returned ever + lands. Never raises; False only when a ledger that should hold the + returns did not take them, so a caller can keep its queue entries.""" + ledger = getattr(registry, "run_ledger", None) if registry else None + if not returned or ledger is None or not run_id: + return True + try: + return bool( + await ledger.record_replay_facts(run_id, steering_returned=returned) + ) + except Exception: + logger.warning( + "returned steering not recorded for replay", + task_run_id=run_id, + exc_info=True, + ) + return False + + async def read_task_meta(thread_id: str, task_id: str) -> dict[str, str] | None: """Read the cross-worker task meta hash written by ``write_task_meta``. diff --git a/src/ptc_agent/agent/middleware/background_subagent/run_executor.py b/src/ptc_agent/agent/middleware/background_subagent/run_executor.py index c552fb065..2d01663cb 100644 --- a/src/ptc_agent/agent/middleware/background_subagent/run_executor.py +++ b/src/ptc_agent/agent/middleware/background_subagent/run_executor.py @@ -28,6 +28,8 @@ from ptc_agent.agent.middleware.background_subagent.outcomes import Outcome from ptc_agent.agent.middleware.background_subagent.redis_stream import ( parse_steering_payload, + record_returned, + returned_entry, steering_queue_key, ) from ptc_agent.agent.middleware.background_subagent.registry import ( @@ -105,14 +107,14 @@ async def _return_unconsumed_steering( ``steering_returned`` events — accepted input a run never consumed is surfaced, not left for a later resume or a silent TTL death. - Read → surface → remove, in that order: an entry leaves the queue only - after it made it into the event archive, so a spill failure (or a crash) - between the two leaves the entries in Redis until TTL instead of - silently destroying acknowledged input. Only the entries read are - removed, one by one, never the key: a producer that pushes behind the - read sees the terminal meta (written before this sweep) on its - post-push verify and reclaims its own entry, and it reads an entry gone - from the queue as one the run settled.""" + Read → surface → record → remove, in that order: an entry leaves the + queue only after it made it into the event archive and onto the run's + ledger row, so a failure of either (or a crash) before the removal + leaves the entries in Redis until TTL instead of silently destroying + acknowledged input. Only the entries read are removed, one by one, never + the key: a producer that pushes behind the read sees the terminal meta + (written before this sweep) on its post-push verify and reclaims its own + entry, and it reads an entry gone from the queue as one the run settled.""" try: from src.utils.cache.redis_cache import get_cache_client @@ -130,11 +132,14 @@ async def _return_unconsumed_steering( ts = time.time() seq_before = task.captured_event_seq appended = 0 + returned: list[dict] = [] for raw in raw_messages: payload = parse_steering_payload(raw) if payload is None: continue appended += 1 + entry = returned_entry(payload, "run_ended") + returned.append(entry) # Identity-exact append: this settle may run after cancel # teardown evicted the entry (or its tool_call_id was reused # by a re-registration) — resolving by id would drop the frame @@ -143,12 +148,7 @@ async def _return_unconsumed_steering( task, { "event": "steering_returned", - "data": { - "agent": f"task:{task.task_id}", - "content": payload["content"], - "input_id": payload["input_id"], - "reason": "run_ended", - }, + "data": {"agent": f"task:{task.task_id}", **entry}, "ts": ts, }, # A kill seals the task's streams; the returned-input record @@ -172,6 +172,14 @@ async def _return_unconsumed_steering( pending=appended, ) return + if not await record_returned(registry, task.task_run_id, returned): + logger.warning( + "unconsumed-steering sweep withheld; returns not recorded", + task_id=task.task_id, + task_run_id=task.task_run_id, + pending=appended, + ) + return for raw in raw_messages: await cache.client.lrem(key, 1, raw) logger.info( diff --git a/src/ptc_agent/agent/middleware/background_subagent/steering.py b/src/ptc_agent/agent/middleware/background_subagent/steering.py index d3c05a408..15915d0e4 100644 --- a/src/ptc_agent/agent/middleware/background_subagent/steering.py +++ b/src/ptc_agent/agent/middleware/background_subagent/steering.py @@ -27,6 +27,8 @@ ) from ptc_agent.agent.middleware.background_subagent.redis_stream import ( parse_steering_payload, + record_returned, + returned_entry, steering_queue_key, ) @@ -110,20 +112,23 @@ async def abefore_model( agent_id = f"task:{task_id}" if task_id else f"subagent:{tool_call_id}" ts = time.time() if returned and self.registry: - for payload in returned: + entries = [returned_entry(p, "run_mismatch") for p in returned] + for entry in entries: await self._capture( tool_call_id, { "event": "steering_returned", - "data": { - "agent": agent_id, - "content": payload["content"], - "input_id": payload["input_id"], - "reason": "run_mismatch", - }, + "data": {"agent": agent_id, **entry}, "ts": ts, }, ) + # ``after`` places each return after the message it followed. + after = _last_message_id(state) + await record_returned( + self.registry, + own_run_id, + [{**entry, "after": after} for entry in entries], + ) if not parsed: return None @@ -182,3 +187,10 @@ async def _capture(self, tool_call_id: str, event: dict[str, Any]) -> None: await self.registry.append_captured_event(tool_call_id, event) except Exception: pass + + +def _last_message_id(state: Any) -> str | None: + messages = state.get("messages") if isinstance(state, dict) else None + if not messages: + return None + return getattr(messages[-1], "id", None) diff --git a/src/ptc_agent/agent/middleware/compaction/compact.py b/src/ptc_agent/agent/middleware/compaction/compact.py index fa0513f92..c171230fb 100644 --- a/src/ptc_agent/agent/middleware/compaction/compact.py +++ b/src/ptc_agent/agent/middleware/compaction/compact.py @@ -59,7 +59,7 @@ select_offloads, tool_call_ids, ) -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.pointer import ( TranscriptTurns, aexport_transcript, @@ -147,10 +147,11 @@ async def compact( fallback: Any | None, thread_id: str | None = None, notes_dir: str | None = None, + window: Window, ) -> Compaction: """Summarize ``view`` before ``cutoff``, ``messages`` being the - agent's whole checkpoint list, from this model, ``fallback`` or the - server (see ``summarize``). + agent's checkpoint list (what ``window`` trimmed of it aside), from + this model, ``fallback`` or the server (see ``summarize``). The transcript is saved first: the summary cites its turn files and ends pointing at it only if that save lands and ``workspace_id``'s @@ -163,9 +164,14 @@ async def compact( started = time.monotonic() budget = get_compaction_timeout() exported = await aexport_transcript( - backend, transcript, messages, workspace_id=workspace_id, budget=budget + backend, + transcript, + messages, + workspace_id=workspace_id, + budget=budget, + window=window, ) - turns = TranscriptTurns.of(exported, messages) if exported else None + turns = TranscriptTurns.of(exported, messages, window=window) if exported else None async def render(trimmed: list[AnyMessage]) -> list[AnyMessage]: trimmed = await aoffload_base64_content(backend, trimmed, thread_id=thread_id) @@ -184,7 +190,11 @@ async def render(trimmed: list[AnyMessage]) -> list[AnyMessage]: to_summarize, limit=self.limit, counter=self.counter, render=render ), server=lambda: server_summary( - to_summarize, preserved, raw_messages=messages, turns=turns + to_summarize, + preserved, + raw_messages=messages, + turns=turns, + window=window, ), budget=remaining, ) @@ -261,8 +271,8 @@ async def compact_messages( notes_dir: str | None = None, ) -> Compaction: """Manual /compact: summarize all but the last ``keep_messages`` of the - view, ``messages`` being the thread's whole checkpoint list and ``state`` - its values. The user's main model is tried when the summary model fails. + view, ``messages`` being the thread's checkpoint list and ``state`` its + values. The user's main model is tried when the summary model fails. ``notes_dir`` is the thread's scratchpad notes folder when the scratchpad feature is on. @@ -301,6 +311,7 @@ async def compact_messages( fallback=_main_client(config), thread_id=thread_id, notes_dir=notes_dir, + window=Window.of(state), ) @@ -348,7 +359,11 @@ async def offload_tool_args( if transcript is None: arg_ids = set() elif await aexport_transcript( - backend, transcript, messages, workspace_id=workspace_id + backend, + transcript, + messages, + workspace_id=workspace_id, + window=Window.of(state), ) is None: if not read_ids: # A failure to retry, not "nothing to offload": the caller diff --git a/src/ptc_agent/agent/middleware/compaction/middleware.py b/src/ptc_agent/agent/middleware/compaction/middleware.py index a0f521253..6ac13f691 100644 --- a/src/ptc_agent/agent/middleware/compaction/middleware.py +++ b/src/ptc_agent/agent/middleware/compaction/middleware.py @@ -10,6 +10,9 @@ transcript file that keeps them and stale Read results are hidden. - Tier 2: when the context reaches its threshold, the start of the view is summarized (see ``compact``). +- The window: on the main agent, at each pass through ``before_agent``, the + runs a summary stands in for are trimmed from the checkpoint once the + server holds them (see ``window``). Each step is reported as a ``context_window`` event, whose ``action`` values are wire protocol: ``token_usage`` after each model call, ``summarize`` @@ -19,10 +22,11 @@ from __future__ import annotations +import asyncio import copy import logging import time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any from langchain_core.exceptions import ContextOverflowError @@ -55,7 +59,13 @@ recorded_offloads, select_offloads, ) -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.middleware.compaction.window import ( + WindowCarry, + WindowCoverage, + trim, + window_cut, +) +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.pointer import aexport_transcript, transcript_target if TYPE_CHECKING: @@ -64,6 +74,11 @@ logger = logging.getLogger(__name__) +# The coverage read sits in front of a turn. It walks the thread's turn +# boundaries, which takes tens of milliseconds; past this the trim waits for +# the next pass rather than hold the turn. +_COVERAGE_TIMEOUT_S = 5.0 + class CompactionMiddleware(AgentMiddleware): """Keeps an agent's context under its threshold (see the module docstring). @@ -106,6 +121,9 @@ def __init__( self._workspace_id = workspace_id # The notes folder this agent's summaries name; see with_scratchpad_notes. self._notes_dir: str | None = None + # The main agent's window; see with_window. + self._coverage: WindowCoverage | None = None + self._carries: tuple[WindowCarry, ...] = () @classmethod def for_agent( @@ -207,6 +225,7 @@ async def _compact( transcript=transcript, fallback=request.model, notes_dir=self._notes_dir, + window=Window.of(request.state), ) except BaseException as e: self._emit_context_signal("summarize", "error", error=str(e)) @@ -269,20 +288,31 @@ def _extract_token_usage(self, response: ModelResponse) -> tuple[int, int]: @override async def abefore_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None: + """Tier 1 when the turn starts after a long pause, then the window's + trim. A resume after an interrupt does not pass through here, since + the graph picks up where it stopped; an orchestrator re-entry does.""" + update = await self._offload_idle(state) + trimmed = await self._trim({**state, **update}) + return {**update, **(trimmed or {})} or None + + async def _offload_idle(self, state: Mapping[str, Any]) -> dict[str, Any]: """Tier 1, once per turn and only when the model last answered more than the idle threshold ago: by then the provider's prompt cache has expired, so hiding part of the prefix costs no cache hit, while doing - it mid-turn would throw a warm one away. A resume after an interrupt - does not pass through here, since the graph picks up where it stopped. - An argument is hidden only once the transcript holding it is saved - where the workspace can read it, since that is then its one copy.""" + it mid-turn would throw a warm one away. An argument is hidden only + once the transcript holding it is saved where the workspace can read + it, since that is then its one copy.""" if not is_idle(state, self._offload, time.time()): - return None + return {} messages = state["messages"] effective = get_effective_messages(messages, state.get("_summarization_event")) args, reads = select_offloads(effective, self._offload, recorded_offloads(state)) if args and await aexport_transcript( - self._backend, self._transcript_target(), messages, workspace_id=self._workspace_id + self._backend, + self._transcript_target(), + messages, + workspace_id=self._workspace_id, + window=Window.of(state), ) is None: args = set() if args: @@ -291,7 +321,64 @@ async def abefore_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None self._emit_context_signal( "offload", "complete", kind="reads", offloaded_reads=len(reads) ) - return record_offloads(state, args, reads) or None + return record_offloads(state, args, reads) + + async def _trim(self, state: Mapping[str, Any]) -> dict[str, Any] | None: + """The window's trim (see ``window``), or None while it may not: + without the server to ask (any stack but the main agent's, see + ``with_window``), before a summary with an anchor, or while the + server's slices do not hold exactly the messages it would drop, in + which case the update keeps the refusal. + + A refusal is kept in the checkpoint with the anchor and cut it was + for, and not asked again until a summary moves them: a thread whose + slices do not match pays for the check once per summary, not per + turn, and one backfilled since trims after its next summary. + """ + if self._coverage is None: + return None + messages = state["messages"] + event = state.get("_summarization_event") + cut = window_cut(messages, event) + if cut <= 0: + return None + asked = {"anchor": event["anchor_message_id"], "cut": cut} + if state.get("_window_refused") == asked: + return None + # Hashing the head is CPU over every message it drops. + head = await asyncio.to_thread(Window.of(state).extend, messages[:cut]) + try: + async with asyncio.timeout(_COVERAGE_TIMEOUT_S): + covered = await self._coverage(head) + except Exception: # noqa: BLE001 - an unanswered check only delays the trim + logger.warning("[Compaction] window coverage unread", exc_info=True) + return None + if not covered: + logger.info( + "[Compaction] window waits for the next summary: slices do not " + "hold runs 1-%d", + head.runs, + ) + return {"_window_refused": asked} + logger.info( + "[Compaction] window trimmed %d messages, runs up to %d", cut, head.runs + ) + return trim(state, cut, self._carries, head=head) + + def with_window( + self, coverage: WindowCoverage, carries: Sequence[WindowCarry] = () + ) -> CompactionMiddleware: + """This middleware, trimming the main agent's checkpoint to its window. + + ``coverage`` asks the server whether turn slices hold the runs a trim + would drop; ``carries`` keep what other readers counted in them. For + the main agent's stack only: a subagent's run is bounded, and its + transcript reads its namespace's whole history. + """ + main = copy.copy(self) + main._coverage = coverage + main._carries = tuple(carries) + return main def with_scratchpad_notes(self, notes_dir: str) -> CompactionMiddleware: """This middleware, with each summary naming ``notes_dir``'s files. diff --git a/src/ptc_agent/agent/middleware/compaction/notes.py b/src/ptc_agent/agent/middleware/compaction/notes.py index a773a29b2..76f038703 100644 --- a/src/ptc_agent/agent/middleware/compaction/notes.py +++ b/src/ptc_agent/agent/middleware/compaction/notes.py @@ -131,17 +131,41 @@ def resolve( user's flag is off or there is no thread to name the folder after.""" if not thread_id or not config.feature_enabled("scratchpad"): return None - short_id = thread_id[:8] - notes_subdir = WorkspaceLayout.scratchpad_subdir( - short_id, WorkspaceLayout.SCRATCHPAD_NOTE_DIR - ) + notes_subdir = thread_notes_subdir(thread_id) return cls( - folder=posixpath.join(workspace, WorkspaceLayout.scratchpad_subdir(short_id)) + "/", + folder=posixpath.join( + workspace, WorkspaceLayout.scratchpad_subdir(thread_id[:8]) + ) + + "/", notes_dir=posixpath.join(workspace, notes_subdir), notes_subdir=notes_subdir, ) +def thread_notes_subdir(thread_id: str) -> str: + """The thread's notes folder, workspace-relative, whether or not the + scratchpad is on.""" + return WorkspaceLayout.scratchpad_subdir(thread_id[:8], WorkspaceLayout.SCRATCHPAD_NOTE_DIR) + + +def _notes_pattern(notes_subdir: str) -> re.Pattern[str]: + return re.compile(re.escape(notes_subdir) + r"(?![\w.-])") + + +def notes_window_carry(thread_id: str) -> Callable[[Sequence[Any], Mapping[str, Any]], dict[str, Any]]: + """The window's carry (see ``compaction.window``) for the check-in's + count, which runs back to the last write to the notes however far back + that was. Kept with the scratchpad off too: it can be turned on later.""" + folder = _notes_pattern(thread_notes_subdir(thread_id)) + + def carry(trimmed: Sequence[Any], state: Mapping[str, Any]) -> dict[str, Any]: + held = state.get("_notes_calls_trimmed") or 0 + total = _calls_since_write(trimmed, folder, held) + return {"_notes_calls_trimmed": total} if total != held else {} + + return carry + + async def anotes_pointer( backend: Any | None, notes_dir: str | None, budget: float = _LIST_TIMEOUT ) -> str: @@ -199,9 +223,12 @@ def _quotable(name: str) -> bool: class NotesDueState(AgentState): """``_notes_below_mark`` names the summary under which a call last measured the context below the reminder's mark, None before the first - summary: the reminder is asked only after one has.""" + summary: the reminder is asked only after one has. + ``_notes_calls_trimmed`` is the tool calls since the last write to the + notes in the runs the window trimmed, for the check-in's count.""" _notes_below_mark: Annotated[NotRequired[str | None], PrivateStateAttr] + _notes_calls_trimmed: Annotated[NotRequired[int], PrivateStateAttr] class NotesDueMiddleware(AgentMiddleware): @@ -233,7 +260,7 @@ def __init__(self, scratchpad: ThreadScratchpad, compaction: CompactionMiddlewar # summarizes at, measured as it measures, so the two cannot drift. self._token_threshold = compaction._token_threshold self._counter = compaction._summarizer.counter - self._notes_folder = re.compile(re.escape(scratchpad.notes_subdir) + r"(?![\w.-])") + self._notes_folder = _notes_pattern(scratchpad.notes_subdir) def wrap_model_call( self, @@ -306,7 +333,9 @@ def _check_in(self, state: Any) -> DurableUpdate | None: since_asked = _calls_since_notes(messages_in_view(state), folder, rows_end=True) if since_asked < NOTES_CHECK_IN_CALLS: return None - count = _calls_since_notes(state.get("messages") or (), folder, rows_end=False) + count = _calls_since_write( + state.get("messages") or (), folder, state.get("_notes_calls_trimmed") or 0 + ) return DurableUpdate( kind=NOTES_CHECK_IN_ROW_KIND, schema_version=UPDATE_SCHEMA_VERSION, @@ -359,15 +388,30 @@ def _calls_since_notes(messages: Sequence[Any], folder: re.Pattern[str], *, rows check-in reports runs from the last write alone, since that is what its text claims. """ + return _count_back(messages, folder, rows_end=rows_end)[0] + + +def _calls_since_write(messages: Sequence[Any], folder: re.Pattern[str], carried: int) -> int: + """The tool calls after the latest write to the notes, ``carried`` being + the count the window kept for the messages it trimmed before these.""" + count, wrote = _count_back(messages, folder, rows_end=False) + return count if wrote else count + carried + + +def _count_back( + messages: Sequence[Any], folder: re.Pattern[str], *, rows_end: bool +) -> tuple[int, bool]: + """``_calls_since_notes``, and whether a write (or row) ended the count + rather than the start of ``messages``.""" count = 0 for message in reversed(messages): if rows_end and _notes_row(message) is not None: - break + return count, True if isinstance(message, AIMessage): if any(_writes_notes(call, folder) for call in message.tool_calls): - break + return count, True count += len(message.tool_calls) - return count + return count, False def _writes_notes(call: Mapping[str, Any], folder: re.Pattern[str]) -> bool: diff --git a/src/ptc_agent/agent/middleware/compaction/offloading.py b/src/ptc_agent/agent/middleware/compaction/offloading.py index fc869f584..bf274cd37 100644 --- a/src/ptc_agent/agent/middleware/compaction/offloading.py +++ b/src/ptc_agent/agent/middleware/compaction/offloading.py @@ -22,7 +22,7 @@ strip_base64_from_messages, truncate_tool_call, ) -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.pointer import TranscriptTurns logger = logging.getLogger(__name__) @@ -120,7 +120,7 @@ def offload_view( effective = get_effective_messages(messages, state.get("_summarization_event")) arg_ids, read_ids = recorded_offloads(state) turns = ( - TranscriptTurns.of(transcript, messages) + TranscriptTurns.of(transcript, messages, window=Window.of(state)) if arg_ids and transcript is not None else None ) diff --git a/src/ptc_agent/agent/middleware/compaction/summarize.py b/src/ptc_agent/agent/middleware/compaction/summarize.py index d016f6ec4..20da255a4 100644 --- a/src/ptc_agent/agent/middleware/compaction/summarize.py +++ b/src/ptc_agent/agent/middleware/compaction/summarize.py @@ -36,6 +36,7 @@ is_run_boundary_message, is_summary_message, ) +from ptc_agent.agent.transcript import Window from ptc_agent.agent.transcript.pointer import TranscriptTurns from ptc_agent.agent.transcript.render import message_turns, visible_text @@ -178,13 +179,17 @@ def server_summary( *, raw_messages: Sequence[AnyMessage], turns: TranscriptTurns | None, + window: Window, ) -> Summary: """A summary built without a model: each request and the reply to it, the newest first to stay when they do not all fit, plus an earlier summary and the todo list. ``turns`` is set when the transcript is - readable, and names the file each turn is in.""" + readable, and names the file each turn is in. ``window`` is what was + trimmed from the head of ``raw_messages``.""" head = list(to_summarize[:1]) if to_summarize and is_summary_message(to_summarize[0]) else [] - numbers = turns.turns if turns is not None else message_turns(raw_messages) + numbers = ( + turns.turns if turns is not None else message_turns(raw_messages, base=window.runs) + ) covered: set[str] = set() sections: list[str] = [] budget = _SERVER_CHARS diff --git a/src/ptc_agent/agent/middleware/compaction/types.py b/src/ptc_agent/agent/middleware/compaction/types.py index 60ecbc7d0..8911376cd 100644 --- a/src/ptc_agent/agent/middleware/compaction/types.py +++ b/src/ptc_agent/agent/middleware/compaction/types.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Iterable from dataclasses import dataclass -from typing import Annotated, NotRequired +from typing import Annotated, Any, NotRequired from langchain_core.messages import MessageLikeRepresentation from langchain_core.messages.human import HumanMessage @@ -12,6 +12,7 @@ from langchain.agents.middleware.types import AgentState, PrivateStateAttr +from ptc_agent.agent.transcript.identity import WindowState from ptc_agent.config.agent import CompactionConfig from ptc_agent.core.paths import AGENT_HISTORY_DIRS, SandboxLayout @@ -84,8 +85,8 @@ class CompactionState(AgentState): Extends AgentState with private fields for tracking compaction events, offloaded tool call IDs, and when the model last answered (epoch seconds), - which gates Tier 1. The PrivateStateAttr annotation hides them from - input/output schemas. + which gates Tier 1, and what the window trimmed (``Window``). The + PrivateStateAttr annotation hides them from input/output schemas. Note: The ``_summarization_event`` field name is preserved because values are stored under that key in the LangGraph checkpointer — renaming it would @@ -100,6 +101,10 @@ class CompactionState(AgentState): _cached_input_tokens: Annotated[NotRequired[int], PrivateStateAttr] _cached_output_tokens: Annotated[NotRequired[int], PrivateStateAttr] _last_model_response_at: Annotated[NotRequired[float], PrivateStateAttr] + _window: Annotated[NotRequired[WindowState], PrivateStateAttr] + #: The summary anchor and cut the server last refused a trim at, so the + #: next turns do not ask again until a summary or a trim moves them. + _window_refused: Annotated[NotRequired[dict[str, Any]], PrivateStateAttr] # Tool names whose arguments carry large payloads (file contents, code strings) diff --git a/src/ptc_agent/agent/middleware/compaction/window.py b/src/ptc_agent/agent/middleware/compaction/window.py new file mode 100644 index 000000000..5786261e4 --- /dev/null +++ b/src/ptc_agent/agent/middleware/compaction/window.py @@ -0,0 +1,95 @@ +"""The sliding window: the main agent's checkpoint keeps the runs it can still see. + +A summary stands in for everything before its boundary, so the messages there +are never sent to the model again, yet every load of the thread still decodes +them. At each pass through ``before_agent`` (a fresh turn, or an orchestrator +re-entry) the list is cut back to the start of the transcript run the latest +summary's boundary falls in. A run is never split, and the current one is +never touched: the boundary is never past its start. The compaction event is +left as it was, since it re-finds its boundary by anchor; rewriting it would +read as a new compaction to everything that fingerprints it. + +The trimmed runs leave their count behind (``Window``), so every turn number +stays the whole thread's, the index text of the newest of them, which the +next summary lists, and their identity. Their messages live +on in the server's turn slices, so the window trims only once those are known +to hold exactly the messages it drops; without the server to ask (the CLI), +it never trims. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Any + +from langchain_core.messages import AnyMessage +from langgraph.types import Overwrite + +from ptc_agent.agent.middleware.compaction.offloading import recorded_offloads, tool_call_ids +from ptc_agent.agent.middleware.compaction.utils import resolve_cutoff_index +from ptc_agent.agent.state import ensure_message_ids +from ptc_agent.agent.transcript.classify import is_run_boundary_message +from ptc_agent.agent.transcript.identity import Window + +#: Whether stored turn slices, read in branch order, start with exactly the +#: messages this window trims. +WindowCoverage = Callable[[Window], Awaitable[bool]] + +#: The state update that keeps what a reader counting over the whole list +#: would have read in the trimmed messages, ``(trimmed, state) -> update``. +#: Each owner declares its own keys. +WindowCarry = Callable[[Sequence[AnyMessage], Mapping[str, Any]], dict[str, Any]] + + +def window_cut(messages: Sequence[AnyMessage], event: Any) -> int: + """Where the window may start: the start of the run holding ``event``'s + boundary, or 0 when nothing may go. + + An event without an anchor (legacy, or a summary that kept nothing) has + only a positional boundary, which a trim would move, so it waits for the + next compaction to write an anchored one. + """ + if not isinstance(event, Mapping): + return 0 + anchor = event.get("anchor_message_id") + if anchor is None: + return 0 + cutoff = resolve_cutoff_index(messages, event) + if not 0 <= cutoff < len(messages) or getattr(messages[cutoff], "id", None) != anchor: + return 0 + openers = (i for i in range(cutoff, -1, -1) if is_run_boundary_message(messages[i])) + start = next(openers, None) + # The first run also holds whatever came before its user message. + if start is None or not any(is_run_boundary_message(m) for m in messages[:start]): + return 0 + return start + + +def trim( + state: Mapping[str, Any], + cut: int, + carries: Sequence[WindowCarry] = (), + *, + head: Window, +) -> dict[str, Any]: + """The update trimming ``state``'s messages before ``cut``, ``head`` + being ``Window.of(state).extend(messages[:cut])``, the one the server + was asked about. + + The recorded offloads shrink to the calls still in the list, as a + compaction drops the ones it summarized. Ids are stamped on the kept + messages here, since an ``Overwrite`` skips the stamping a write through + the reducer gets. + """ + messages = state["messages"] + trimmed, kept = messages[:cut], ensure_message_ids(list(messages[cut:])) + update: dict[str, Any] = {"messages": Overwrite(kept), **head.update()} + live = tool_call_ids(kept) + args, reads = recorded_offloads(state) + if args - live: + update["_offloaded_tool_call_ids"] = args & live + if reads - live: + update["_offloaded_read_result_ids"] = reads & live + for carry in carries: + update.update(carry(trimmed, state)) + return update diff --git a/src/ptc_agent/agent/middleware/skills/content.py b/src/ptc_agent/agent/middleware/skills/content.py index a7881d40c..5929e17b4 100644 --- a/src/ptc_agent/agent/middleware/skills/content.py +++ b/src/ptc_agent/agent/middleware/skills/content.py @@ -445,7 +445,9 @@ def compute_already_loaded( return set() event = summarization_event - if not isinstance(event, dict) or not event.get("cutoff_index"): + # Presence, not the boundary's index: once the window trims the history + # (compaction.window), a summary's boundary can sit at index 0. + if not isinstance(event, dict): return loaded_set from ptc_agent.agent.middleware.compaction import get_effective_messages diff --git a/src/ptc_agent/agent/middleware/subagent_switch.py b/src/ptc_agent/agent/middleware/subagent_switch.py index a4e825c74..8f0f1d072 100644 --- a/src/ptc_agent/agent/middleware/subagent_switch.py +++ b/src/ptc_agent/agent/middleware/subagent_switch.py @@ -14,11 +14,11 @@ import asyncio import logging -from collections.abc import Awaitable, Callable -from typing import Any +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Annotated, Any, NotRequired from langchain.agents.middleware import AgentMiddleware -from langchain.agents.middleware.types import ToolCallRequest +from langchain.agents.middleware.types import AgentState, PrivateStateAttr, ToolCallRequest from langchain_core.messages import ToolMessage from langgraph.types import Command @@ -53,6 +53,13 @@ SubagentSwitchReader = Callable[[], Awaitable[bool | None]] +class SubagentSwitchState(AgentState): + """``_subagents_trimmed`` is set once the window trims a notice (see + ``compaction.window``), so that one was ever written outlives it.""" + + _subagents_trimmed: Annotated[NotRequired[bool], PrivateStateAttr] + + def last_announced(state: Any) -> bool | None: """The switch as the last notice the model can still read stated it. @@ -64,16 +71,32 @@ def last_announced(state: Any) -> bool | None: stated = last_stated(state, SUBAGENTS_ROW_KIND, "allowed") if stated is not None: return stated is not False - for message in state_get(state, "messages") or (): + if state_get(state, "_subagents_trimmed") or _has_notice(state_get(state, "messages") or ()): + return None + return True + + +def _has_notice(messages: Sequence[Any]) -> bool: + for message in reversed(messages): row = runtime_update_from_message(message) if row is not None and row.kind == SUBAGENTS_ROW_KIND: - return None - return True + return True + return False + + +def subagents_window_carry(trimmed: Sequence[Any], state: Mapping[str, Any]) -> dict[str, Any]: + """The window's carry (see ``compaction.window``) for the notices it + trims, which ``last_announced`` would otherwise read as never written.""" + if state.get("_subagents_trimmed") or not _has_notice(trimmed): + return {} + return {"_subagents_trimmed": True} class SubagentSwitchMiddleware(AgentMiddleware): """Announces the user's subagent switch in history and holds launches to it.""" + state_schema = SubagentSwitchState + def __init__(self, read: SubagentSwitchReader) -> None: super().__init__() self._read = read diff --git a/src/ptc_agent/agent/state.py b/src/ptc_agent/agent/state.py index d701814db..ababc40d6 100644 --- a/src/ptc_agent/agent/state.py +++ b/src/ptc_agent/agent/state.py @@ -4,9 +4,10 @@ defines `DeltaAgentState` (an `AgentState` whose `messages` uses `DeltaChannel` for O(1)-per-step checkpoint storage instead of re-serializing the full list). `DeltaAgentState` is structurally identical to deepagents 0.6.11's -`DeepAgentState` (`AgentState` + `DeltaChannel(reducer, snapshot_frequency=50)`); -we vendor the reducer rather than import the private `deepagents._messages_reducer` -symbol so on-disk reconstruction stays frozen to our release. The parity test +`DeepAgentState` (`AgentState` + `DeltaChannel(reducer)`) apart from a sparser +snapshot frequency; we vendor the reducer rather than import the private +`deepagents._messages_reducer` symbol so on-disk reconstruction stays frozen to +our release. The parity test `test_messages_delta_reducer.py::test_vendored_reducer_matches_deepagents` fails CI if our copy drifts. @@ -35,10 +36,15 @@ from langgraph.graph.ui import AnyUIMessage, ui_message_reducer from typing_extensions import NotRequired, Required -# A full snapshot blob is written every N updates, bounding delta replay depth -# (matches deepagents' tested default). Single source of truth for the -# `DeltaChannel` snapshot frequency, consumed by `DeltaAgentState` below. -MESSAGES_SNAPSHOT_FREQUENCY = 50 +# A full snapshot blob is written every N updates, bounding delta replay depth. +# Storage per message written is about 1 + W/N for a list of W messages. At 250 +# a long thread stores a fifth of the snapshot bytes it did at 50 while a tip +# load still reads one scan page; at 1000 loads cost half again as much. A +# thread from before DeltaChannel walks to its root until its first snapshot, +# so run scripts/ops/thin_message_snapshots.py --keep all over those first. +# Single source of truth for the `DeltaChannel` snapshot frequency, consumed +# by `DeltaAgentState` below. +MESSAGES_SNAPSHOT_FREQUENCY = 250 def ensure_message_ids(messages: list[AnyMessage]) -> list[AnyMessage]: diff --git a/src/ptc_agent/agent/transcript/__init__.py b/src/ptc_agent/agent/transcript/__init__.py index 0e37e4f0d..f730a9d19 100644 --- a/src/ptc_agent/agent/transcript/__init__.py +++ b/src/ptc_agent/agent/transcript/__init__.py @@ -5,18 +5,22 @@ TASK_META, TASKS_DIR, Directory, + EarlierTurnsMissing, TranscriptTarget, build_directory, load_manifest, transcript_subdir, ) +from ptc_agent.agent.transcript.identity import Window __all__ = [ "MANIFEST", "TASKS_DIR", "TASK_META", "Directory", + "EarlierTurnsMissing", "TranscriptTarget", + "Window", "build_directory", "load_manifest", "transcript_subdir", diff --git a/src/ptc_agent/agent/transcript/identity.py b/src/ptc_agent/agent/transcript/identity.py new file mode 100644 index 000000000..a9627371e --- /dev/null +++ b/src/ptc_agent/agent/transcript/identity.py @@ -0,0 +1,158 @@ +"""What the window has trimmed from the head of a message list: how many runs, +which messages, as one order-sensitive digest, and what a summary's index +says of the newest of them. + +The window trims runs from the head of the main agent's messages (see +``compaction.window``) once the server's turn slices hold them, and from then +on those slices are where the trimmed messages are read back from. A trim +checks that the slices hold exactly the messages it drops, and a read back +checks that it got them, both against this digest rather than by counting +turns or runs, which a thread can have in more than one shape. + +A message is keyed by what it says and the tool calls it makes or answers, +never by its id: one saved before ids were stable loads with a different id, +or none, from one checkpoint to the next. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +from langchain_core.messages import AnyMessage +from typing_extensions import TypedDict + +from ptc_agent.agent.transcript.render import split_runs, turn_map + +# The index lists the newest turns by request and one line stands for the +# rest, so a long thread does not grow every summary without bound. +_INDEX_LIMIT = 20 +# Each request is copied from the checkpoint at every compaction, so the +# user's instructions keep their wording here while summaries paraphrase +# them; the cap only stops a paste from riding in every summary. +_REQUEST_CHARS = 1_000 +_NO_REQUEST = "(no request)" + + +def _feed(step: Any, message: AnyMessage) -> None: + content = message.content + text = content if isinstance(content, str) else json.dumps( + content, sort_keys=True, default=str + ) + # A streamed message can be stored as its chunk class. + kind = message.type.removesuffix("MessageChunk").lower() + calls = [call.get("id") for call in getattr(message, "tool_calls", None) or ()] + answers = getattr(message, "tool_call_id", None) + data = text.encode("utf-8", "surrogatepass") + # The text's length frames it, so it is hashed as it is rather than + # escaped into the header: a transcript render hashes every message. + step.update(json.dumps([kind, calls, answers, len(data)], default=str).encode()) + step.update(data) + + +def extend(digest: str, messages: Iterable[AnyMessage]) -> str: + """``digest`` carried on over ``messages``. A list's digest is its + head's carried on over the rest, so a trim adds to the digest the trims + before it left instead of needing what they dropped.""" + for message in messages: + step = hashlib.blake2b(digest.encode(), digest_size=16) + _feed(step, message) + digest = step.hexdigest() + return digest + + +class WindowState(TypedDict): + """``Window`` as the checkpoint keeps it: ``requests`` is its index text, + the last line being run ``runs``'s, a list since a checkpoint keeps no + int keys.""" + + runs: int + messages: int + digest: str + requests: list[str] + + +@dataclass(frozen=True) +class Window: + """What the window has trimmed from the head of the main agent's + messages. Every numbering of turns counts on from ``runs``, so a file + name or an index line reads the same before and after a trim.""" + + #: How many transcript runs were trimmed; the list's first is ``runs + 1``. + runs: int = 0 + #: How many messages were trimmed, and their digest. + messages: int = 0 + digest: str = "" + #: The index text of the newest trimmed runs, by number. + earlier: Mapping[int, str] = field(default_factory=dict) + + @classmethod + def of(cls, state: Mapping[str, Any]) -> Window: + kept: WindowState | None = state.get("_window") + if not kept: + return cls() + runs, requests = kept["runs"], kept["requests"] + first = runs - len(requests) + 1 + return cls( + runs, + kept["messages"], + kept["digest"], + {first + i: line for i, line in enumerate(requests)}, + ) + + def update(self) -> dict[str, WindowState]: + """The state update that keeps it.""" + return { + "_window": WindowState( + runs=self.runs, + messages=self.messages, + digest=self.digest, + requests=list(self.earlier.values()), + ) + } + + def lines(self, requests: Mapping[int, str]) -> dict[int, str]: + """The index text of each turn, by number: the trimmed ones' as + kept, and the rest's from ``requests``, the text of the user message + that opened each (``turn_map``).""" + return {**self.earlier, **{n: _index_request(text) for n, text in requests.items()}} + + def extend(self, messages: Sequence[AnyMessage]) -> Window: + """This window with ``messages`` trimmed after it, which run up to + the start of a run. Hashes every one of them.""" + _, requests = turn_map(messages, base=self.runs) + runs = self.runs + len(split_runs(messages)) + return Window( + runs, + self.messages + len(messages), + extend(self.digest, messages), + newest_lines(self.lines(requests), runs), + ) + + def holds(self, messages: Sequence[AnyMessage]) -> bool: + """Whether ``messages`` are exactly the ones trimmed.""" + return len(messages) == self.messages and extend("", messages) == self.digest + + +def newest_lines(lines: Mapping[int, str], last: int) -> dict[int, str]: + """The index text of the newest turns up to ``last`` a summary lists, + by number.""" + start = max(1, last - _INDEX_LIMIT + 1) + return {n: lines.get(n, _NO_REQUEST) for n in range(start, last + 1)} + + +def _index_request(request: str) -> str: + """What the index says of a turn opened by ``request``.""" + return _one_line(request) if request else _NO_REQUEST + + +def _one_line(text: str) -> str: + from ptc_agent.agent.middleware.skills.content import skill_blocks_as_names + + line = " ".join(skill_blocks_as_names(text).split()) + if len(line) <= _REQUEST_CHARS: + return line + return line[: _REQUEST_CHARS - 1].rstrip() + "…" diff --git a/src/ptc_agent/agent/transcript/pointer.py b/src/ptc_agent/agent/transcript/pointer.py index 9f71d14c9..5ec1bdfca 100644 --- a/src/ptc_agent/agent/transcript/pointer.py +++ b/src/ptc_agent/agent/transcript/pointer.py @@ -13,12 +13,13 @@ import asyncio import logging from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING from langchain_core.messages import AnyMessage from ptc_agent.agent.transcript.classify import is_summary_message +from ptc_agent.agent.transcript.identity import Window, newest_lines from ptc_agent.agent.transcript.render import turn_map from ptc_agent.agent.transcript.store import TranscriptTarget, segment_file @@ -51,31 +52,27 @@ ) _INDEX_HEAD = "\n\nEach {unit}'s file and the request that opened it:\n" -# The index lists the newest turns by request and one line stands for the -# rest, so a long thread does not grow every summary without bound. -_INDEX_LIMIT = 20 -# Each request is copied from the checkpoint at every compaction, so the -# user's instructions keep their wording here while summaries paraphrase -# them; the cap only stops a paste from riding in every summary. -_REQUEST_CHARS = 1_000 - @dataclass(frozen=True) class TranscriptTurns: """Which transcript file holds each message. - Numbered over the agent's whole checkpoint list, as the renderer numbers - its files: a trimmed or summarized view would number from the wrong turn. + Numbered over the agent's checkpoint list from ``Window.runs`` on, as + the renderer numbers its files: a summarized view would number from the + wrong turn, and so would a list the window trimmed, counted from one. """ target: TranscriptTarget turns: Mapping[str, int] - #: The text of the user message that opened each turn, by turn number. - requests: Mapping[int, str] = field(default_factory=dict) + #: The index text of each turn, by number (``Window.lines``). + lines: Mapping[int, str] @classmethod - def of(cls, target: TranscriptTarget, messages: Sequence[AnyMessage]) -> TranscriptTurns: - return cls(target, *turn_map(messages)) + def of( + cls, target: TranscriptTarget, messages: Sequence[AnyMessage], *, window: Window + ) -> TranscriptTurns: + turns, requests = turn_map(messages, base=window.runs) + return cls(target, turns, window.lines(requests)) def file(self, message_id: str | None) -> str | None: number = self.turns.get(message_id or "") @@ -97,10 +94,10 @@ def index(self, last: int) -> list[str]: them and the agent find them. """ unit = self.target.unit - start = max(1, last - _INDEX_LIMIT + 1) + newest = newest_lines(self.lines, last) + earlier = next(iter(newest), 1) - 1 lines = [] - if start > 1: - earlier = start - 1 + if earlier: names = f"`{segment_file(unit, 1)}`" if earlier > 1: names += f" to `{segment_file(unit, earlier)}`" @@ -108,12 +105,7 @@ def index(self, last: int) -> list[str]: f"- {names}: {earlier} earlier {unit}{'s' if earlier > 1 else ''}, " f"not listed here; `{self.target.manifest}` lists every file" ) - for number in range(start, last + 1): - request = self.requests.get(number) - lines.append( - f"- `{segment_file(unit, number)}`: " - f"{_one_line(request) if request else '(no request)'}" - ) + lines += [f"- `{segment_file(unit, n)}`: {text}" for n, text in newest.items()] return lines @@ -127,15 +119,6 @@ class SummarySpan: gap: tuple[int, int] | None = None -def _one_line(text: str) -> str: - from ptc_agent.agent.middleware.skills.content import skill_blocks_as_names - - line = " ".join(skill_blocks_as_names(text).split()) - if len(line) <= _REQUEST_CHARS: - return line - return line[: _REQUEST_CHARS - 1].rstrip() + "\u2026" - - def _mount(backend: SandboxBackend | None) -> MountHandle | None: return backend.livefs if backend is not None else None @@ -155,10 +138,13 @@ async def aexport_transcript( *, workspace_id: str | None, budget: float = _EXPORT_TIMEOUT, + window: Window, ) -> TranscriptTarget | None: """Bring the transcript up to date within ``budget`` seconds (at most ``_EXPORT_TIMEOUT``); the one a summary may point at, or None unless the save landed and the mount serves ``workspace_id``'s folder. Never raises. + ``window`` is what was trimmed from the head of ``messages``, whose + files the store keeps. The turn-end export comes only after this agent finishes, so a pointer left on a failed save sends it for history that is not there for the @@ -174,7 +160,7 @@ async def save_and_settle() -> tuple[bool, MountHandle | None]: # A task group, not gather: when one fails, the other is cancelled # rather than left saving past the timeout. async with asyncio.TaskGroup() as group: - saved = group.create_task(mount.save_transcript(transcript, messages)) + saved = group.create_task(mount.save_transcript(transcript, messages, window=window)) settled = group.create_task(backend.settled_livefs(workspace_id)) return saved.result(), settled.result() diff --git a/src/ptc_agent/agent/transcript/render.py b/src/ptc_agent/agent/transcript/render.py index 24c2cdfce..3ddb14a95 100644 --- a/src/ptc_agent/agent/transcript/render.py +++ b/src/ptc_agent/agent/transcript/render.py @@ -251,17 +251,24 @@ def _args_checksum(args: Any) -> str: ) -def message_turns(messages: Iterable[AnyMessage]) -> dict[str, int]: +def message_turns(messages: Iterable[AnyMessage], *, base: int) -> dict[str, int]: """The turn each message belongs to, by message id.""" - return turn_map(messages)[0] + return turn_map(messages, base=base)[0] -def turn_map(messages: Iterable[AnyMessage]) -> tuple[dict[str, int], dict[int, str]]: +def turn_map( + messages: Iterable[AnyMessage], *, base: int +) -> tuple[dict[str, int], dict[int, str]]: """The turn each message belongs to, by message id, and the text of the - user message that opened each turn, by turn number.""" + user message that opened each turn, by turn number. + + ``base`` is how many turns the window trimmed from the head of + ``messages`` (``Window.runs``): numbers stay those of the whole thread, + the ones its files and earlier summaries already name. + """ turns: dict[str, int] = {} requests: dict[int, str] = {} - for number, run in enumerate(split_runs(messages), start=1): + for number, run in enumerate(split_runs(messages), start=base + 1): for message in run: if message.id: turns[message.id] = number diff --git a/src/ptc_agent/agent/transcript/store.py b/src/ptc_agent/agent/transcript/store.py index b58f33394..1f61f00a6 100644 --- a/src/ptc_agent/agent/transcript/store.py +++ b/src/ptc_agent/agent/transcript/store.py @@ -15,6 +15,7 @@ from langchain_core.messages import AnyMessage +from ptc_agent.agent.transcript.identity import Window, extend from ptc_agent.agent.transcript.render import ( SCHEMA_VERSION, Segment, @@ -105,7 +106,7 @@ def dump_manifest(manifest: dict[str, Any]) -> str: return json.dumps(manifest, ensure_ascii=False, indent=1, default=str) -def _entry(segment: Segment, unit: str, shape: str) -> dict[str, Any]: +def _entry(segment: Segment, unit: str, shape: str, through: str) -> dict[str, Any]: return { "file": segment_file(unit, segment.number), unit: segment.number, @@ -116,13 +117,18 @@ def _entry(segment: Segment, unit: str, shape: str) -> dict[str, Any]: "bytes": len(segment.data), "sha256": segment.sha256, "shape": shape, + "through": through, } def _reusable(entry: Any, unit: str, number: int, shape: str) -> bool: + return _carriable(entry, unit, number) and entry.get("shape") == shape + + +def _carriable(entry: Any, unit: str, number: int) -> bool: return ( isinstance(entry, dict) - and entry.get("shape") == shape + and isinstance(entry.get("shape"), str) and entry.get("file") == segment_file(unit, number) and isinstance(entry.get("sha256"), str) and len(entry["sha256"]) == 64 @@ -142,12 +148,18 @@ class Directory: rendered: dict[str, bytes] +class EarlierTurnsMissing(LookupError): + """The previous manifest cannot stand in for the turns trimmed from the + head of the messages: they have to be passed to render them.""" + + def build_directory( messages: list[AnyMessage], *, unit: str = "turn", header: dict[str, Any] | None = None, previous: dict[str, Any] | None = None, + window: Window, ) -> Directory: """Render the segments that changed since ``previous``, the manifest the stored copy was rendered with, and carry the rest over from it. @@ -155,6 +167,15 @@ def build_directory( A turn end adds to one turn, so this renders one segment where a whole thread's render costs its size in CPU. With no previous manifest, or one of another schema, every segment renders. + + ``window`` was trimmed from the head of ``messages`` (see + ``compaction.window``): its runs' entries are carried from ``previous`` + as they are, and the first turn of ``messages`` is ``window.runs + 1``. + Each entry records the digest of every message through its turn + (``identity``), and the last carried one must match ``window``'s: a + copy rendered from another branch (a live save labelled with an edit's + fork point) shares its numbering, not its turns. Raises + ``EarlierTurnsMissing`` when ``previous`` does not hold them. """ held: list[Any] = [] if previous and previous.get("schema") == SCHEMA_VERSION: @@ -162,14 +183,34 @@ def build_directory( names: dict[str, str] = {} entries: list[dict[str, Any]] = [] rendered: dict[str, bytes] = {} - shape = "" - for number, run in enumerate(split_runs(messages), start=1): + shape = through = "" + base = window.runs + if base: + carried = held[:base] + if ( + len(carried) < base + or not all( + _carriable(entry, unit, number) + for number, entry in enumerate(carried, start=1) + ) + or carried[-1].get("through") != window.digest + ): + raise EarlierTurnsMissing(f"{len(carried)} of {base} trimmed turns held") + entries.extend(carried) + # The key chains through every turn before it, the trimmed ones too. + shape, through = carried[-1]["shape"], window.digest + for number, run in enumerate(split_runs(messages), start=base + 1): shape = segment_shape(number, run, unit=unit, previous=shape, tool_names=names) + through = extend(through, run) entry = held[number - 1] if number <= len(held) else None if not _reusable(entry, unit, number, shape): segment = render_segment(number, run, unit=unit, tool_names=names) - entry = _entry(segment, unit, shape) + entry = _entry(segment, unit, shape, through) rendered[entry["file"]] = segment.data + elif entry.get("through") != through: + # Its render reads the same in messages that differ where no + # render looks (reasoning, say), or it predates the field. + entry = {**entry, "through": through} entries.append(entry) manifest = {"schema": SCHEMA_VERSION, **(header or {}), "segments": entries} return Directory( diff --git a/src/ptc_agent/core/sandbox/livefs_mount.py b/src/ptc_agent/core/sandbox/livefs_mount.py index 2c71fd1a1..18a525291 100644 --- a/src/ptc_agent/core/sandbox/livefs_mount.py +++ b/src/ptc_agent/core/sandbox/livefs_mount.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from langchain_core.messages import AnyMessage - from ptc_agent.agent.transcript import TranscriptTarget + from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.core.sandbox.ptc_sandbox import PTCSandbox # Covers a cold start's wait for its mount, plus the probe and links. @@ -82,10 +82,11 @@ async def report( ``output`` showing the mount dead restarts it.""" async def save_transcript( - self, target: TranscriptTarget, messages: Sequence[AnyMessage] + self, target: TranscriptTarget, messages: Sequence[AnyMessage], *, window: Window ) -> bool: - """Store one agent's transcript from the messages in hand, where the - mount serves it; whether it landed.""" + """Store one agent's transcript from the messages in hand, ``window`` + having been trimmed from their head, where the mount serves it; + whether it landed.""" async def through_mount( diff --git a/src/server/AGENTS.md b/src/server/AGENTS.md index 033484a5b..4fa8aaf3c 100644 --- a/src/server/AGENTS.md +++ b/src/server/AGENTS.md @@ -14,7 +14,7 @@ Inside those tiers, the run machinery groups into **domain namespaces**: `app/th - **Two separate async psycopg3 pools, opened/closed in the `app/setup.py` lifespan** — the LangGraph checkpointer (workflow state; the LangGraph Store shares it) and the conversation-history (app-data) pool. Don't reuse one for the other. - **Mid-flight reconnection is Redis-only, not the checkpoint.** Live events are buffered in a per-run Redis **Stream** `workflow:stream:{thread_id}:{run_id}` (XADD + `MAXLEN` trim, 24h TTL, 150k-event cap). `GET /api/v1/threads/{id}/messages/stream` XREAD-BLOCK-replays the buffer then tails live; `last_event_id` is a resume cursor (resumes at seq N+1), not post-hoc dedup. An abandoned running workflow is reaped after ~6h; a completed workflow's reconnect returns `410` once its task-info key expires. -- **Historical replay is projected from the LangGraph checkpoint, not from stored SSE events** (re-homed in #315). `GET /api/v1/threads/{id}/messages/replay` (`source=auto`) rebuilds the transcript from checkpoint state via `CheckpointHistoryReader` + a pure projector (`services/history/`), falling back to persisted `conversation_responses.sse_events` only when checkpoint coverage is missing (that column is still dual-written every turn, but it's now the transitional fallback). **So the durable replay contract is the checkpoint state schema: `messages` (a `DeltaChannel`) + the `ui` channel** (`langgraph.graph.ui`; an id-keyed upsert accumulator of compact, non-rederivable records — image-URL maps, `model_fallback` notices). +- **Historical replay is projected from the LangGraph checkpoint, not from stored SSE events** (re-homed in #315). `GET /api/v1/threads/{id}/messages/replay` (`source=auto`) rebuilds the transcript from checkpoint state via `CheckpointHistoryReader` + a pure projector (`services/history/`), falling back to persisted `conversation_responses.sse_events` only when checkpoint coverage is missing (that column is still dual-written every turn, but it's now the transitional fallback). A turn settled before its row carried replay facts also merges what only its stored events kept (widget data, evicted results, reasoning durations, signal rows in position); `scripts/utils/backfill_replay_facts.py` resolves those once into the row's `replay_facts.legacy` (`services/history/replay/legacy.py`), after which its replay never reads the column, and any write to the row's stored events drops them until the next run. **So the durable replay contract is the checkpoint state schema: `messages` (a `DeltaChannel`) + the `ui` channel** (`langgraph.graph.ui`; an id-keyed upsert accumulator of compact, non-rederivable records — image-URL maps, `model_fallback` notices). - **What silently breaks replay of already-stored checkpoints:** renaming a `ui`-record's `name`/`props`, or emitting one without a stable pre-stamped `id` (the reducer upserts by id — keep the projector's field whitelists in sync); putting rederivable or bulk data in graph state (the projector rebuilds artifacts from tool-call *args* + message `additional_kwargs`, and >32KB widget data must be a content-addressed `data_ref`, not inlined); or adding a chat-wire SSE type without a checkpoint- or table-sourced home — `tests/unit/server/services/history/test_event_ledger.py` fails CI until every type is classified, so build the replay source before you emit. - **Error funnel (`handlers/chat/error_handling.py`, `classify_error()`):** recoverable errors (DB/connection drop, timeout, network, API 5xx/429) emit a `retry` SSE event (`auto_retry`; thread → `interrupted`, and the client/gateway re-submits) up to 3× (`get_max_workflow_retries`), then convert to a terminal `error`. Non-recoverable *code bugs* (`AttributeError`/`NameError`/`TypeError`/`ImportError`/`SyntaxError`/`KeyError`) fail immediately with a stable `error` payload. Each of these frames names its `recovery`: `retry` when a failed run is on the ledger (`POST .../retry` with its `run_id`; a message that never reached the graph runs again from the turn's query row, else the checkpoint replays) or `resend` when no run started (the client sends the same request again). A reader error after the executor holds the run emits neither: the stream ends and the client reconnects. Don't add ad-hoc try/except that swallows this path. - **Cancelled ≠ disconnected:** an explicit cancel stops the workflow; a client disconnect lets it keep running in the background (`LocalRunExecutor` keeps executing; the run's `in_progress` ledger row is the durable record), and a reconnect resumes the live stream. diff --git a/src/server/app/public.py b/src/server/app/public.py index b8c91935d..841affb28 100644 --- a/src/server/app/public.py +++ b/src/server/app/public.py @@ -26,7 +26,6 @@ """ import json -import logging from typing import Any from fastapi import APIRouter, HTTPException, Query, Response @@ -46,19 +45,12 @@ signed_url_expires_at, with_preview_path, ) -from src.server.database.conversation import ( - get_queries_for_thread, - get_responses_for_thread, -) +from src.server.database.conversation.replay_rows import get_replay_thread_data from src.server.database.share_links import KIND_APP from src.server.services.file_grants import grant_prefix, mint_file_grant, seconds_left from src.server.utils.api import PageViewer, Viewer -from src.server.services.history.replay.items import run_completed_at -from src.server.services.history.replay.stopped import stop_close_item from src.server.services.share_redaction import ShareRedaction -logger = logging.getLogger(__name__) - router = APIRouter(prefix="/api/v1/public", tags=["Public Sharing"]) # The file routes live next door and mount here, so the token prefix and the # tag stay in one place. @@ -164,120 +156,44 @@ async def get_shared_thread_metadata( async def replay_shared_thread(share_token: str): """Replay a shared thread as SSE. No auth required. - Same replay logic as the authenticated endpoint, but resolves - thread via share_token and strips sensitive fields. + The same assembly as the owner's replay, so a viewer sees the turns the + owner sees; the thread resolves through the share token and every event + passes ``ShareRedaction`` on the way out. """ + from src.server.services.history.replay import finish_lines, read_replay_page + thread = await get_shared_thread(share_token) thread_id = str(thread["conversation_thread_id"]) - queries, _ = await get_queries_for_thread(thread_id) - responses, _ = await get_responses_for_thread(thread_id) - responses_by_turn = {r.get("turn_index"): r for r in responses if isinstance(r, dict)} - + rows = await get_replay_thread_data(thread_id) + if rows is None: + raise HTTPException(status_code=404, detail="Shared thread not found") + page, _ = await read_replay_page(rows) # Public replay has no /status reconciliation at all, so without the # stamp its task cards would be stuck "running" forever (see - # history/task_status.py). Only the whitelisted status value is added. - from src.server.services.history.task_status import ( - collect_task_ids, - resolve_task_details, - stamp_task_artifact_data, - ) - - stored_events = [ - item - for r in responses_by_turn.values() - if isinstance(r.get("sse_events"), list) - for item in r["sse_events"] - if isinstance(item, dict) - ] - redaction = ShareRedaction(stored_events) - - task_details: dict[str, dict] = {} - try: - task_details = await resolve_task_details( - thread_id, collect_task_ids(stored_events) - ) - except Exception: - logger.warning( - f"[PUBLIC REPLAY] task-status stamping failed for {thread_id}", - exc_info=True, - ) + # history/task_status.py). Only the whitelisted status value is added, + # which redaction leaves as it is. + await finish_lines(thread_id, page.lines, status_only=True) + items = [line.item() for line in page.lines] + redaction = ShareRedaction(items) async def event_generator(): seq = 0 - - for q in queries: - if not isinstance(q, dict): - continue - - turn_index = q.get("turn_index") - seq += 1 - - content, metadata = redaction.query(q) - payload = { - "thread_id": thread_id, - "turn_index": turn_index, - "content": content, - "timestamp": q.get("created_at"), - "metadata": metadata, - } - # Tag system queries so the frontend can hide the user bubble - query_type = q.get("type") - if query_type == "system": - payload["query_type"] = "system" - # The turn's end, paired with the query timestamp above to give the - # fold row its duration. This payload is hand-built rather than - # taken from the replay builder, so the field has to be mirrored - # here or a shared transcript folds with no duration to show. The - # run id the builder also stamps stays out: it exists for the - # report-back catch-up, which a public viewer never runs. - completed_at = run_completed_at(responses_by_turn.get(turn_index)) - if completed_at is not None: - payload["run_completed_at"] = completed_at - - yield ( - f"id: {seq}\n" - f"event: user_message\n" - f"data: {json.dumps(payload, ensure_ascii=False, default=str)}\n\n" - ) - - response = responses_by_turn.get(turn_index) - if not response: - continue - - sse_events = response.get("sse_events") - if not isinstance(sse_events, list): - sse_events = [] - # A stop during bring-up can archive nothing and still owes its close. - stop_close = stop_close_item(thread_id, response, sse_events) - if stop_close: - sse_events = [*sse_events, stop_close] - - for item in sse_events: - if not isinstance(item, dict): - continue - event_type = item.get("event") - data = item.get("data") - if not event_type or not isinstance(data, dict): - continue + for item in items: + event_type, data = item["event"], item["data"] + if event_type == "user_message": + replay_data = _shared_user_message(redaction, data) + else: replay_data = redaction.event(event_type, data) if replay_data is None: continue - - seq += 1 replay_data.setdefault("thread_id", thread_id) - replay_data["turn_index"] = turn_index - replay_data["response_id"] = str(response.get("conversation_response_id")) - if task_details: - replay_data = stamp_task_artifact_data( - replay_data, task_details, status_only=True - ) - - yield ( - f"id: {seq}\n" - f"event: {event_type}\n" - f"data: {json.dumps(replay_data, ensure_ascii=False, default=str)}\n\n" - ) + seq += 1 + yield ( + f"id: {seq}\n" + f"event: {event_type}\n" + f"data: {json.dumps(replay_data, ensure_ascii=False, default=str)}\n\n" + ) seq += 1 yield f"id: {seq}\nevent: replay_done\ndata: {json.dumps({'thread_id': thread_id}, default=str)}\n\n" @@ -287,3 +203,23 @@ async def event_generator(): media_type="text/event-stream", headers={"Cache-Control": "no-cache"}, ) + + +def _shared_user_message(redaction: ShareRedaction, data: dict[str, Any]) -> dict[str, Any]: + """A turn's user message as a viewer may read it. + + The run id stays out: it exists for the owner's report-back catch-up, + which a public viewer never runs. + """ + content, metadata = redaction.query( + { + "type": data.get("query_type"), + "content": data.get("content"), + "metadata": data.get("metadata"), + } + ) + return { + **{k: v for k, v in data.items() if k != "run_id"}, + "content": content, + "metadata": metadata, + } diff --git a/src/server/app/threads/messaging.py b/src/server/app/threads/messaging.py index 8afff9641..26e00c06b 100644 --- a/src/server/app/threads/messaging.py +++ b/src/server/app/threads/messaging.py @@ -824,139 +824,109 @@ async def replay_thread_messages( description=( "Replay source: 'checkpoint' projects the transcript from LangGraph " "checkpoints, 'sse' replays persisted sse_events, 'auto' prefers " - "checkpoint and falls back to sse when coverage is incomplete." + "checkpoint, replaying from sse a turn it cannot project, or the " + "whole thread when coverage is incomplete." ), ), limit: int | None = Query( None, ge=1, description=( - "Windowed replay: build only the most recent N turns from " - "checkpoints (bounds initial-load latency to the window). " - "Checkpoint-sourced only; ignored for source='sse'." + "Paged replay: only the most recent N turns (before " + "``before_turn`` when given), plus a leading ``history_page`` " + "event naming the oldest turn sent and whether older ones remain." + ), + ), + before_turn: int | None = Query( + None, + ge=0, + description=( + "Paged replay: only turns older than this turn_index, the " + "previous page's ``history_page.first_turn_index``." ), ), ): """Replay a thread as SSE. Stream includes: + - history_page: first, only when ``limit`` or ``before_turn`` is given - user_message: emitted once per turn_index (query content) - message_chunk/tool_* events: projected from checkpoints or emitted from stored sse_events, per ``source`` + - snapshot: cursors for the runs replay could not project (newest page only) - replay_done: terminal sentinel """ try: - owner_id, thread, queries, responses, usages, provenance = ( - await get_replay_thread_data(thread_id) - ) + rows = await get_replay_thread_data(thread_id) # Preserve existing 404/403 semantics from require_thread_owner - if owner_id is None: + if rows is None: raise HTTPException(status_code=404, detail="Thread not found") - if owner_id != x_user_id: + if rows.owner_id != x_user_id: raise HTTPException(status_code=403, detail="Forbidden") - if not thread: - raise HTTPException( - status_code=404, detail=f"Thread not found: {thread_id}" - ) - - responses_by_turn = { - r.get("turn_index"): r for r in responses if isinstance(r, dict) - } from src.server.services.history.replay import ( CheckpointReplayUnavailable, - build_checkpoint_replay_items, - build_sse_replay_items, - ) - - checkpoint_items: list[dict] | None = None - if source in ("auto", "checkpoint"): - try: - if thread.get("latest_checkpoint_id") is None: - # The commit pointer (stamped at turn persist) is the only - # tip checkpoint replay may read — without it the reader - # would walk the newest checkpoint, which mid-run is - # uncommitted partial state. - raise CheckpointReplayUnavailable( - "thread has no committed checkpoint pointer" - ) - checkpoint_items = await build_checkpoint_replay_items( - thread_id, - queries, - responses_by_turn, - branch_tip_checkpoint_id=thread.get("latest_checkpoint_id"), - last_n_turns=limit, - usages=usages, - provenance=provenance, - ) - except CheckpointReplayUnavailable as e: - if source == "checkpoint": - raise HTTPException( - status_code=409, - detail=f"Checkpoint replay unavailable: {e}", - ) - logger.info( - f"[REPLAY] Checkpoint replay unavailable for {thread_id}, " - f"falling back to sse: {e}" - ) - except HTTPException: - raise - except Exception as e: - if source == "checkpoint": - raise - logger.warning( - f"[REPLAY] Checkpoint replay failed for {thread_id}, " - f"falling back to sse: {e}", - exc_info=True, - ) - - replay_items = ( - checkpoint_items - if checkpoint_items is not None - else build_sse_replay_items(thread_id, queries, responses_by_turn) + finish_lines, + read_replay_page, ) - # After source selection and outside the projection cache: task - # artifacts replay their spawn-time payload, so the card's status is - # stamped here from current liveness (see history/task_status.py). - from src.server.services.history.task_status import ( - stamp_replay_task_status, - ) + try: + page, resolved_source = await read_replay_page( + rows, source=source, before_turn=before_turn, limit=limit + ) + except CheckpointReplayUnavailable as e: + raise HTTPException( + status_code=409, + detail=f"Checkpoint replay unavailable: {e}", + ) - await stamp_replay_task_status(thread_id, replay_items) + # After source selection and outside the stored lines: task artifacts + # replay their spawn-time payload, so the card's status is stamped + # here from current liveness (see history/task_status.py). + replay_lines = page.lines + await finish_lines(thread_id, replay_lines) + paged = limit is not None or before_turn is not None async def event_generator(): seq = 0 - for item in replay_items: + if paged: seq += 1 + history_page = { + "first_turn_index": page.first_turn_index, + "has_more": page.has_more, + } yield ( f"id: {seq}\n" - f"event: {item['event']}\n" - f"data: {json.dumps(item['data'], ensure_ascii=False, default=str)}\n\n" + f"event: history_page\n" + f"data: {json.dumps(history_page, default=str)}\n\n" ) + for line in replay_lines: + seq += 1 + yield f"id: {seq}\nevent: {line.event}\ndata: {line.data_json}\n\n" # Cursors for the runs replay could not project — an in-flight run # belongs to its stream, so the snapshot hands the client where to # resume instead of its content. Additive: a snapshot outage (None) # just omits the frame, and v1 clients ignore the unknown event. - from src.server.services.history.snapshot import ( - build_thread_snapshot, - ) - - snapshot = await build_thread_snapshot(thread_id) - if snapshot is not None: - seq += 1 - yield ( - f"id: {seq}\n" - f"event: snapshot\n" - f"data: {json.dumps(snapshot, ensure_ascii=False, default=str)}\n\n" + # An older page holds no in-flight run. + if before_turn is None: + from src.server.services.history.snapshot import ( + build_thread_snapshot, ) + snapshot = await build_thread_snapshot(thread_id) + if snapshot is not None: + seq += 1 + yield ( + f"id: {seq}\n" + f"event: snapshot\n" + f"data: {json.dumps(snapshot, ensure_ascii=False, default=str)}\n\n" + ) + seq += 1 yield f"id: {seq}\nevent: replay_done\ndata: {json.dumps({'thread_id': thread_id}, default=str)}\n\n" - resolved_source = "checkpoint" if checkpoint_items is not None else "sse" return StreamingResponse( observe_replay_stream(event_generator(), source="private"), media_type="text/event-stream", diff --git a/src/server/app/threads/tasks.py b/src/server/app/threads/tasks.py index 15da781b1..43986d4bb 100644 --- a/src/server/app/threads/tasks.py +++ b/src/server/app/threads/tasks.py @@ -60,7 +60,7 @@ async def get_subagent_task_history( transcript belongs to its live stream). """ await auth_api.require_thread_owner(thread_id, x_user_id) - from src.server.services.history.replay.task_lane import ( + from src.server.services.history.replay.run_lane import ( project_task_transcript, ) diff --git a/src/server/contracts/status.py b/src/server/contracts/status.py index a94450a3b..1750e2934 100644 --- a/src/server/contracts/status.py +++ b/src/server/contracts/status.py @@ -186,12 +186,24 @@ def _classify_one(request: dict) -> Optional[str]: return "approval_required" if request.get("name") else None -def classify_interrupt_reason(interrupts: Any) -> Optional[str]: - """Classify HITL interrupt payloads into the ``interrupt_reason`` column. - - One authority for the spelling: the live streaming path and the recovery - scanner must never drift apart on it. ``None`` when no payload carries a - discriminator at all, which the column already stores as NULL. +def classify_interrupts(interrupts: Any) -> tuple[Optional[str], Optional[str]]: + """Classify HITL interrupt payloads -> (``interrupt_reason``, the credit + pause's denial). + + One authority for both: the live producer, the finalize fallback and the + recovery scanner must never drift apart on the spelling or on which + words a pause stamps. The reason is ``None`` when no payload carries a + discriminator, which the column stores as NULL. The denial is the quota + service's own words, relayed verbatim or not at all; the last pause wins. """ - reasons = [r for r in map(_classify_one, _action_requests(interrupts)) if r] - return min(reasons, key=INTERRUPT_REASONS.index, default=None) + reasons: list[str] = [] + message: Optional[str] = None + for request in _action_requests(interrupts): + reason = _classify_one(request) + if reason: + reasons.append(reason) + if reason == INTERRUPT_REASON_CREDIT_PAUSE and isinstance( + request.get("message"), str + ): + message = request["message"] + return min(reasons, key=INTERRUPT_REASONS.index, default=None), message diff --git a/src/server/database/automation_executions.py b/src/server/database/automation_executions.py index 6602877d7..8f5b80ae6 100644 --- a/src/server/database/automation_executions.py +++ b/src/server/database/automation_executions.py @@ -470,16 +470,19 @@ async def get_settling_run(run_id: str) -> Optional[Dict[str, Any]]: """The columns of a run's ledger row that settling its firing reads. ``sse_events`` is the turn's whole event archive, which can run to - megabytes, so it is read only for a credit pause, whose interrupt - carries the denial the user is told. + megabytes, so it is read only for a credit pause settled before its + finalize stamped the denial into ``metadata``: the pause's stored + interrupt is then the one place that still has it. """ async with get_db_connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute(""" SELECT conversation_response_id, status, interrupt_reason, metadata, errors, - CASE WHEN interrupt_reason = %s THEN sse_events END - AS sse_events + CASE WHEN interrupt_reason = %s + AND NOT (COALESCE(metadata, '{}'::jsonb) + ? 'credit_pause_message') + THEN sse_events END AS sse_events FROM conversation_responses WHERE conversation_response_id = %s """, (INTERRUPT_REASON_CREDIT_PAUSE, run_id)) diff --git a/src/server/database/conversation/__init__.py b/src/server/database/conversation/__init__.py index 202b7d6f8..f1d07ca4e 100644 --- a/src/server/database/conversation/__init__.py +++ b/src/server/database/conversation/__init__.py @@ -9,6 +9,7 @@ from src.server.database.conversation._sql import ( _RESPONSE_COLUMNS, _SETTLED_ATTEMPTS, + settled_attempts, ) from src.server.database.conversation.errors import ( EXTERNAL_ID_CONFLICT_ERROR_TYPE, @@ -75,6 +76,7 @@ upsert_feedback, ) from src.server.database.conversation.replay_rows import ( + get_replay_responses, get_replay_thread_data, ) @@ -99,6 +101,7 @@ "get_latest_turn_index", "get_queries_for_thread", "get_recent_responses_for_thread", + "get_replay_responses", "get_replay_thread_data", "get_responses_for_thread", "get_thread_auth_meta", @@ -120,6 +123,7 @@ "promote_flash_thread", "_RESPONSE_COLUMNS", "_SETTLED_ATTEMPTS", + "settled_attempts", "_sse_has_provenance", "_sync_provenance_for_response", "thread_exists_key", diff --git a/src/server/database/conversation/_sql.py b/src/server/database/conversation/_sql.py index 522bb084a..431dd7ab0 100644 --- a/src/server/database/conversation/_sql.py +++ b/src/server/database/conversation/_sql.py @@ -56,13 +56,21 @@ def stored_subagents_allowed(value_sql: str, workspace_id_sql: str) -> str: # ``usage_settled_at`` rides along because it is the settle instant itself, and # replay pairs it with the query timestamp to say how long a turn took. Readers # that go through ``_SETTLED_ATTEMPTS`` directly get it from the ``*``; one that -# projects this list would otherwise be left with the start-plus-duration +# projects these lists would otherwise be left with the start-plus-duration # estimate, which measures from before the row existed. -_RESPONSE_COLUMNS = ( +# +# Replay reads every settled row of a thread on each open, but needs a row's +# ``sse_events`` and ``replay_facts`` only for a turn it projects afresh. Those +# re-read their rows whole by id (``get_replay_responses``), so a cached open +# never moves the stored events, or the legacy facts backfilled from them, of +# its whole history over the wire. A row read without ``sse_events`` has them +# still to read. +_LIGHT_RESPONSE_COLUMNS = ( "conversation_response_id, conversation_thread_id, turn_index, status, " "interrupt_reason, metadata, warnings, errors, execution_time, created_at, " - "usage_settled_at, sse_events, attempt_no, retry_of_run_id" + "usage_settled_at, attempt_no, retry_of_run_id" ) +_RESPONSE_COLUMNS = f"{_LIGHT_RESPONSE_COLUMNS}, sse_events, replay_facts" # 1.6: retries append attempt rows at the SAME turn_index, and the live run is # an in_progress row. History readers must see ONE row per turn — the newest @@ -77,9 +85,17 @@ def stored_subagents_allowed(value_sql: str, workspace_id_sql: str) -> str: # append) writes a new tuple version. It is a tuple-header read, so it costs # nothing — and a subquery cannot expose it, which is why it belongs here # rather than at the replay call site. -_SETTLED_ATTEMPTS = """ - SELECT DISTINCT ON (turn_index) *, xmin +def settled_attempts(columns: str, thread: str = "%s") -> str: + """The query for the row each turn of a thread reads. ``thread`` is what + the thread id is matched against: a bound parameter, or an outer column + such as ``t.conversation_thread_id`` when this is a lateral join, so a + job counting turns across threads selects the rows replay reads.""" + return f""" + SELECT DISTINCT ON (turn_index) {columns}, xmin FROM conversation_responses - WHERE conversation_thread_id = %s AND status <> 'in_progress' + WHERE conversation_thread_id = {thread} AND status <> 'in_progress' ORDER BY turn_index ASC, attempt_no DESC """ + + +_SETTLED_ATTEMPTS = settled_attempts("*") diff --git a/src/server/database/conversation/replay_rows.py b/src/server/database/conversation/replay_rows.py index f92dfdb73..32f3335ec 100644 --- a/src/server/database/conversation/replay_rows.py +++ b/src/server/database/conversation/replay_rows.py @@ -1,31 +1,73 @@ """One-shot read of every table a thread replay needs.""" import logging +from dataclasses import dataclass +from typing import Any from psycopg.rows import dict_row from src.server.database import pool from src.server.database.conversation import _sql +from src.server.database.replay_facts import LEGACY_VERSION_SQL, RUN_FACTS_SQL from src.server.utils.pg_sanitize import normalize_uuid logger = logging.getLogger(__name__) -async def get_replay_thread_data( - thread_id: str, -) -> tuple[str | None, dict | None, list[dict], list[dict], list[dict], list[dict]]: - """Fetch all data needed for thread replay in a single connection. +@dataclass(frozen=True) +class ThreadRows: + """Every row a thread's replay reads. A turn's stored lines are keyed by + the versions of these rows, so a read checks and projects against one + set of them.""" + + # Canonical, as the rows store it. + thread_id: str + owner_id: str + # The thread row's summary; ``latest_checkpoint_id`` is the commit pointer. + thread: dict[str, Any] + queries: list[dict[str, Any]] + # One settled response per turn, without its stored events. + responses_by_turn: dict[Any, dict[str, Any]] + # Main-workflow usage rows. + usages: list[dict[str, Any]] + provenance: list[dict[str, Any]] + # The thread's runs that hold replay facts, each with its launching turn. + run_facts: list[dict[str, Any]] + + +async def get_branch_pointer(thread_id: str) -> tuple[str, list[Any]] | None: + """A thread's commit pointer and its persisted turn indexes ascending: + what pairing its branch's turns takes, without the rows a replay reads. + None when the thread is unknown or never committed a turn.""" + thread_id = normalize_uuid(thread_id) + if thread_id is None: + return None + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + """SELECT latest_checkpoint_id, + ARRAY(SELECT DISTINCT turn_index + FROM conversation_queries q + WHERE q.conversation_thread_id = t.conversation_thread_id + AND turn_index IS NOT NULL + ORDER BY turn_index) AS turn_indexes + FROM conversation_threads t + WHERE conversation_thread_id = %s""", + (thread_id,), + ) + row = await cur.fetchone() + if row is None or not row["latest_checkpoint_id"]: + return None + return row["latest_checkpoint_id"], list(row["turn_indexes"]) - Consolidates the separate connection acquisitions (owner check, thread - summary, queries, responses, usage rows, provenance rows) into 1, reducing - pool pressure during concurrent replays. - Returns: - (owner_user_id, thread_summary, queries, responses, usages, provenance) - """ +async def get_replay_thread_data(thread_id: str) -> ThreadRows | None: + """Every row a thread's replay reads, on one connection, which keeps + concurrent replays from each holding several. None when the thread does + not exist or the id is not a UUID.""" thread_id = normalize_uuid(thread_id) if thread_id is None: - return None, None, [], [], [], [] + return None async with pool.get_db_connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: @@ -41,7 +83,7 @@ async def get_replay_thread_data( owner_id = owner_row["user_id"] if owner_row else None if owner_id is None: - return None, None, [], [], [], [] + return None # 2. Thread summary await cur.execute( @@ -52,10 +94,12 @@ async def get_replay_thread_data( (thread_id,), ) thread = await cur.fetchone() + if thread is None: + return None - # 3. Queries + # 3. Queries, with the row version replay keys stored lines by await cur.execute( - """SELECT * FROM conversation_queries + """SELECT *, xmin FROM conversation_queries WHERE conversation_thread_id = %s ORDER BY turn_index ASC""", (thread_id,), @@ -63,9 +107,11 @@ async def get_replay_thread_data( queries = [dict(r) for r in await cur.fetchall()] # 4. Responses — one settled row per turn; an in_progress slot or - # superseded attempt must never drive replay rendering. + # superseded attempt must never drive replay rendering. Without + # their stored events: a turn replay projects afresh re-reads its + # row whole (see get_replay_responses). await cur.execute( - _sql._SETTLED_ATTEMPTS, + _sql.settled_attempts(_sql._LIGHT_RESPONSE_COLUMNS), (thread_id,), ) responses = [dict(r) for r in await cur.fetchall()] @@ -100,11 +146,55 @@ async def get_replay_thread_data( ) provenance = [dict(r) for r in await cur.fetchall()] - return ( - owner_id, - dict(thread) if thread else None, - queries, - responses, - usages, - provenance, + # 7. Replay facts of the thread's runs, keyed by launching turn + await cur.execute(RUN_FACTS_SQL, (thread_id,)) + run_facts = [dict(r) for r in await cur.fetchall()] + + return ThreadRows( + thread_id=thread_id, + owner_id=owner_id, + thread=dict(thread), + queries=queries, + responses_by_turn={r.get("turn_index"): r for r in responses}, + usages=usages, + provenance=provenance, + run_facts=run_facts, ) + + +async def get_replay_responses( + response_ids: list[str], *, legacy_facts: int | None = None +) -> dict[str, dict]: + """Response rows by id, whole: ``sse_events``, ``replay_facts`` and + ``xmin`` included. + + One read per row keeps a row's stored events and its version (the + projection cache's key) from the same tuple, so a projection keyed by + ``xmin`` was built from exactly that version's events. + + ``legacy_facts``: a row backfilled with legacy facts of this version + comes without its stored events, which those facts stand in for. The + events are never detoasted for it. + """ + ids = [normalize_uuid(r) for r in response_ids] + ids = [r for r in ids if r] + if not ids: + return {} + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + f"""SELECT {_sql._LIGHT_RESPONSE_COLUMNS}, replay_facts, + CASE WHEN %(version)s::text IS NULL + OR ({LEGACY_VERSION_SQL}) + IS DISTINCT FROM %(version)s::text + THEN sse_events END AS sse_events, + xmin + FROM conversation_responses + WHERE conversation_response_id = ANY(%(ids)s::uuid[])""", + { + "ids": ids, + "version": None if legacy_facts is None else str(legacy_facts), + }, + ) + rows = await cur.fetchall() + return {str(r["conversation_response_id"]): dict(r) for r in rows} diff --git a/src/server/database/conversation/responses.py b/src/server/database/conversation/responses.py index b2f032eed..861e9c060 100644 --- a/src/server/database/conversation/responses.py +++ b/src/server/database/conversation/responses.py @@ -1,12 +1,14 @@ """conversation_responses: SSE-event patches, provenance sync, and per-thread readers.""" import logging +from collections.abc import Collection from typing import Optional, List, Dict, Any, Tuple from psycopg.rows import dict_row from src.server.database import pool from src.server.database.conversation import _sql +from src.server.database.replay_facts import DROP_LEGACY_SQL from src.server.utils.pg_sanitize import SafeJson, safe_jsonb_array logger = logging.getLogger(__name__) @@ -30,12 +32,14 @@ async def _sync_provenance_for_response( sse_events: Optional[Any], strict: bool = False, appended_only: bool = True, + lanes: Collection[str], ) -> None: """(Re)derive provenance_records from sse_events on the caller's connection. Imported lazily to avoid a circular import (provenance imports this module). Best-effort by default; ``strict=True`` re-raises so a transaction-bound caller (the finalize CAS) aborts instead of committing over a poisoned txn. + ``lanes`` scopes the rewrite to those agents' rows. """ # Most persists carry no provenance (a turn with no external data access, or # a non-provenance event drain). A write that only appends cannot have @@ -54,6 +58,7 @@ async def _sync_provenance_for_response( turn_index=turn_index, sse_events=sse_events, strict=strict, + lanes=lanes, ) @@ -76,19 +81,17 @@ def _event_agents(events: List[Dict[str, Any]]) -> frozenset[str]: # One statement, so the row lock it takes orders it against append_sse_event: # under READ COMMITTED a waiting UPDATE re-reads the newest row. The strip is # one jsonpath filter over the stored array, so the kept rows are never -# unnested or sorted, and the blob never leaves Postgres: only its provenance -# entries come back, read from the written value, for the sync that follows. -_REPLACE_AGENT_EVENTS_SQL = """ +# unnested or sorted, and the blob never leaves Postgres. +_REPLACE_AGENT_EVENTS_SQL = f""" UPDATE conversation_responses SET sse_events = jsonb_path_query_array( COALESCE(sse_events, '[]'::jsonb), '$[*] ? (!exists(@.data.agent ? (@ == $drop[*])))', jsonb_build_object('drop', %(drop)s::text[]) - ) || %(append)s::jsonb + ) || %(append)s::jsonb, + {DROP_LEGACY_SQL} WHERE conversation_response_id = %(id)s - RETURNING conversation_thread_id, turn_index, - jsonb_path_query_array(sse_events, '$[*] ? (@.event == "provenance")') - AS provenance + RETURNING conversation_thread_id, turn_index """ @@ -104,6 +107,7 @@ async def replace_agent_events( when the row is missing. """ appended = await safe_jsonb_array(events) + agents = _event_agents(events) async with pool.get_db_connection() as conn: async with conn.transaction(): async with conn.cursor(row_factory=dict_row) as cur: @@ -111,7 +115,7 @@ async def replace_agent_events( _REPLACE_AGENT_EVENTS_SQL, { "id": conversation_response_id, - "drop": sorted(_event_agents(events)), + "drop": sorted(agents), "append": appended, }, ) @@ -122,16 +126,18 @@ async def replace_agent_events( f"response_id={conversation_response_id}" ) return False - # Re-derive provenance_records inside the same transaction: this - # is the choke point for the background subagent drain, which - # bypasses the turn-finalize path. + # The batch is the whole record of every lane it names, so those + # lanes' provenance is rewritten from it in the same transaction, + # and a lane it empties loses its rows with them. Other lanes keep + # theirs: the main lane is written when the run settles. await _sync_provenance_for_response( conn, conversation_response_id=conversation_response_id, conversation_thread_id=str(urow["conversation_thread_id"]), turn_index=urow["turn_index"], - sse_events=urow["provenance"], + sse_events=events, appended_only=False, + lanes=agents, ) logger.info( f"[conversation_db] replace_agent_events response_id=" @@ -154,9 +160,10 @@ async def append_sse_event( (the read-modify-write it replaces is not). Returns True when a row was updated, False when the thread has no response row yet. """ - sql = """ + sql = f""" UPDATE conversation_responses - SET sse_events = COALESCE(sse_events, '[]'::jsonb) || %s::jsonb + SET sse_events = COALESCE(sse_events, '[]'::jsonb) || %s::jsonb, + {DROP_LEGACY_SQL} WHERE conversation_response_id = ( SELECT conversation_response_id FROM conversation_responses @@ -241,13 +248,14 @@ async def get_recent_responses_for_thread( Selects the newest ``limit`` settled turns (latest attempt each) via ``turn_index DESC`` (so a window keeps the latest turns, not the oldest) and reverses them to chronological order. ``limit=None`` returns every - turn. + turn. Without ``sse_events`` and ``replay_facts``, which grow with the + turn: ``get_replay_responses`` reads a row whole when needed. """ try: async with pool.get_db_connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: base = f""" - SELECT {_sql._RESPONSE_COLUMNS} + SELECT {_sql._LIGHT_RESPONSE_COLUMNS} FROM ({_sql._SETTLED_ATTEMPTS}) r ORDER BY turn_index DESC """ @@ -263,3 +271,4 @@ async def get_recent_responses_for_thread( except Exception as e: logger.error(f"Error getting recent responses for thread: {e}") raise + diff --git a/src/server/database/conversation/threads_write.py b/src/server/database/conversation/threads_write.py index 7902dc08c..754066e6b 100644 --- a/src/server/database/conversation/threads_write.py +++ b/src/server/database/conversation/threads_write.py @@ -373,11 +373,12 @@ async def advance_thread_checkpoint_id( from_checkpoint_id: Optional[str], to_checkpoint_id: str, ) -> bool: - """CAS tip advance for root-namespace ui appends after a turn finalized. + """CAS tip advance for a root-namespace write outside a turn: a ui append + after a turn finalized, a manual compaction or offload. Not the turn-lifecycle writer (``database/runs/lifecycle.py`` owns that): this only moves the tip when it still points at the exact checkpoint the - append built on, so a concurrent turn or branch switch keeps ownership. + write built on, so a concurrent turn or branch switch keeps ownership. """ try: sql = """ diff --git a/src/server/database/conversation/turn_slices.py b/src/server/database/conversation/turn_slices.py new file mode 100644 index 000000000..e10d7b6c9 --- /dev/null +++ b/src/server/database/conversation/turn_slices.py @@ -0,0 +1,286 @@ +"""Stored per-turn and per-subagent-run checkpoint slices (derived data). + +Writes take a batch of rows in one statement. A row whose parent is already +gone (a concurrent fork or edit deletes the response, a thread delete takes +its runs) fails the whole batch on its foreign key, so the batch is retried a +row at a time and only the orphan is dropped: there is nothing left to keep. +""" + +import logging +from typing import Any, NamedTuple + +from psycopg.errors import ForeignKeyViolation +from psycopg.rows import dict_row + +from src.server.database import pool +from src.server.utils.pg_sanitize import SafeJson + +logger = logging.getLogger(__name__) + + +class StoredSlice(NamedTuple): + """A slice as stored: the checkpoints it spans, the code that cut it, its + bytes (None from a read that leaves them out) and, for a turn, the lines + and run claims stored beside it.""" + + input_checkpoint_id: str + tail_checkpoint_id: str + slice_key: str + codec: str | None = None + data: bytes | None = None + lines_key: str | None = None + lines: str | None = None + claims: list[str | None] | None = None + + +class TurnRow(NamedTuple): + response_id: str + thread_id: str + slice: StoredSlice + + +class TurnUpdate(NamedTuple): + """Lines and run claims to set beside a stored slice; a None leaves its + column. The tail keeps what was projected from a superseded slice from + landing on its successor.""" + + response_id: str + tail_checkpoint_id: str + lines_key: str | None + lines: str | None + claims: list[str | None] | None + + +class RunRow(NamedTuple): + thread_id: str + task_id: str + task_run_id: str | None + slice: StoredSlice + + +def _stored(row: dict[str, Any]) -> StoredSlice: + data = row.get("slice") + return StoredSlice( + input_checkpoint_id=row["input_checkpoint_id"], + tail_checkpoint_id=row["tail_checkpoint_id"], + slice_key=row["slice_key"], + codec=row.get("slice_codec"), + data=bytes(data) if data is not None else None, + lines_key=row.get("lines_key"), + lines=row.get("lines"), + claims=row.get("run_claims"), + ) + + +def _json(claims: list[str | None] | None) -> SafeJson | None: + return SafeJson(claims) if claims is not None else None + + +_VALIDITY_COLUMNS = ( + "conversation_response_id::text AS response_id, input_checkpoint_id, " + "tail_checkpoint_id, slice_key, lines_key, lines, run_claims" +) + + +async def _read_turns(response_ids: list[str], columns: str) -> dict[str, StoredSlice]: + if not response_ids: + return {} + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + f"SELECT {columns} FROM turn_slices " + "WHERE conversation_response_id = ANY(%s::uuid[])", + (response_ids,), + ) + return {r["response_id"]: _stored(r) for r in await cur.fetchall()} + + +async def get_turn_lines(response_ids: list[str]) -> dict[str, StoredSlice]: + """Turns' stored slices by response id, never their bytes.""" + return await _read_turns(response_ids, _VALIDITY_COLUMNS) + + +async def get_turn_slices(response_ids: list[str]) -> dict[str, StoredSlice]: + """Turns' stored slices by response id, bytes included.""" + return await _read_turns(response_ids, f"{_VALIDITY_COLUMNS}, slice_codec, slice") + + +async def has_turn_slices(thread_id: str, slice_key: str) -> bool: + """Whether a thread holds any turn slice under ``slice_key``: one indexed + probe that settles a thread not yet backfilled.""" + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + "SELECT EXISTS (SELECT 1 FROM turn_slices " + "WHERE conversation_thread_id = %s::uuid AND slice_key = %s) AS held", + (thread_id, slice_key), + ) + row = await cur.fetchone() + return bool(row and row["held"]) + + +async def get_slices_at( + thread_id: str, slice_key: str, input_checkpoint_ids: list[str] +) -> list[StoredSlice]: + """A thread's turn slices under ``slice_key`` that start at these + checkpoints, bytes included: what a reader holding turn anchors, and no + response ids, reads.""" + if not input_checkpoint_ids: + return [] + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + "SELECT input_checkpoint_id, tail_checkpoint_id, slice_key, " + "slice_codec, slice FROM turn_slices " + "WHERE conversation_thread_id = %s::uuid " + "AND slice_key = %s AND input_checkpoint_id = ANY(%s)", + (thread_id, slice_key, input_checkpoint_ids), + ) + return [_stored(r) for r in await cur.fetchall()] + + +async def _write_many(sql: str, params: list[tuple[Any, ...]], what: str) -> None: + async with pool.get_db_connection() as conn: + async with conn.cursor() as cur: + try: + await cur.executemany(sql, params) + return + except ForeignKeyViolation: + if len(params) == 1: + logger.debug("%s not stored, its parent row is gone", what) + return + for row in params: + try: + await cur.execute(sql, row) + except ForeignKeyViolation: + logger.debug("%s not stored, its parent row is gone", what) + + +# Claims are written as given, never kept from the row: they were placed +# against the slice beside them, and a reader trusts them only while that +# slice is still the turn's. +_UPSERT_TURN = """ + INSERT INTO turn_slices + (conversation_response_id, conversation_thread_id, input_checkpoint_id, + tail_checkpoint_id, slice_key, slice_codec, slice, lines_key, lines, + run_claims) + VALUES (%s::uuid, %s::uuid, %s, %s, %s, %s, %s, %s, %s, %s) + ON CONFLICT (conversation_response_id) DO UPDATE SET + conversation_thread_id = EXCLUDED.conversation_thread_id, + input_checkpoint_id = EXCLUDED.input_checkpoint_id, + tail_checkpoint_id = EXCLUDED.tail_checkpoint_id, + slice_key = EXCLUDED.slice_key, + slice_codec = EXCLUDED.slice_codec, + slice = EXCLUDED.slice, + lines_key = EXCLUDED.lines_key, + lines = EXCLUDED.lines, + run_claims = EXCLUDED.run_claims, + built_at = NOW() +""" + + +async def upsert_turn_slices(rows: list[TurnRow]) -> None: + """Store turns' slices, each with its lines and run claims when known.""" + if not rows: + return + await _write_many( + _UPSERT_TURN, + [ + ( + r.response_id, + r.thread_id, + r.slice.input_checkpoint_id, + r.slice.tail_checkpoint_id, + r.slice.slice_key, + r.slice.codec, + r.slice.data, + r.slice.lines_key, + r.slice.lines, + _json(r.slice.claims), + ) + for r in rows + ], + "turn slice", + ) + + +async def update_turns(rows: list[TurnUpdate]) -> None: + """Set lines and run claims on stored slices.""" + if not rows: + return + async with pool.get_db_connection() as conn: + async with conn.cursor() as cur: + await cur.executemany( + """UPDATE turn_slices + SET lines_key = COALESCE(%s, lines_key), + lines = COALESCE(%s, lines), + run_claims = COALESCE(%s::jsonb, run_claims), + built_at = NOW() + WHERE conversation_response_id = %s::uuid + AND tail_checkpoint_id = %s""", + [ + ( + r.lines_key, + r.lines, + _json(r.claims), + r.response_id, + r.tail_checkpoint_id, + ) + for r in rows + ], + ) + + +async def get_run_slices(thread_id: str, task_ids: list[str]) -> list[RunRow]: + """Every stored run slice of the given tasks.""" + if not task_ids: + return [] + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + """SELECT task_id, input_checkpoint_id, + task_run_id::text AS task_run_id, tail_checkpoint_id, + slice_key, slice_codec, slice + FROM task_run_slices + WHERE conversation_thread_id = %s::uuid + AND task_id = ANY(%s)""", + (thread_id, task_ids), + ) + return [ + RunRow(thread_id, r["task_id"], r["task_run_id"], _stored(r)) + for r in await cur.fetchall() + ] + + +async def upsert_run_slices(rows: list[RunRow]) -> None: + """Store subagent runs' slices.""" + if not rows: + return + await _write_many( + """INSERT INTO task_run_slices + (conversation_thread_id, task_id, input_checkpoint_id, + task_run_id, tail_checkpoint_id, slice_key, slice_codec, slice) + VALUES (%s::uuid, %s, %s, %s::uuid, %s, %s, %s, %s) + ON CONFLICT (conversation_thread_id, task_id, input_checkpoint_id) + DO UPDATE SET + task_run_id = EXCLUDED.task_run_id, + tail_checkpoint_id = EXCLUDED.tail_checkpoint_id, + slice_key = EXCLUDED.slice_key, + slice_codec = EXCLUDED.slice_codec, + slice = EXCLUDED.slice, + built_at = NOW()""", + [ + ( + r.thread_id, + r.task_id, + r.slice.input_checkpoint_id, + r.task_run_id, + r.slice.tail_checkpoint_id, + r.slice.slice_key, + r.slice.codec, + r.slice.data, + ) + for r in rows + ], + "run slice", + ) diff --git a/src/server/database/pool.py b/src/server/database/pool.py index f776578db..6fc971f26 100644 --- a/src/server/database/pool.py +++ b/src/server/database/pool.py @@ -10,7 +10,7 @@ import anyio from psycopg import AsyncConnection -from psycopg_pool import AsyncConnectionPool +from psycopg_pool import AsyncConnectionPool, PoolTimeout from src.config.env import DB_SSLMODE from src.config.settings import get_conversation_pool_max @@ -103,6 +103,20 @@ async def _configure_postgres_connection(conn): ) +class AppDataPoolTimeout(PoolTimeout): + """No app-data connection came free in time. Told apart from the + checkpointer pool's timeout, which leaves a reader that falls back to + stored rows a pool to read them through.""" + + +class _AppDataPool(AsyncConnectionPool): + async def getconn(self, timeout: float | None = None) -> AsyncConnection: + try: + return await super().getconn(timeout) + except PoolTimeout as e: + raise AppDataPoolTimeout(*e.args) from e + + def get_or_create_pool() -> AsyncConnectionPool: """ Get or create the shared connection pool for conversation database operations. @@ -121,7 +135,7 @@ def get_or_create_pool() -> AsyncConnectionPool: f"Creating PostgreSQL connection pool for conversations (max_size={pool_max})" ) # Create pool with minimal configuration matching LangGraph pool - _conversation_db_pool_cache[db_uri] = AsyncConnectionPool( + _conversation_db_pool_cache[db_uri] = _AppDataPool( conninfo=db_uri, min_size=1, max_size=pool_max, diff --git a/src/server/database/provenance.py b/src/server/database/provenance.py index 28c0cd9d4..f66f51b2f 100644 --- a/src/server/database/provenance.py +++ b/src/server/database/provenance.py @@ -1,14 +1,16 @@ """Provenance records: a derived index of external data the agent accessed. -Rows are extracted from the accumulated `conversation_responses.sse_events` -(top-level `event == "provenance"` entries) and written delete-then-insert -keyed by conversation_response_id, so the same turn can be re-persisted -(background subagent drains overwrite sse_events repeatedly) without -duplicating rows. All binds pass through the shared NUL sanitizers. +Rows are extracted from the ``provenance`` events an agent lane emitted and +written delete-then-insert per lane of a response: the main lane when the run +settles, a subagent's lane each time its archive lands. A lane is rewritten +whole by the one writer that holds all of it, so re-persisting it (subagent +drains repeat) never duplicates rows, and never touches another lane's. +All binds pass through the shared NUL sanitizers. """ import logging import math +from collections.abc import Collection from datetime import datetime from typing import Any @@ -33,6 +35,10 @@ "provider", ) +# The main agent's records, pinned to this name at emit time whatever graph +# node produced them; a subagent's are ``task:{id}``. +MAIN_PROVENANCE_LANE = "main" + # Signed 64-bit range of the result_size BIGINT column. _BIGINT_MIN = -(2**63) _BIGINT_MAX = 2**63 - 1 @@ -221,6 +227,12 @@ def _row_binds( ) +def record_lane(record: dict[str, Any]) -> str: + """The lane a record belongs to: its agent, ``""`` when it has none.""" + agent = record.get("agent") + return agent if isinstance(agent, str) else "" + + def provenance_row_to_event(row: dict[str, Any]) -> dict[str, Any]: """Wire ``provenance`` event data from one table row (inverse of extraction). @@ -255,6 +267,7 @@ async def insert_provenance_records( conversation_thread_id: str, turn_index: int, records: list[dict[str, Any]], + lanes: Collection[str], ) -> int: """Delete-then-insert provenance rows for one response (idempotent, atomic). @@ -271,6 +284,10 @@ async def insert_provenance_records( than executemany, which re-parses per row because finalize SQL is never prepared. Every TEXT bind is NUL-stripped and the JSONB ``args_fingerprint`` / ``args`` binds are wrapped in ``SafeJson``. Returns the number of rows inserted. + + ``lanes`` limits the delete to those agents' rows (an agentless row counts + as lane ``""``), and every record's own agent joins it, so a write always + replaces what it inserts. """ async with conn.cursor() as cur, conn.transaction(): # Serialize concurrent writers for this response (xact-scoped lock). @@ -278,11 +295,15 @@ async def insert_provenance_records( "SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))", (conversation_response_id,), ) - # Delete-then-insert keyed by response_id: re-runs (background drains) - # replace the prior set rather than accumulating duplicates. + # Delete-then-insert keyed by response_id and lane: re-runs + # (background drains) replace the prior set rather than accumulating + # duplicates. + scope = set(lanes) | {record_lane(r) for r in records} await cur.execute( - "DELETE FROM provenance_records WHERE conversation_response_id = %s", - (conversation_response_id,), + "DELETE FROM provenance_records " + "WHERE conversation_response_id = %s " + "AND COALESCE(agent, '') = ANY(%s)", + (conversation_response_id, sorted(scope)), ) if not records: @@ -317,6 +338,7 @@ async def sync_provenance_for_response( turn_index: int, sse_events: list[dict[str, Any]] | None, strict: bool = False, + lanes: Collection[str], ) -> int: """Extract provenance from sse_events and (re)write rows for one response. @@ -324,7 +346,8 @@ async def sync_provenance_for_response( default so a provenance failure can't break a standalone persist; ``strict=True`` re-raises for transaction-bound callers, where a swallowed SQL error would poison the enclosing transaction and turn its - commit into a silent rollback. + commit into a silent rollback. ``lanes`` as in + ``insert_provenance_records``. """ try: records = extract_provenance_from_sse_events(sse_events) @@ -334,6 +357,7 @@ async def sync_provenance_for_response( conversation_thread_id=conversation_thread_id, turn_index=turn_index, records=records, + lanes=lanes, ) except Exception as e: if strict: diff --git a/src/server/database/replay_facts.py b/src/server/database/replay_facts.py new file mode 100644 index 000000000..9d3a4125e --- /dev/null +++ b/src/server/database/replay_facts.py @@ -0,0 +1,105 @@ +"""The ``replay_facts`` column: what a settled run shows on replay that its +checkpoint does not keep. + +Both ledgers carry it. A response row holds its main agent's reasoning +durations (written by the finalize CAS) and, once backfilled, the facts its +stored events resolve to under ``LEGACY_KEY``. A subagent run's row holds the +captured URLs of the images its lane referenced (``images``) and the steering +it handed back unread (``steering_returned``), both learned after its settle. +""" + +from typing import Any, Dict, List, Optional + +from psycopg.rows import dict_row + +from src.server.database import pool +from src.server.utils.pg_sanitize import SafeJson + +LEGACY_KEY = "legacy" + +# Legacy facts derived from a row's stored events stand in for them on +# replay, so a write to the events takes the facts away in the same +# statement: the row replays from its events until the backfill derives the +# facts again. +DROP_LEGACY_SQL = f"replay_facts = replay_facts - '{LEGACY_KEY}'" + +# The version a row's legacy facts were resolved at, as text. +LEGACY_VERSION_SQL = f"replay_facts #>> '{{{LEGACY_KEY},v}}'" + + +async def append_returned_steering( + task_run_id: str, returned: List[Dict[str, Any]] +) -> Optional[Dict[str, Any]]: + """Append steering a settled run handed back, returning its ``thread_id`` + and ``parent_run_id`` (None when no row matched). + + Not the settle CAS: the sweep that follows it reads the queue. The merge + is one statement, so an images write landing on the row at the same + time keeps its key. + """ + if not returned: + return None + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + """ + UPDATE subagent_runs + SET replay_facts = COALESCE(replay_facts, '{}'::jsonb) + || jsonb_build_object('steering_returned', + COALESCE(replay_facts -> 'steering_returned', '[]'::jsonb) + || %(returned)s::jsonb) + WHERE task_run_id = %(id)s + RETURNING thread_id, parent_run_id + """, + {"id": task_run_id, "returned": SafeJson(returned)}, + ) + row = await cur.fetchone() + return dict(row) if row else None + + +async def append_lane_images( + thread_id: str, response_id: str, task_id: str, images: Dict[str, str] +) -> bool: + """Merge a lane's captured images into every run the turn launched for + that task; False when it launched none. + + The archive's records are fenced to the turn, so they come from exactly + these runs, and replay reads a run's facts only through the turn that + launched it. ``thread_id`` is there for the ``(thread_id, task_id)`` + index, as ``parent_run_id`` has none of its own. + """ + if not images: + return False + async with pool.get_db_connection() as conn: + async with conn.cursor() as cur: + await cur.execute( + """ + UPDATE subagent_runs + SET replay_facts = COALESCE(replay_facts, '{}'::jsonb) + || jsonb_build_object('images', + COALESCE(replay_facts -> 'images', '{}'::jsonb) + || %(images)s::jsonb) + WHERE thread_id = %(thread)s AND task_id = %(task)s + AND parent_run_id = %(response)s + """, + { + "thread": thread_id, + "task": task_id, + "response": response_id, + "images": SafeJson(images), + }, + ) + return cur.rowcount > 0 + + +# Every run of a thread that holds replay facts, with the turn index of the +# response that launched it: replay keys a turn's stored lines on the facts of +# the runs it shows, so they are read with the thread's other replay rows. +RUN_FACTS_SQL = """ + SELECT r.turn_index, s.task_run_id, s.replay_facts + FROM subagent_runs s + JOIN conversation_responses r + ON r.conversation_response_id = s.parent_run_id + WHERE s.thread_id = %s AND s.replay_facts IS NOT NULL + ORDER BY s.task_run_id +""" diff --git a/src/server/database/runs/lifecycle.py b/src/server/database/runs/lifecycle.py index f316a8094..b27c63264 100644 --- a/src/server/database/runs/lifecycle.py +++ b/src/server/database/runs/lifecycle.py @@ -11,7 +11,7 @@ import logging from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any, Awaitable, Callable, Dict, List, Optional @@ -22,6 +22,7 @@ from src.server.database import order_attempts as oa_db from src.server.database import pool from src.server.database import thread_transcripts as tt_db +from src.server.database.runs import subagent_repair from src.server.database.runs import subagent_runs as sr_db from src.server.database.runs.outbox import ( build_finalize_jobs_from_run_row, @@ -29,6 +30,7 @@ release_deferred_jobs, ) from src.server.contracts.status import TERMINAL_STATUSES +from src.server.database.provenance import MAIN_PROVENANCE_LANE from src.server.utils.pg_sanitize import SafeJson, normalize_uuid logger = logging.getLogger(__name__) @@ -91,10 +93,50 @@ class ForkSpec: preserve_query_at_fork: bool = False +# The row a finalize hands back: the terminal status plus what the post-CAS +# effects read (outbox jobs, the transcript export). Never sse_events: a long +# turn's archive runs to tens of MB, and decoding it holds the event loop. +_FINALIZE_ROW_COLUMNS = """ + conversation_response_id, conversation_thread_id, turn_index, run_seq, + status, interrupt_reason, metadata +""" + + @dataclass class FinalizeResult: applied: bool # False = run was already terminal (idempotent no-op) - run: Optional[Dict[str, Any]] = None + run: Optional[Dict[str, Any]] = None # _FINALIZE_ROW_COLUMNS only + + +@dataclass(frozen=True) +class ProducerRecord: + """What a run's stream kept apart for its row, beside the status. + + ``provenance`` is every provenance record the stream held: the main + lane's, plus the task lanes a stop drained. ``reasoning_ms`` is None when + nothing timed the main agent's thinking (a recovered run), so the row's + facts stay unset and replay takes the durations from its stored events. + """ + + provenance: List[Dict[str, Any]] = field(default_factory=list) + reasoning_ms: Optional[Dict[str, List[int]]] = None + + +@dataclass +class RunOutcome: + """Everything finalize writes, resolved in-band by the run's own executor.""" + + status: str # completed | interrupted | error | cancelled + interrupt_reason: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + warnings: Optional[List[str]] = None + errors: Optional[List[str]] = None + execution_time: Optional[float] = None + sse_events: Optional[List[Dict[str, Any]]] = None + per_call_records: Optional[list] = None + tool_usage: Optional[Dict[str, int]] = None + # None for a run with no stream to keep apart (one that never started). + record: Optional[ProducerRecord] = None @asynccontextmanager @@ -238,7 +280,7 @@ async def start_run( # branch no longer owns; re-anchor the task rows in the # same transaction so no window exists where a resume can # read an empty latest_run_id. - await sr_db.repair_task_chains(thread_id, conn=conn) + await subagent_repair.repair_task_chains(thread_id, conn=conn) # update_thread_checkpoint_id swallows failures into # False; inside this transaction that must abort loudly, # not commit a truncation with an unpinned checkpoint. @@ -348,13 +390,7 @@ async def finalize_run( *, run_id: str, thread_id: str, - status: str, - interrupt_reason: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - warnings: Optional[List[str]] = None, - errors: Optional[List[str]] = None, - execution_time: Optional[float] = None, - sse_events: Optional[List[Dict[str, Any]]] = None, + outcome: RunOutcome, checkpoint_id: Optional[str] = None, usage_writer: Optional[Callable[[Any, str], Awaitable[None]]] = None, conn=None, @@ -377,9 +413,21 @@ async def finalize_run( flag is the user-provenance marker, the timestamp only the decision time. usage_writer runs inside the transaction on the same connection so a persist failure aborts the terminal transition with it (no more swallowed-exception zombie turns). + + The outcome's record lands with the status it describes, so a settled + row never lacks the facts its replay reads, and its provenance replaces + the main lane's rows. Subagent lanes are their archive's to write. """ + status = outcome.status if status not in TERMINAL_STATUSES: raise ValueError(f"finalize_run: {status!r} is not a terminal status") + record = outcome.record or ProducerRecord() + replay_facts = ( + {"reasoning_ms": record.reasoning_ms} + if record.reasoning_ms is not None + else None + ) + sse_events = outcome.sse_events try: async with _lifecycle_connection(conn) as conn: @@ -422,24 +470,28 @@ async def finalize_run( -- Merge, never replace: mid-run appenders -- (append_sse_event, e.g. manual compact/offload -- context_window persists) may have durably written - -- to the open row already. - sse_events = CASE - WHEN %s::jsonb IS NULL THEN sse_events - ELSE COALESCE(sse_events, '[]'::jsonb) || %s::jsonb - END + -- to the open row already. The events bind once, + -- since every bind encodes them on the event loop; + -- jsonb || NULL is NULL, so a finalize without + -- events leaves the column as it was. + sse_events = COALESCE( + COALESCE(sse_events, '[]'::jsonb) || %s::jsonb, + sse_events + ), + replay_facts = COALESCE(%s::jsonb, replay_facts) WHERE conversation_response_id = %s AND status = 'in_progress' - RETURNING * - """, + RETURNING """ + + _FINALIZE_ROW_COLUMNS, ( status, - interrupt_reason, + outcome.interrupt_reason, status, - SafeJson(metadata or {}), - warnings or [], - errors or [], - execution_time, - SafeJson(sse_events) if sse_events else None, + SafeJson(outcome.metadata or {}), + outcome.warnings or [], + outcome.errors or [], + outcome.execution_time, SafeJson(sse_events) if sse_events else None, + SafeJson(replay_facts) if replay_facts else None, run_id, ), ) @@ -475,11 +527,9 @@ async def finalize_run( conversation_response_id=run_id, conversation_thread_id=thread_id, turn_index=run_row["turn_index"], - # The RETURNING row carries the MERGED archive (pre-existing - # mid-run appends || this finalize's events) — provenance - # must derive from what was actually persisted. - sse_events=run_row.get("sse_events"), + sse_events=record.provenance, strict=True, + lanes=(MAIN_PROVENANCE_LANE,), ) if usage_writer is not None: @@ -527,7 +577,7 @@ async def finalize_run( ) return FinalizeResult(applied=True, run=run_row) except _AlreadyTerminal: - run = await get_run(run_id) + run = await _get_finalize_row(run_id) if run is None: raise RunNotFoundError(run_id) logger.info( @@ -537,6 +587,20 @@ async def finalize_run( return FinalizeResult(applied=False, run=run) +async def _get_finalize_row(run_id: str) -> Optional[Dict[str, Any]]: + """The survivor of a lost finalize, in the shape a won one returns.""" + async with pool.get_db_connection() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute( + "SELECT " + + _FINALIZE_ROW_COLUMNS + + " FROM conversation_responses WHERE conversation_response_id = %s", + (run_id,), + ) + row = await cur.fetchone() + return dict(row) if row else None + + async def request_run_cancel( run_id: str, thread_id: Optional[str] = None ) -> Dict[str, Any]: diff --git a/src/server/database/runs/subagent_repair.py b/src/server/database/runs/subagent_repair.py new file mode 100644 index 000000000..9d997779e --- /dev/null +++ b/src/server/database/runs/subagent_repair.py @@ -0,0 +1,105 @@ +"""Task-chain repair: re-anchor subagent_tasks rows to their surviving runs. + +Deleting conversation_responses cascades their subagent_runs away, which can +strand a task pointing at nothing (rewind to the newest survivor) or with +nothing left at all (delete the row). Left unrepaired, a later resume reads +the empty pointer and starts an unchained run with no predecessor and no +start pin. +""" + +import logging +from typing import Any, Dict + +from src.server.database import pool + +logger = logging.getLogger(__name__) + +# The lazy global sweep runs unfenced, so it ignores rows younger than this. +# A task row is born without runs for the width of the START transaction; the +# FK from subagent_runs already makes deleting a run-bearing task impossible, +# and this makes the sweep's zero-run read stale-proof as well. +SWEEP_MIN_AGE_SECONDS = 300 + +# A task's latest_run_id is dangling when nothing in subagent_runs answers to +# it. NULL counts: the FK is ON DELETE SET NULL, so a cascade-deleted run +# leaves the pointer empty rather than pointing at a tombstone. Both shapes +# read the same to every consumer — the task has no resolvable latest run. +_DANGLING_LATEST = """ + NOT EXISTS ( + SELECT 1 FROM subagent_runs r WHERE r.task_run_id = t.latest_run_id + ) +""" + +_HAS_SURVIVING_RUN = """ + EXISTS ( + SELECT 1 FROM subagent_runs r + WHERE r.thread_id = t.thread_id AND r.task_id = t.task_id + ) +""" + + +async def _rechain(conn, scope: str, param: Any) -> Dict[str, int]: + """Rewind the dangling tasks ``scope`` selects to their newest surviving + run, and delete those with none left.""" + async with conn.cursor() as cur: + await cur.execute( + f""" + UPDATE subagent_tasks t + SET latest_run_id = ( + SELECT r.task_run_id FROM subagent_runs r + WHERE r.thread_id = t.thread_id + AND r.task_id = t.task_id + ORDER BY r.started_at DESC + LIMIT 1 + ), + updated_at = NOW() + WHERE {scope} + AND {_DANGLING_LATEST} + AND {_HAS_SURVIVING_RUN} + """, + (param,), + ) + rewound = cur.rowcount + + await cur.execute( + f""" + DELETE FROM subagent_tasks t + WHERE {scope} + AND NOT {_HAS_SURVIVING_RUN} + """, + (param,), + ) + return {"rewound": rewound, "deleted": cur.rowcount} + + +async def repair_task_chains(thread_id: str, conn=None) -> Dict[str, int]: + """Re-anchor a thread's task rows to their surviving runs. + + Idempotent and safe on any thread — a healthy thread matches neither + statement — so truncation paths can call it unconditionally inside their + own transaction. + """ + async with pool.get_db_connection(conn) as conn: + result = await _rechain(conn, "t.thread_id = %s", thread_id) + if result["rewound"] or result["deleted"]: + logger.info( + f"[subagent_runs] REPAIR thread={thread_id} " + f"rewound={result['rewound']} deleted={result['deleted']}" + ) + return result + + +async def repair_dangling_task_chains( + min_age_seconds: int = SWEEP_MIN_AGE_SECONDS, +) -> Dict[str, int]: + """The global, thread-agnostic form of repair_task_chains. + + Heals damage that predates the transactional rewind — the shadow-deploy + window, and any truncation path that ever escapes the guard. Drives off + the dangling-pointer anti-join so a healthy ledger scans a small table and + updates nothing. + """ + async with pool.get_db_connection() as conn: + return await _rechain( + conn, "t.updated_at < NOW() - MAKE_INTERVAL(secs => %s)", min_age_seconds + ) diff --git a/src/server/database/runs/subagent_runs.py b/src/server/database/runs/subagent_runs.py index 23d21c742..d7133ab42 100644 --- a/src/server/database/runs/subagent_runs.py +++ b/src/server/database/runs/subagent_runs.py @@ -745,126 +745,6 @@ async def get_latest_run_details( } -# ------------------------------------------------------------------ repair - -# The lazy global sweep runs unfenced, so it ignores rows younger than this. -# A task row is born without runs for the width of the START transaction; the -# FK from subagent_runs already makes deleting a run-bearing task impossible, -# and this makes the sweep's zero-run read stale-proof as well. -SWEEP_MIN_AGE_SECONDS = 300 - -# A task's latest_run_id is dangling when nothing in subagent_runs answers to -# it. NULL counts: the FK is ON DELETE SET NULL, so a cascade-deleted run -# leaves the pointer empty rather than pointing at a tombstone. Both shapes -# read the same to every consumer — the task has no resolvable latest run. -_DANGLING_LATEST = """ - NOT EXISTS ( - SELECT 1 FROM subagent_runs r WHERE r.task_run_id = t.latest_run_id - ) -""" - -_HAS_SURVIVING_RUN = """ - EXISTS ( - SELECT 1 FROM subagent_runs r - WHERE r.thread_id = t.thread_id AND r.task_id = t.task_id - ) -""" - - -async def repair_task_chains(thread_id: str, conn=None) -> Dict[str, int]: - """Re-anchor a thread's task rows to their surviving runs. - - Deleting conversation_responses cascades their subagent_runs away, which - can strand a task pointing at nothing (rewind to the newest survivor) or - with nothing left at all (delete the row). Left unrepaired, a later resume - reads the empty pointer and starts an unchained run with no predecessor - and no start pin. Idempotent and safe on any thread — a healthy thread - matches neither statement — so truncation paths can call it unconditionally - inside their own transaction. - """ - async with _ledger_connection(conn) as conn: - async with conn.cursor() as cur: - await cur.execute( - f""" - UPDATE subagent_tasks t - SET latest_run_id = ( - SELECT r.task_run_id FROM subagent_runs r - WHERE r.thread_id = t.thread_id - AND r.task_id = t.task_id - ORDER BY r.started_at DESC - LIMIT 1 - ), - updated_at = NOW() - WHERE t.thread_id = %s - AND {_DANGLING_LATEST} - AND {_HAS_SURVIVING_RUN} - """, - (thread_id,), - ) - rewound = cur.rowcount - - await cur.execute( - f""" - DELETE FROM subagent_tasks t - WHERE t.thread_id = %s - AND NOT {_HAS_SURVIVING_RUN} - """, - (thread_id,), - ) - deleted = cur.rowcount - - if rewound or deleted: - logger.info( - f"[subagent_runs] REPAIR thread={thread_id} rewound={rewound} " - f"deleted={deleted}" - ) - return {"rewound": rewound, "deleted": deleted} - - -async def repair_dangling_task_chains( - min_age_seconds: int = SWEEP_MIN_AGE_SECONDS, -) -> Dict[str, int]: - """The global, thread-agnostic form of repair_task_chains. - - Heals damage that predates the transactional rewind — the shadow-deploy - window, and any truncation path that ever escapes the guard. Drives off - the dangling-pointer anti-join so a healthy ledger scans a small table and - updates nothing. - """ - async with pool.get_db_connection() as conn: - async with conn.cursor() as cur: - await cur.execute( - f""" - UPDATE subagent_tasks t - SET latest_run_id = ( - SELECT r.task_run_id FROM subagent_runs r - WHERE r.thread_id = t.thread_id - AND r.task_id = t.task_id - ORDER BY r.started_at DESC - LIMIT 1 - ), - updated_at = NOW() - WHERE t.updated_at < NOW() - MAKE_INTERVAL(secs => %s) - AND {_DANGLING_LATEST} - AND {_HAS_SURVIVING_RUN} - """, - (min_age_seconds,), - ) - rewound = cur.rowcount - - await cur.execute( - f""" - DELETE FROM subagent_tasks t - WHERE t.updated_at < NOW() - MAKE_INTERVAL(secs => %s) - AND NOT {_HAS_SURVIVING_RUN} - """, - (min_age_seconds,), - ) - deleted = cur.rowcount - - return {"rewound": rewound, "deleted": deleted} - - # ------------------------------------------------------------ mutation guard diff --git a/src/server/handlers/chat/ptc_run.py b/src/server/handlers/chat/ptc_run.py index 83e425d33..af6123cc6 100644 --- a/src/server/handlers/chat/ptc_run.py +++ b/src/server/handlers/chat/ptc_run.py @@ -28,6 +28,7 @@ from src.server.database.user import get_user_profile_for_prompt from src.server.database.workspace import update_workspace_activity from src.server.services.computer_disk import TURN_MEASURE_MIN_INTERVAL_SECONDS +from src.server.services.history.window import runs_held from src.server.services.runs.sse_producer import RunSSEProducer from src.server.models.chat import ( ChatRequest, @@ -664,6 +665,7 @@ async def cache_workspace_preview( # A port, not a value: the switch can flip mid-turn from any # worker, so the agent re-reads the row on every call. subagent_switch=partial(read_thread_subagents_allowed, thread_id), + window_coverage=partial(runs_held, thread_id), ) _mark_phase("graph_build") diff --git a/src/server/handlers/checkpoint_handler.py b/src/server/handlers/checkpoint_handler.py index ce0688ea8..a51967121 100644 --- a/src/server/handlers/checkpoint_handler.py +++ b/src/server/handlers/checkpoint_handler.py @@ -8,13 +8,12 @@ """ import logging -from bisect import bisect_right -from typing import Any from fastapi import HTTPException from src.server.database.conversation import queries as queries_db from src.server.database.conversation import threads_read +from src.server.services.history.reader import pair_turns, turn_anchors from src.server.utils.checkpoint_helpers import ( build_checkpoint_config, get_checkpointer, @@ -28,29 +27,6 @@ logger = logging.getLogger(__name__) -def _turn_numbers(boundaries: list[Any], persisted: list[int]) -> list[int | None]: - """Each boundary's persisted turn_index. - - A turn's input checkpoint carries the turn_index its run was admitted - under. A resume's checkpoint belongs to the run it answers, and threads - from before the stamp carry none, so those take the next persisted turn - after the one before: wherever replay's pairing holds, this is the turn it - names. None when no persisted turn is left to name. - """ - numbers: list[int | None] = [] - previous: int | None = None - for cp_tuple in boundaries: - metadata = cp_tuple.metadata or {} - number = metadata.get("turn_index") if metadata.get("source") == "input" else None - if number is None: - at = 0 if previous is None else bisect_right(persisted, previous) - number = persisted[at] if at < len(persisted) else None - numbers.append(number) - if number is not None: - previous = number - return numbers - - async def get_thread_turns( thread_id: str, branch_tip_checkpoint_id: str | None = None ) -> ThreadTurnsResponse: @@ -58,11 +34,9 @@ async def get_thread_turns( Edit/regenerate fork the checkpoint graph, so only ancestors of the branch tip count (tip: ``branch_tip_checkpoint_id`` when on the graph, else the - newest checkpoint). Each turn is named by its persisted ``turn_index``, - the number its rows and replayed bubbles carry, never by its position - among the boundaries: a turn whose run died before its first checkpoint - has rows and no boundary, so every later turn sits a place lower than its - number. A boundary with no persisted turn to name is left out. + newest checkpoint). Each turn is named by its persisted ``turn_index`` + (``pair_turns``), the number its rows and replayed bubbles carry; a turn + with no persisted turn to name is left out. """ try: boundaries, tip_id = await walk_current_branch_boundaries( @@ -75,23 +49,16 @@ async def get_thread_turns( logger.error(f"[CHECKPOINT] Failed to list checkpoints for thread {thread_id}: {e}") raise HTTPException(status_code=500, detail="Failed to retrieve checkpoint history") - turns = [] - for cp_tuple, turn_index in zip(boundaries, _turn_numbers(boundaries, persisted)): - if turn_index is None: - continue - cp_id = cp_tuple.config["configurable"]["checkpoint_id"] - # The parent checkpoint is the state BEFORE this turn — only meaningful - # for a source=input turn (a resume shares its interrupted turn's state). - is_source_input = (cp_tuple.metadata or {}).get("source") == "input" - edit_checkpoint_id = None - if is_source_input and cp_tuple.parent_config: - edit_checkpoint_id = cp_tuple.parent_config["configurable"].get("checkpoint_id") - - turns.append(TurnCheckpointInfo( + anchors = turn_anchors(boundaries, tip_id) if tip_id else [] + turns = [ + TurnCheckpointInfo( turn_index=turn_index, - edit_checkpoint_id=edit_checkpoint_id, - regenerate_checkpoint_id=cp_id, - )) + edit_checkpoint_id=None if anchor.is_resume else anchor.parent_checkpoint_id, + regenerate_checkpoint_id=anchor.input_checkpoint_id, + ) + for anchor, turn_index in zip(anchors, pair_turns(anchors, persisted)) + if turn_index is not None + ] return ThreadTurnsResponse( thread_id=thread_id, diff --git a/src/server/handlers/thread_maintenance.py b/src/server/handlers/thread_maintenance.py index 1e9959f99..5bd280985 100644 --- a/src/server/handlers/thread_maintenance.py +++ b/src/server/handlers/thread_maintenance.py @@ -16,6 +16,7 @@ from src.server.utils.checkpoint_helpers import ( build_checkpoint_config, get_checkpointer, + update_at_commit, ) # Import setup module to access initialized globals @@ -43,7 +44,7 @@ async def _resolve_graph_and_state( folder in place until the caller's block ends. Returns: - (graph, lg_config, state, messages, workspace_id, backend) + (graph, state, messages, workspace_id, backend) """ from src.server.database import conversation as qr_db from src.server.database.workspace import get_workspace @@ -101,9 +102,10 @@ async def _resolve_graph_and_state( ) # State with timeout - lg_config = build_checkpoint_config(thread_id) try: - state = await asyncio.wait_for(graph.aget_state(lg_config), timeout=10.0) + state = await asyncio.wait_for( + graph.aget_state(build_checkpoint_config(thread_id)), timeout=10.0 + ) except asyncio.TimeoutError: logger.error(f"aget_state timed out for thread {thread_id} during {verb}") raise HTTPException( @@ -140,21 +142,36 @@ async def _resolve_graph_and_state( layout = SandboxLayout(session.sandbox.working_dir).for_workspace(dir_name) backend = SandboxBackend(session.sandbox, layout.workspace) - return graph, lg_config, state, messages, workspace_id, backend + return graph, state, messages, workspace_id, backend async def _update_graph_state( - graph, config: dict, values: dict, thread_id: str, verb: str + graph, state, values: dict, thread_id: str, verb: str ) -> None: - """Timeout-wrapped aupdate_state call.""" + """Write ``values`` onto the checkpoint ``state`` was read at, and move + the thread's commit pointer onto the write when it still points there + (``update_at_commit``). + + Left beside the pointer, the write is seen by the next turn, which runs + from the latest checkpoint, but not by anything reading the committed + branch until that turn ends: replay, and the window's slice coverage. + """ + built_on = ((state.config or {}).get("configurable") or {}).get("checkpoint_id") try: - await asyncio.wait_for(graph.aupdate_state(config, values), timeout=10.0) + advanced = await update_at_commit( + graph, thread_id, built_on, values, timeout=10.0 + ) except asyncio.TimeoutError: logger.error(f"aupdate_state timed out for thread {thread_id} during {verb}") raise HTTPException( status_code=504, detail=f"Timed out updating state for thread: {thread_id}", ) + if not advanced: + logger.info( + f"[{verb}] thread {thread_id}: commit pointer is not at {built_on}, " + "left in place" + ) @asynccontextmanager @@ -253,7 +270,7 @@ async def resolve(model: str | None): agent_cfg = setup.agent_config ( - graph, lg_config, state, messages, workspace_id, backend + graph, state, messages, workspace_id, backend ) = await _resolve_graph_and_state( thread_id, "compact", config=agent_cfg, checkpointer=mutation.saver, user_id=user_id, held=held, @@ -282,7 +299,11 @@ async def resolve(model: str | None): raise HTTPException(status_code=400, detail=str(e)) await _update_graph_state( - graph, lg_config, compaction.update(state.values), thread_id, "compact" + graph, + state, + compaction.update(state.values), + thread_id, + "compact", ) # The view that was compacted, as automatic compaction counts it, @@ -370,7 +391,7 @@ async def trigger_offload(thread_id: str, *, user_id: str | None = None) -> dict AsyncExitStack() as held, ): ( - graph, lg_config, state, messages, workspace_id, backend + graph, state, messages, workspace_id, backend ) = await _resolve_graph_and_state( thread_id, "offload", checkpointer=mutation.saver, user_id=user_id, held=held ) @@ -398,7 +419,7 @@ async def trigger_offload(thread_id: str, *, user_id: str | None = None) -> dict # user asked for it now. await _update_graph_state( graph, - lg_config, + state, record_offloads(state.values, offloaded_arg_ids, offloaded_read_ids), thread_id, "offload", diff --git a/src/server/services/automation_excerpt.py b/src/server/services/automation_excerpt.py index 04881e48e..0ceefc41a 100644 --- a/src/server/services/automation_excerpt.py +++ b/src/server/services/automation_excerpt.py @@ -81,18 +81,17 @@ async def read_run_excerpt(thread_id: str, run_id: str) -> Optional[str]: from src.server.services.history.reader import CheckpointHistoryReader try: - history = await CheckpointHistoryReader.get_instance().aget_recent_history( - thread_id, 1 - ) + reader = CheckpointHistoryReader.get_instance() + anchors, _ = await reader.aget_turn_anchors(thread_id) + if not anchors or anchors[-1].run_id != run_id: + return None + (turn,) = await reader.aget_turn_slices(thread_id, anchors[-1:]) except Exception as e: logger.warning( f"[AUTOMATION] Excerpt read failed: thread_id={thread_id} " f"run_id={run_id}: {e}" ) return None - turn = history.turns[-1] if history.turns else None - if turn is None or turn.run_id != run_id: - return None for message in reversed(turn.messages): if not isinstance(message, AIMessage): continue diff --git a/src/server/services/automation_settlement.py b/src/server/services/automation_settlement.py index db4717da3..a97c00277 100644 --- a/src/server/services/automation_settlement.py +++ b/src/server/services/automation_settlement.py @@ -224,11 +224,14 @@ def run_failure_message(run: Dict[str, Any]) -> Optional[str]: error, else a line naming it. A credit pause records no error. Its words are the quota service's - denial, which the pause's interrupt carries on the row, and they are - relayed verbatim or not at all. + denial, which every finalize stamps into the row's metadata, and they + are relayed verbatim or not at all. """ if run.get("interrupt_reason") == INTERRUPT_REASON_CREDIT_PAUSE: - return _pause_message(run.get("sse_events")) + stamped = (run.get("metadata") or {}).get("credit_pause_message") + if isinstance(stamped, str): + return stamped + return stored_pause_message(run.get("sse_events")) errors = run.get("errors") if errors: return clean_error_text(str(errors[-1])) @@ -238,7 +241,14 @@ def run_failure_message(run: Dict[str, Any]) -> Optional[str]: ) -def _pause_message(events: Any) -> Optional[str]: +def stored_pause_message(events: Any) -> Optional[str]: + """A credit pause's denial, read off its stored ``interrupt`` event. + + For rows a release before the stamp finalized: turns an old worker + drains across a deploy, and older rows the replay-facts backfill has not + reached yet. Both still settle through here, so this goes with the + stored events column, after the backfill's last run. + """ for event in reversed(events if isinstance(events, list) else []): if not isinstance(event, dict) or event.get("event") != "interrupt": continue diff --git a/src/server/services/history/__init__.py b/src/server/services/history/__init__.py index 2435e0489..5d73cec2f 100644 --- a/src/server/services/history/__init__.py +++ b/src/server/services/history/__init__.py @@ -1,9 +1,6 @@ """Checkpoint-sourced thread history: reader (I/O) + projector (pure).""" -from src.server.services.history.reader import ( - CheckpointHistoryReader, - ThreadHistory, - TurnSlice, -) +from src.server.services.history.reader import CheckpointHistoryReader +from src.server.services.history.slices import TurnSlice -__all__ = ["CheckpointHistoryReader", "ThreadHistory", "TurnSlice"] +__all__ = ["CheckpointHistoryReader", "TurnSlice"] diff --git a/src/server/services/history/committed.py b/src/server/services/history/committed.py new file mode 100644 index 000000000..1e9244c12 --- /dev/null +++ b/src/server/services/history/committed.py @@ -0,0 +1,114 @@ +"""What a settled turn's main agent said, for a tool reading another thread. + +The secretary's ``agent_output`` reads a thread's newest turns while its own +turn waits. A turn's stored slice serves it while the slice is the turn's on +the current branch, by the rule replay applies (``slices.stored_spans``); +only a turn without one waits for its projection. Replay's stored lines would +serve the same text, but checking them reads every row of the thread, which +costs this path several times what the slice read does. +""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from typing import Any + +from src.server.database.conversation import turn_slices as slices_db +from src.server.database.conversation.replay_rows import get_branch_pointer +from src.server.services.history import slices +from src.server.services.history.projector import ( + history_events_to_sse, + messages_to_history_events, +) +from src.server.services.history.reader import CheckpointHistoryReader, TurnAnchor +from src.server.services.history.replay.errors import CheckpointReplayUnavailable +from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound + +logger = logging.getLogger(__name__) + + +async def committed_texts(thread_id: str, turns: Mapping[str, Any]) -> dict[str, str]: + """The main agent's text in each of ``turns`` (response id to turn + index), cut from the turn's stored slice, which a turn without one has + projected first, as its settle would have. A response left out has no + slice to read (no checkpoint boundary, a turn waiting on the thread's + claims pass, a failed read), so its stored events are the caller's to + read.""" + from src.server.services.history.replay.refresh import refresh_thread + + anchors = await _anchors(thread_id, turns) + texts = await _slice_texts(thread_id, anchors) + missing = {rid for rid in anchors if rid not in texts} + if missing: + await refresh_thread(thread_id, missing, claims_pass=False) + texts.update( + await _slice_texts( + thread_id, {rid: anchors[rid] for rid in missing} + ) + ) + return texts + + +async def _anchors( + thread_id: str, turns: Mapping[str, Any] +) -> dict[str, TurnAnchor]: + """Each response's turn on the current branch, paired as replay pairs + them; none when the branch cannot be read or paired.""" + from src.server.services.history.replay import pair_anchors + + try: + pointer = await get_branch_pointer(thread_id) + if pointer is None: + return {} + tip, persisted = pointer + anchors, _ = await CheckpointHistoryReader.get_instance().aget_turn_anchors( + thread_id, tip + ) + paired = dict(pair_anchors(anchors, persisted)) + except (CheckpointBranchTipNotFound, CheckpointReplayUnavailable) as e: + logger.debug("no branch to read for %s: %s", thread_id, e) + return {} + except Exception: + logger.warning( + "branch read failed for %s; using stored events", thread_id, exc_info=True + ) + return {} + return {rid: paired[ti] for rid, ti in turns.items() if ti in paired} + + +async def _slice_texts( + thread_id: str, anchors: Mapping[str, TurnAnchor] +) -> dict[str, str]: + """Main-lane text by response id, for the turns with a usable slice.""" + if not anchors: + return {} + try: + serde = CheckpointHistoryReader.get_instance().serde + rows = await slices_db.get_turn_slices(list(anchors)) + except Exception: + logger.warning("turn slice read failed; using stored events", exc_info=True) + return {} + out: dict[str, str] = {} + for response_id, span in slices.stored_spans(serde, anchors, rows).items(): + try: + items = history_events_to_sse( + messages_to_history_events(span.messages), thread_id=thread_id + ) + except Exception: + logger.warning( + "turn slice projection failed for response %s; using stored events", + response_id, + exc_info=True, + ) + continue + out[response_id] = "".join( + d["content"] + for i in items + if i.get("event") == "message_chunk" + and isinstance(d := i.get("data"), dict) + and d.get("content_type") == "text" + and isinstance(d.get("content"), str) + and d["content"] + ) + return out diff --git a/src/server/services/history/projection_cache.py b/src/server/services/history/projection_cache.py deleted file mode 100644 index a95627ca4..000000000 --- a/src/server/services/history/projection_cache.py +++ /dev/null @@ -1,312 +0,0 @@ -"""Per-turn replay projection cache. - -Wire-ready replay items for one settled turn, keyed by the turn's own last -checkpoint id (its *tail*) **and** a fingerprint of the turn's table-sourced -replay inputs. The tail exists at persist time and survives forks of later -turns (edits/regenerates fork new tails → new keys; orphaned entries expire -by TTL), but it alone does not identify a complete projection: a settled -turn's stored events, provenance rows and usage row keep being rewritten -after finalize (the subagent collector's archive drain, atomic -``context_window`` appends) without ever moving the tail. Folding those -inputs into the key makes a superseded entry *unreachable* rather than -merely stale, so no invalidation call has to succeed for correctness — and -a reader that built from pre-drain state can only ever write its own -generation's key, never poison the current one. - -Entries are therefore immutable, except the interrupted-turn refresh: a -turn's ``ending_interrupts`` ride the resume boundary created later, so the -post-persist refresh rebuilds a two-turn window to overwrite the previous -turn's entry under its unchanged key. - -Entries are canonicalized to the wire's JSON form (``default=str``) before -storing so a cache hit is byte-identical to a fresh projection. Purely an -optimization: any miss falls back to a full checkpoint rebuild + backfill. -""" - -import asyncio -import hashlib -import json -import logging -from typing import Any - -from src.config.settings import get_replay_projection_cache_ttl -from src.utils.cache.redis_cache import get_cache_client - -logger = logging.getLogger(__name__) - -# v7: on a thread past 1024 checkpoints, v6 entries could hold an early turn -# projected empty and the next turn holding both (utils/checkpointer.py). -# v6: a user-stopped turn's last assistant message carries the stop close. -# v5: a stopped or failed turn's main lane resurrects its uncommitted partial -# (v4 entries of those turns hold the lossy projection under an unchanged key). -# v4: user_message items carry run_completed_at (v3 entries predate the field). -# v3: fingerprinted keys (v1/v2 entries are tail-only, so a v2 hit could -# serve a projection built before the turn's archive drain landed). -_KEY_PREFIX = "replay:turn:v7" -# Entries beyond this are pathological (widget-heavy legacy stored events); -# skip caching rather than bloat Redis — replay just rebuilds those threads. -_MAX_ENTRY_BYTES = 512 * 1024 - -# One strong-referenced runner per thread (create_task results are otherwise -# GC-eligible). A trigger that arrives while its runner is active marks the -# thread dirty; the same runner performs one coalesced follow-up pass after the -# current pass finishes. This preserves refresh generation order, so an older -# snapshot can never finish after and overwrite a newer one. -_refresh_tasks: dict[str, asyncio.Task[None]] = {} -_refresh_dirty: set[str] = set() - - -def _key(thread_id: str, tail_checkpoint_id: str, fingerprint: str) -> str: - return f"{_KEY_PREFIX}:{thread_id}:{tail_checkpoint_id}:{fingerprint}" - - -def turn_fingerprint( - response: dict[str, Any] | None, - provenance_rows: list[dict[str, Any]] | None, - usage_row: dict[str, Any] | None, -) -> str: - """Digest of the post-finalize-mutable replay inputs the tail can't identify. - - Hashes each source whole rather than a chosen subset of fields: an - over-inclusive digest only costs an extra rebuild when something - irrelevant changes, while an under-inclusive one silently resurrects the - bug this key exists to prevent. The response row is the one source read - through a proxy — see ``_response_generation``. - """ - digest = hashlib.blake2b(digest_size=16) - for source in (_response_generation(response), provenance_rows, usage_row): - digest.update( - json.dumps(source, default=str, sort_keys=True).encode("utf-8") - ) - digest.update(b"\x1e") - return digest.hexdigest() - - -def _response_generation( - response: dict[str, Any] | None, -) -> dict[str, Any] | None: - """The response row reduced to its identity plus its Postgres row version. - - Every fingerprinted field of a settled turn lives on that one row, so any - rewrite — the archive drain replacing ``sse_events``, an atomic - ``context_window`` append — lands a new tuple with a new ``xmin``. Reading - the version is strictly more inclusive than hashing the columns and costs - a tuple-header read, where serializing the megabyte-scale transcript here - ran per turn, per replay request, on the event loop. Rows fetched without - ``xmin`` keep the whole-row hash. - """ - if not response or "xmin" not in response: - return response - return { - "conversation_response_id": response.get("conversation_response_id"), - "xmin": response["xmin"], - } - - -def cache_active() -> bool: - client = get_cache_client() - return bool( - client.enabled and client.client and get_replay_projection_cache_ttl() > 0 - ) - - -async def get_cached_turns( - thread_id: str, turn_keys: list[tuple[str, str]] -) -> dict[str, list[dict[str, Any]] | None]: - """Fetch entries for ``(tail, fingerprint)`` pairs; None per tail on miss. - - Keyed by fingerprint, indexed back by tail — callers hold a tail per turn - and a fingerprint is only ever derived from it. - """ - keys = [_key(thread_id, tail, fp) for tail, fp in turn_keys] - values = await get_cache_client().mget(keys) - return dict(zip([tail for tail, _ in turn_keys], values)) - - -async def store_turn( - thread_id: str, - tail_checkpoint_id: str | None, - fingerprint: str, - items: list[dict[str, Any]], -) -> None: - """Best-effort write of one turn's wire-ready items.""" - if not tail_checkpoint_id or not cache_active(): - return - try: - # Canonicalize exactly as the SSE endpoint serializes (default=str), - # so cached datetimes replay with the same wire text as fresh ones. - serialized = json.dumps(items, ensure_ascii=False, default=str) - if len(serialized.encode("utf-8")) > _MAX_ENTRY_BYTES: - logger.debug( - f"[ProjectionCache] entry too large for {thread_id} " - f"tail={tail_checkpoint_id}, skipping" - ) - return - await get_cache_client().set( - _key(thread_id, tail_checkpoint_id, fingerprint), - json.loads(serialized), - ttl=get_replay_projection_cache_ttl(), - ) - except Exception as e: - logger.warning(f"[ProjectionCache] store failed for {thread_id}: {e}") - - -async def delete_turns(thread_id: str, turn_keys: list[tuple[str, str]]) -> None: - """Best-effort eviction — keeps salvage-bearing turns out of the cache so - the fast path can't replay them without their trailing salvage.""" - if not turn_keys or not cache_active(): - return - try: - client = get_cache_client() - for tail_checkpoint_id, fingerprint in turn_keys: - await client.delete(_key(thread_id, tail_checkpoint_id, fingerprint)) - except Exception as e: - logger.warning(f"[ProjectionCache] delete failed for {thread_id}: {e}") - - -async def live_task_streams(thread_id: str, task_ids: set[str]) -> set[str]: - """The subset of tasks whose per-task Redis stream is still unfinalized. - - A turn whose projected task state is still growing must not be cached: - task-ns writes never move the main-branch tail, so a frozen partial - would never self-heal. Only an explicit end sentinel proves terminal - here. Missing streams and verification failures count as live - (conservative — skip caching rather than freeze a partial). - - Probe BEFORE reading task-namespace state, and gate caching on this - pre-read verdict: sealing is monotonic, so "sealed before the read" - proves the read saw final state, while a post-read probe can pass for a - stream that sealed after the read returned pre-terminal data — freezing - it for the cache TTL. - """ - if not task_ids: - return set() - client = get_cache_client().client - if client is None: - return set(task_ids) - - async def _task_live(task_id: str) -> bool: - try: - entries = await client.xrevrange( - f"subagent:stream:{thread_id}:{task_id}", count=1 - ) - except Exception as e: - logger.warning( - f"[ProjectionCache] task stream check failed for " - f"{thread_id}/{task_id}: {e}" - ) - return True # conservative: skip caching rather than freeze a partial - if not entries: - return True - raw = entries[0][1].get(b"event") or entries[0][1].get("event") - if isinstance(raw, bytes): - raw = raw.decode("utf-8", errors="replace") - return not _is_stream_end_sentinel(raw) - - ordered = list(task_ids) - liveness = await asyncio.gather(*(_task_live(task_id) for task_id in ordered)) - return {task_id for task_id, live in zip(ordered, liveness) if live} - - -def _is_stream_end_sentinel(raw: Any) -> bool: - """Match forwarder.finalize()'s ``{"event": "subagent_stream_end"}`` - sentinel (a JSON dict without ``seq`` — same test as the reconnect - consumer's payload classifier).""" - from ptc_agent.agent.middleware.background_subagent.redis_stream import ( - SUBAGENT_STREAM_END_EVENT, - ) - - if not isinstance(raw, str) or not raw.startswith("{"): - return False - try: - record = json.loads(raw) - except json.JSONDecodeError: - return False - return ( - isinstance(record, dict) - and record.get("event") == SUBAGENT_STREAM_END_EVENT - and "seq" not in record - ) - - -def _refresh_done(thread_id: str, task: asyncio.Task[None]) -> None: - """Drop a completed runner without clobbering a newer replacement task.""" - if _refresh_tasks.get(thread_id) is task: - _refresh_tasks.pop(thread_id, None) - _refresh_dirty.discard(thread_id) - - -async def _run_projection_refreshes(thread_id: str) -> None: - """Run refreshes serially, coalescing overlapping triggers per thread.""" - while True: - # Triggers arriving after this point mark the pass dirty and force one - # follow-up pass. Multiple triggers during the same pass coalesce. - _refresh_dirty.discard(thread_id) - await refresh_thread_projection(thread_id) - if thread_id not in _refresh_dirty: - return - - -def schedule_projection_refresh(thread_id: str) -> None: - """Fire-and-forget post-persist warm-up of the thread's last two turns. - - At most one runner executes per thread. A trigger overlapping an active - pass requests a coalesced follow-up pass, preserving snapshot write order. - """ - if not cache_active(): - return - try: - loop = asyncio.get_running_loop() - except RuntimeError: # no running loop (sync test contexts) - return - - current = _refresh_tasks.get(thread_id) - if current is not None and not current.done(): - _refresh_dirty.add(thread_id) - return - - # A completed task's callback may not have run yet. Install the replacement - # first; _refresh_done's identity guard prevents the old callback from - # removing it. - _refresh_dirty.discard(thread_id) - task = loop.create_task(_run_projection_refreshes(thread_id)) - _refresh_tasks[thread_id] = task - task.add_done_callback( - lambda completed, tid=thread_id: _refresh_done(tid, completed) - ) - - -async def refresh_thread_projection(thread_id: str) -> None: - """Rebuild the last two turns so their entries land while checkpoints are - hot — two turns because the previous turn's ``ending_interrupts`` only - become known at this turn's (resume) boundary.""" - from src.server.database.conversation import get_replay_thread_data - from src.server.services.history.replay import ( - CheckpointReplayUnavailable, - build_checkpoint_replay_items, - ) - - try: - _owner, thread, queries, responses, usages, provenance = ( - await get_replay_thread_data(thread_id) - ) - branch_tip = (thread or {}).get("latest_checkpoint_id") - if not branch_tip: - return - responses_by_turn = { - r.get("turn_index"): r for r in responses if isinstance(r, dict) - } - # Backfills the cache as a side effect of the rebuild. - await build_checkpoint_replay_items( - thread_id, - queries, - responses_by_turn, - branch_tip_checkpoint_id=branch_tip, - last_n_turns=2, - usages=usages, - provenance=provenance, - ) - except CheckpointReplayUnavailable as e: - logger.debug(f"[ProjectionCache] refresh skipped for {thread_id}: {e}") - except Exception as e: - logger.warning(f"[ProjectionCache] refresh failed for {thread_id}: {e}") diff --git a/src/server/services/history/projector.py b/src/server/services/history/projector.py index 1ceafbb0c..9260ae1af 100644 --- a/src/server/services/history/projector.py +++ b/src/server/services/history/projector.py @@ -16,7 +16,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage @@ -34,6 +34,9 @@ ) from src.server.utils.text_phase import agreed_phase, block_phase +if TYPE_CHECKING: + from src.server.services.history.slices import SpanDelta + MAIN_AGENT = "main" HistoryEventKind = Literal[ @@ -532,7 +535,7 @@ def _count_lines(text: str) -> int: def context_signal_items( - thread_id: str, turn: Any, *, agent: str = MAIN_AGENT + thread_id: str, turn: SpanDelta, *, agent: str = MAIN_AGENT ) -> list[dict[str, Any]]: """Project a turn's compaction signals from its private-state deltas. @@ -577,7 +580,7 @@ def context_signal_items( def workflow_run_items( - history: Any, + history: SpanDelta, *, task_id: str | None = None, ledger_status: str | None = None, @@ -662,7 +665,7 @@ def _reconcile_terminal_frame( def model_fallback_items( - thread_id: str, turn: Any, *, agent: str = MAIN_AGENT + thread_id: str, turn: SpanDelta, *, agent: str = MAIN_AGENT ) -> list[dict[str, Any]]: """Project a turn's model_fallback notices from its new ``ui`` records. diff --git a/src/server/services/history/reader.py b/src/server/services/history/reader.py index 399b11352..e9a40ed2a 100644 --- a/src/server/services/history/reader.py +++ b/src/server/services/history/reader.py @@ -3,7 +3,7 @@ ``messages`` is a ``DeltaChannel`` (see ``ptc_agent.agent.state``), so a raw ``checkpointer.aget_tuple`` cannot materialize it: deltas live in ``checkpoint_writes`` and are only replayed by a compiled graph. This module -compiles a no-op ``StateGraph(_ReaderState)`` against the server checkpointer +compiles a no-op ``StateGraph(MainAgentState)`` against the server checkpointer purely to read state — in the ``task_namespace_graph`` shape, since background subagents checkpoint under the parent ``thread_id`` with ``checkpoint_ns="task:{task_id}"``. @@ -14,31 +14,28 @@ import asyncio import logging import uuid +from bisect import bisect_right from dataclasses import dataclass, field -from typing import Annotated, Any +from typing import Any -from langchain.agents.middleware.types import PrivateStateAttr -from langchain_core.messages import AnyMessage, HumanMessage from langgraph.graph import START, StateGraph -from typing_extensions import NotRequired -from ptc_agent.agent.middleware.compaction.types import CompactionEvent -from ptc_agent.agent.state import DeltaAgentState -from ptc_agent.agent.transcript.classify import is_run_boundary_message -from src.server.utils.checkpoint_helpers import walk_current_branch_boundaries +from ptc_agent.agent.main_state import MainAgentState +from src.server.services.history import slices +from src.server.utils.checkpoint_helpers import ( + INTERRUPT_CHANNEL, + Boundary, + interrupt_records, + update_at_commit, + walk_current_branch_boundaries, +) logger = logging.getLogger(__name__) -# LangGraph records a pending write on this channel when a node pauses via -# interrupt(). We match the stored channel name directly rather than importing -# the constant, which langgraph made private in v1.0 (deprecated import); the -# on-disk channel name is a stable storage-format detail. -_INTERRUPT_CHANNEL = "__interrupt__" - def _has_pending_interrupt(tup: Any) -> bool: return tup is not None and any( - w[1] == _INTERRUPT_CHANNEL for w in (tup.pending_writes or ()) + w[1] == INTERRUPT_CHANNEL for w in (tup.pending_writes or ()) ) # Concurrent aget_state reads per history materialization. Each read holds a @@ -47,22 +44,6 @@ def _has_pending_interrupt(tup: Any) -> bool: _STATE_READ_CONCURRENCY = 8 -class _ReaderState(DeltaAgentState): - """DeltaAgentState plus middleware-private channels replay materializes. - - ``aget_state`` only surfaces channels the reading graph declares. Live - agents get ``_summarization_event`` from the compaction middleware's state - schema; the reader graph must declare it itself (same annotation, so - channel semantics match the checkpointed writes). - """ - - _summarization_event: Annotated[ - NotRequired[CompactionEvent | None], PrivateStateAttr - ] - _offloaded_tool_call_ids: Annotated[NotRequired[set[str]], PrivateStateAttr] - _offloaded_read_result_ids: Annotated[NotRequired[set[str]], PrivateStateAttr] - - def _silence_pending_sends_noise() -> None: """Suppress langgraph's 'unknown node name … in pending sends' warnings. @@ -89,67 +70,11 @@ def _filter(record: logging.LogRecord) -> bool: @dataclass -class TurnSlice: - """Messages of one conversational turn on the current checkpoint branch.""" - - turn_ordinal: int - input_checkpoint_id: str - end_checkpoint_id: str - user_message: HumanMessage | None - messages: list[AnyMessage] = field(default_factory=list) - run_id: str | None = None - turn_index: int | None = None - # The ``_summarization_event`` that landed during this turn (end-state - # event differing from start-state), or None. Compaction's summary message - # lives in this state key — never in the messages channel — so replay - # re-emits the summarize signal from here. - new_summarization_event: dict[str, Any] | None = None - # Per-turn growth of the compaction offload sets — replay re-emits the - # offload signals from these counts (one aggregated event per kind). - newly_offloaded_args: int = 0 - newly_offloaded_reads: int = 0 - # Answered interrupts this turn raised, read from the FOLLOWING resume - # boundary's ``__interrupt__`` writes. Pending (unanswered) interrupts at - # the branch tip surface via ``ThreadHistory.interrupts`` instead. - ending_interrupts: list[dict[str, Any]] = field(default_factory=list) - # ``ui``-channel records that landed during this turn (id-diff between the - # boundary states) — e.g. model_fallback notices pushed by middleware. - new_ui_records: list[dict[str, Any]] = field(default_factory=list) - # The turn's own last checkpoint — the projection-cache key (see TurnAnchor). - tail_checkpoint_id: str | None = None - - -@dataclass -class ThreadHistory: - thread_id: str - turns: list[TurnSlice] = field(default_factory=list) - interrupts: list[dict[str, Any]] = field(default_factory=list) - ui: list[dict[str, Any]] = field(default_factory=list) +class TaskHistory(slices.SpanDelta): + """A background task's whole namespace: what it added from nothing.""" - -@dataclass -class TaskHistory: - """Materialized state for one background-task checkpoint namespace.""" - - messages: list[AnyMessage] = field(default_factory=list) #: The namespace checkpoint these were read at. checkpoint_id: str | None = None - new_summarization_event: dict[str, Any] | None = None - newly_offloaded_args: int = 0 - newly_offloaded_reads: int = 0 - new_ui_records: list[dict[str, Any]] = field(default_factory=list) - - -@dataclass -class _InputBoundary: - checkpoint_id: str - metadata: dict[str, Any] - is_resume: bool = False - # ``__interrupt__`` pending writes riding a resume boundary — the answered - # interrupts of the turn this boundary resumes (interrupt and resume writes - # attach to the same checkpoint). - interrupts: list[dict[str, Any]] = field(default_factory=list) - parent_checkpoint_id: str | None = None @dataclass @@ -167,76 +92,94 @@ class TurnAnchor: tail_checkpoint_id: str | None turn_index: Any | None = None run_id: str | None = None + # Where the turn's message diff ends: the next boundary, or the tip for + # the last turn. A boundary's state is the thread before its turn + # applies, so the next boundary's state is this turn's result. + end_checkpoint_id: str | None = None + # Answered interrupts this turn raised, from the next resume boundary. + # Pending ones at the branch tip come from ``aget_tip_interrupts``. + ending_interrupts: list[dict[str, Any]] = field(default_factory=list) + # A HITL resume opens on the interrupt's answer, never a message of its + # own: it carries on the turn before it. + is_resume: bool = False + # The state before the turn, where an edit of its message forks; a + # resume has no message of its own to edit. + parent_checkpoint_id: str | None = None -def _tail_of_turn( - boundaries: list[_InputBoundary], i: int, tip_id: str -) -> str | None: - """Turn *i*'s last checkpoint: the tip for the last turn; otherwise the - next boundary's parent — except a resume boundary IS the interrupted - turn's tip (interrupt and resume writes ride the same checkpoint).""" - if i == len(boundaries) - 1: - return tip_id - nxt = boundaries[i + 1] - return nxt.checkpoint_id if nxt.is_resume else nxt.parent_checkpoint_id - - -def _interrupt_writes(cp_tuple: Any) -> list[dict[str, Any]]: - """``{"id", "value"}`` records from a checkpoint's ``__interrupt__`` writes.""" - interrupts: list[dict[str, Any]] = [] - for _task_id, channel, value in (cp_tuple.pending_writes or []) if cp_tuple else []: - if channel != _INTERRUPT_CHANNEL: - continue - values = value if isinstance(value, (list, tuple)) else [value] - for intr in values: - interrupts.append( - { - "id": getattr(intr, "id", None), - "value": getattr(intr, "value", None), - } - ) - return interrupts - - -def _new_summarization_event(start_state: Any, end_state: Any) -> dict[str, Any] | None: - """The turn's freshly landed ``_summarization_event``, or None. +def pair_turns(anchors: list[TurnAnchor], persisted: list[int]) -> list[int | None]: + """Each turn's persisted turn_index, given the thread's persisted ones + ascending. - Events are identified by their summary message id (uuid-stamped at build); - an end-state event matching the start state predates this turn. + A turn's input checkpoint carries the turn_index its run was admitted + under. A resume's checkpoint belongs to the run it answers, and threads + from before the stamp carry none, so those take the next persisted turn + after the one before: a turn whose run died before its first checkpoint + has rows and no boundary, so position alone would name every later turn + one too low. None when no persisted turn is left to name. """ + numbers: list[int | None] = [] + previous: int | None = None + for anchor in anchors: + number = anchor.turn_index + if number is None: + at = 0 if previous is None else bisect_right(persisted, previous) + number = persisted[at] if at < len(persisted) else None + numbers.append(number) + if number is not None: + previous = number + return numbers - def _identity(state: Any) -> tuple[Any, Any] | None: - event = state.values.get("_summarization_event") - if not isinstance(event, dict): - return None - summary_message = event.get("summary_message") - return (getattr(summary_message, "id", None), event.get("cutoff_index")) - - end_event = end_state.values.get("_summarization_event") - if isinstance(end_event, dict) and _identity(end_state) != _identity(start_state): - return end_event - return None +@dataclass +class TaskRun: + """One run of a background task: the span its namespace wrote between + its own input boundary and the next run's (or the namespace tip).""" -def _set_growth(start_state: Any, end_state: Any, key: str) -> int: - """How many ids ``key``'s set gained between the two states.""" - start = start_state.values.get(key) or set() - end = end_state.values.get(key) or set() - return len(set(end) - set(start)) + task_id: str + ordinal: int + input_checkpoint_id: str + end_checkpoint_id: str + # The run's own last checkpoint, which validates a stored run slice. + tail_checkpoint_id: str + # The ledger's run id stamped on the boundary; None before the ledger. + task_run_id: str | None = None -def _new_ui_records(start_state: Any, end_state: Any) -> list[dict[str, Any]]: - """``ui``-channel records the end state has that the start state lacks.""" - start_ids = { - r.get("id") - for r in (start_state.values.get("ui") or []) - if isinstance(r, dict) - } - return [ - r - for r in (end_state.values.get("ui") or []) - if isinstance(r, dict) and r.get("id") not in start_ids - ] +def _tail_of_turn(boundaries: list[Boundary], i: int, tip_id: str) -> str | None: + """Turn *i*'s last checkpoint: the tip for the last turn; otherwise the + next boundary's parent — except a resume boundary IS the interrupted + turn's tip (interrupt and resume writes ride the same checkpoint).""" + if i == len(boundaries) - 1: + return tip_id + nxt = boundaries[i + 1] + return nxt.parent_checkpoint_id if nxt.is_input else nxt.checkpoint_id + + +def turn_anchors(boundaries: list[Boundary], tip_id: str) -> list[TurnAnchor]: + """The branch's turns, one per boundary of its walk.""" + anchors: list[TurnAnchor] = [] + for i, boundary in enumerate(boundaries): + is_last = i == len(boundaries) - 1 + # A resume checkpoint's metadata belongs to the interrupted run, + # not the resume turn — don't propagate its run_id/turn_index. + metadata = boundary.metadata if boundary.is_input else {} + anchors.append( + TurnAnchor( + turn_ordinal=i, + input_checkpoint_id=boundary.checkpoint_id, + tail_checkpoint_id=_tail_of_turn(boundaries, i, tip_id), + turn_index=metadata.get("turn_index"), + run_id=metadata.get("run_id"), + end_checkpoint_id=( + tip_id if is_last else boundaries[i + 1].checkpoint_id + ), + ending_interrupts=[] if is_last else boundaries[i + 1].interrupts, + is_resume=not boundary.is_input, + parent_checkpoint_id=boundary.parent_checkpoint_id, + ) + ) + return anchors class CheckpointHistoryReader: @@ -253,15 +196,20 @@ def __init__(self, checkpointer: Any): _silence_pending_sends_noise() self._checkpointer = checkpointer - # Same shape the snapshot writer uses, so what it wrote into - # task:{id} is what this reads back. - self._graph = task_namespace_graph(_ReaderState, checkpointer) + # The main agent's own schema, for two reasons: ``aget_state`` only + # surfaces channels the reading graph declares, and the ui-record + # append writes through this graph, which would erase any primitive + # field it left undeclared (see ``main_state``). Same shape the + # snapshot writer uses, so what it wrote into task:{id} is what this + # reads back. Public: checkpoint storage maintenance decodes channels + # through the same declaration. + self.graph = task_namespace_graph(MainAgentState, checkpointer) # Separate single-node graph for ui-record appends: with exactly one # node, aupdate_state auto-attributes the write (no as_node needed), # and the update checkpoint carries source="update" — never a turn # boundary. self._updater = ( - StateGraph(_ReaderState) + StateGraph(MainAgentState) .add_node("noop", lambda state: {}) .add_edge(START, "noop") .compile(checkpointer=checkpointer) @@ -281,127 +229,100 @@ def get_instance(cls) -> CheckpointHistoryReader: def reset_instance(cls) -> None: cls._instance = None - async def aget_thread_history( - self, thread_id: str, branch_tip_checkpoint_id: str | None = None - ) -> ThreadHistory: - """Materialize per-turn message slices for the current branch. - - Turn boundaries are ``source=input`` checkpoints plus HITL resume - points (``__resume__`` pending writes) — each persists its own query - row, so boundaries stay 1:1 with turns. A boundary checkpoint's state - is the thread BEFORE its turn applies (the input/resume rides pending - writes), so turn *i* = id-diff between boundary *i* and boundary *i+1* - (or the tip). Id-diff, not count-diff, so compaction/REMOVE_ALL is - safe. - """ - boundaries, tip_id = await self._current_branch_input_boundaries( - thread_id, branch_tip_checkpoint_id + @property + def serde(self) -> Any: + """The checkpointer's serializer, which stored slices share.""" + return self._checkpointer.serde + + async def aget_turn_slices( + self, thread_id: str, anchors: list[TurnAnchor] + ) -> list[slices.TurnSlice]: + """Materialize only these turns: each reads its boundary and the + state its diff ends at, shared between adjacent turns, so the cost + follows the turns asked for rather than the thread's length.""" + ids = list( + dict.fromkeys( + cid + for a in anchors + for cid in (a.input_checkpoint_id, a.end_checkpoint_id) + if cid + ) ) - history = ThreadHistory(thread_id=thread_id) - if not boundaries or tip_id is None: - return history + states = dict(zip(ids, await self._aget_states_at(thread_id, ids))) + return [ + slices.turn_slice( + a, + states[a.input_checkpoint_id].values, + states[a.end_checkpoint_id].values, + ) + for a in anchors + ] - # Materialize each boundary state once, plus the branch tip. - *boundary_states, tip_state = await self._aget_states_at( - thread_id, [b.checkpoint_id for b in boundaries] + [tip_id] - ) - history.turns = self._build_turn_slices( - boundaries, boundary_states, tip_id, tip_state - ) - history.interrupts = await self._extract_interrupts(thread_id, tip_id) - history.ui = list(tip_state.values.get("ui", []) or []) - return history - - async def aget_recent_history( - self, - thread_id: str, - n_turns: int, - branch_tip_checkpoint_id: str | None = None, - ) -> ThreadHistory: - """Materialize only the last ``n_turns`` turns (windowed replay). - - Same id-diff slicing as ``aget_thread_history``, but materializes just - the tail boundaries + tip — ``n_turns + 1`` state reads instead of one - per turn — so initial-load latency is bounded by the window, not by - thread length. ``turn_ordinal`` stays absolute. + async def aget_task_runs(self, thread_id: str, task_id: str) -> list[TaskRun]: + """A task namespace's runs, oldest first, from a light walk. + + Each spawn or resume opens a run at a ``source=input`` boundary + stamped with its ledger run id; any other boundary resumes the run + before it. """ - boundaries, tip_id = await self._current_branch_input_boundaries( - thread_id, branch_tip_checkpoint_id + boundaries, tip_id = await walk_current_branch_boundaries( + self._checkpointer, thread_id, checkpoint_ns=f"task:{task_id}" ) - history = ThreadHistory(thread_id=thread_id) if not boundaries or tip_id is None: - return history - - n = max(1, min(n_turns, len(boundaries))) - kept = boundaries[-n:] - *kept_states, tip_state = await self._aget_states_at( - thread_id, [b.checkpoint_id for b in kept] + [tip_id] - ) - history.turns = self._build_turn_slices( - kept, kept_states, tip_id, tip_state, ordinal_offset=len(boundaries) - n - ) - history.interrupts = await self._extract_interrupts(thread_id, tip_id) - history.ui = list(tip_state.values.get("ui", []) or []) - return history - - @staticmethod - def _build_turn_slices( - boundaries: list[_InputBoundary], - states: list[Any], - tip_id: str, - tip_state: Any, - ordinal_offset: int = 0, - ) -> list[TurnSlice]: - """Id-diff each boundary against the next (or the tip) into a TurnSlice. - - ``boundaries``/``states`` are a contiguous run; the last boundary's end - is the tip (callers pass the run that ends at the branch tip). - """ - turns: list[TurnSlice] = [] - for i, boundary in enumerate(boundaries): - start_msgs: list[AnyMessage] = states[i].values.get("messages", []) - is_last = i == len(boundaries) - 1 - end_state = tip_state if is_last else states[i + 1] - end_id = tip_id if is_last else boundaries[i + 1].checkpoint_id - end_msgs: list[AnyMessage] = end_state.values.get("messages", []) - - start_ids = {m.id for m in start_msgs if m.id is not None} - slice_msgs = [m for m in end_msgs if m.id not in start_ids] - - # A stamped injection is not the turn's input: a runtime-update - # row lands in the same slice, right after the real user message. - user_message = next( - (m for m in slice_msgs if is_run_boundary_message(m)), None - ) - # A resume checkpoint's metadata belongs to the interrupted run, - # not the resume turn — don't propagate its run_id/turn_index. - metadata = {} if boundary.is_resume else (boundary.metadata or {}) - turns.append( - TurnSlice( - turn_ordinal=ordinal_offset + i, - input_checkpoint_id=boundary.checkpoint_id, - end_checkpoint_id=end_id, - user_message=user_message, - messages=slice_msgs, - run_id=metadata.get("run_id"), - turn_index=metadata.get("turn_index"), - new_summarization_event=_new_summarization_event( - states[i], end_state - ), - newly_offloaded_args=_set_growth( - states[i], end_state, "_offloaded_tool_call_ids" - ), - newly_offloaded_reads=_set_growth( - states[i], end_state, "_offloaded_read_result_ids" + return [] + starts = [b for i, b in enumerate(boundaries) if i == 0 or b.is_input] + runs: list[TaskRun] = [] + for k, start in enumerate(starts): + nxt = starts[k + 1] if k + 1 < len(starts) else None + stamp = start.metadata.get("task_run_id") + runs.append( + TaskRun( + task_id=task_id, + ordinal=k, + input_checkpoint_id=start.checkpoint_id, + end_checkpoint_id=nxt.checkpoint_id if nxt else tip_id, + tail_checkpoint_id=( + (nxt.parent_checkpoint_id or nxt.checkpoint_id) + if nxt + else tip_id ), - ending_interrupts=( - [] if is_last else boundaries[i + 1].interrupts - ), - new_ui_records=_new_ui_records(states[i], end_state), - tail_checkpoint_id=_tail_of_turn(boundaries, i, tip_id), + task_run_id=str(stamp) if stamp else None, + ) + ) + return runs + + async def aget_run_slices( + self, thread_id: str, task_id: str, runs: list[TaskRun] + ) -> list[slices.SpanDelta]: + """What each run added to its namespace, two state reads per run, + shared between adjacent runs. Diffing each run's own boundaries keeps + early runs whole after the namespace compacts them away.""" + ids = list( + dict.fromkeys( + cid for r in runs for cid in (r.input_checkpoint_id, r.end_checkpoint_id) + ) + ) + semaphore = asyncio.Semaphore(_STATE_READ_CONCURRENCY) + + async def _one(checkpoint_id: str) -> Any: + async with semaphore: + return await self.graph.aget_state( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": f"task:{task_id}", + "checkpoint_id": checkpoint_id, + } + } ) + + states = dict(zip(ids, await asyncio.gather(*(_one(c) for c in ids)))) + return [ + slices.delta( + states[r.input_checkpoint_id].values, states[r.end_checkpoint_id].values ) - return turns + for r in runs + ] async def append_ui_record( self, thread_id: str, name: str, props: dict[str, Any] @@ -417,16 +338,10 @@ async def append_ui_record( The image-capture hook runs after the turn finalized, so the append sits beyond the recorded branch tip where the replay walk cannot see - it — the recorded tip is CAS-advanced onto the new checkpoint (no-op - when a concurrent turn or branch switch moved the tip first). - - The append is anchored to the tip it read rather than to the - checkpointer's latest, because the same value is also the CAS guard: - left unanchored, a turn that starts on another worker between the read - and the write re-parents the record onto that turn's uncommitted - checkpoint while the CAS still passes, publishing partial state as the - thread's commit pointer. Anchoring costs a dead-branch record when the - tip moves — which is the no-op case this already accepts. + it: it is written at the tip read here and the recorded tip + CAS-advanced onto it (``update_at_commit``). A concurrent turn or + branch switch that moved the tip first leaves a dead-branch record, + and a thread with no checkpoint gets none. """ try: tip = await self._checkpointer.aget_tuple( @@ -457,52 +372,23 @@ async def append_ui_record( "props": props, "metadata": {}, } - tip_cfg = (tip.config.get("configurable") or {}) if tip is not None else {} - built_on = tip_cfg.get("checkpoint_id") - configurable: dict[str, Any] = {"thread_id": thread_id} - if built_on: - # The saver keys a write by (thread_id, checkpoint_ns, checkpoint_id), - # so the namespace has to ride along with the id. - configurable["checkpoint_ns"] = tip_cfg.get("checkpoint_ns", "") - configurable["checkpoint_id"] = built_on - new_config = await self._updater.aupdate_state( - {"configurable": configurable}, {"ui": [record]} - ) - await self._advance_branch_tip(thread_id, built_on, new_config) - - async def _advance_branch_tip( - self, thread_id: str, built_on: str | None, new_config: Any - ) -> None: - """CAS the recorded tip from the checkpoint the append was anchored to. - - ``built_on`` is the caller's write anchor, not a re-read: passing the - same value to both is what keeps the guard honest. - """ - from src.server.database.conversation.threads_write import ( - advance_thread_checkpoint_id, - ) - - new_id = ((new_config or {}).get("configurable") or {}).get( - "checkpoint_id" - ) - if not new_id: - return - await advance_thread_checkpoint_id( - thread_id, from_checkpoint_id=built_on, to_checkpoint_id=new_id + built_on = ( + (tip.config.get("configurable") or {}).get("checkpoint_id") if tip else None ) + await update_at_commit(self._updater, thread_id, built_on, {"ui": [record]}) async def aget_state(self, thread_id: str, checkpoint_id: str | None = None): """The thread's state at ``checkpoint_id``, or at its latest checkpoint.""" configurable = {"thread_id": thread_id} if checkpoint_id: configurable["checkpoint_id"] = checkpoint_id - return await self._graph.aget_state({"configurable": configurable}) + return await self.graph.aget_state({"configurable": configurable}) async def aget_task_history( self, thread_id: str, task_id: str ) -> TaskHistory: """Materialize replay-relevant state from a ``task:{task_id}`` namespace.""" - snapshot = await self._graph.aget_state( + snapshot = await self.graph.aget_state( { "configurable": { "thread_id": thread_id, @@ -510,29 +396,11 @@ async def aget_task_history( } } ) - values = snapshot.values - summarization_event = values.get("_summarization_event") return TaskHistory( - messages=list(values.get("messages", []) or []), checkpoint_id=(snapshot.config or {}) .get("configurable", {}) .get("checkpoint_id"), - new_summarization_event=( - summarization_event - if isinstance(summarization_event, dict) - else None - ), - newly_offloaded_args=len( - set(values.get("_offloaded_tool_call_ids") or ()) - ), - newly_offloaded_reads=len( - set(values.get("_offloaded_read_result_ids") or ()) - ), - new_ui_records=[ - record - for record in (values.get("ui") or []) - if isinstance(record, dict) - ], + **vars(slices.delta({}, snapshot.values)), ) async def alatest_checkpoint_id( @@ -548,49 +416,19 @@ async def alatest_checkpoint_id( return None return (tip.config.get("configurable") or {}).get("checkpoint_id") - async def aget_task_run_stamps( - self, thread_id: str, task_id: str - ) -> list[str | None]: - """Per-run ``task_run_id`` stamps for a task namespace, oldest-first. - - Each spawn/resume writes a ``source=input`` boundary whose metadata - carries the execution's ledger identity (stamped at spawn via the - child-graph config) — stamp k belongs to run segment k of the - namespace transcript. None entries are pre-ledger runs. - """ - boundaries, _tip = await walk_current_branch_boundaries( - self._checkpointer, - thread_id, - checkpoint_ns=f"task:{task_id}", - ) - return [(b.metadata or {}).get("task_run_id") for b in boundaries] - async def aget_turn_anchors( self, thread_id: str, branch_tip_checkpoint_id: str | None = None ) -> tuple[list[TurnAnchor], str | None]: - """Turn identities on the current branch — light walk, no state reads. - - Lets the projection cache pair and key turns without materializing - any checkpoint state. - """ - boundaries, tip_id = await self._current_branch_input_boundaries( - thread_id, branch_tip_checkpoint_id + """Turn identities on the current branch — light walk, no state reads.""" + boundaries, tip_id = await walk_current_branch_boundaries( + self._checkpointer, + thread_id, + branch_tip_checkpoint_id, + strict_branch_tip=branch_tip_checkpoint_id is not None, ) if not boundaries or tip_id is None: return [], tip_id - anchors: list[TurnAnchor] = [] - for i, boundary in enumerate(boundaries): - metadata = {} if boundary.is_resume else (boundary.metadata or {}) - anchors.append( - TurnAnchor( - turn_ordinal=i, - input_checkpoint_id=boundary.checkpoint_id, - tail_checkpoint_id=_tail_of_turn(boundaries, i, tip_id), - turn_index=metadata.get("turn_index"), - run_id=metadata.get("run_id"), - ) - ) - return anchors, tip_id + return turn_anchors(boundaries, tip_id), tip_id async def aget_tip_interrupts( self, thread_id: str, tip_checkpoint_id: str @@ -598,17 +436,8 @@ async def aget_tip_interrupts( """Pending (unanswered) interrupts at the branch tip.""" return await self._extract_interrupts(thread_id, tip_checkpoint_id) - async def acount_input_boundaries( - self, thread_id: str, branch_tip_checkpoint_id: str | None = None - ) -> int: - """Number of turn boundaries on the current branch (auto-mode coverage).""" - boundaries, _ = await self._current_branch_input_boundaries( - thread_id, branch_tip_checkpoint_id - ) - return len(boundaries) - async def _aget_state_at(self, thread_id: str, checkpoint_id: str): - return await self._graph.aget_state( + return await self.graph.aget_state( {"configurable": {"thread_id": thread_id, "checkpoint_id": checkpoint_id}} ) @@ -624,37 +453,6 @@ async def _one(checkpoint_id: str) -> Any: return list(await asyncio.gather(*(_one(cid) for cid in checkpoint_ids))) - async def _current_branch_input_boundaries( - self, thread_id: str, branch_tip_checkpoint_id: str | None - ) -> tuple[list[_InputBoundary], str | None]: - """Turn boundaries on the current branch, as ``_InputBoundary`` records. - - Wraps the canonical ``walk_current_branch_boundaries`` — every boundary - is a ``source=input`` checkpoint or (by construction of the walk) a HITL - resume, so ``is_resume`` is simply "not input". - """ - boundaries, tip_id = await walk_current_branch_boundaries( - self._checkpointer, - thread_id, - branch_tip_checkpoint_id, - strict_branch_tip=branch_tip_checkpoint_id is not None, - ) - records = [ - _InputBoundary( - checkpoint_id=cp.config["configurable"]["checkpoint_id"], - metadata=dict(cp.metadata or {}), - is_resume=(cp.metadata or {}).get("source") != "input", - interrupts=_interrupt_writes(cp), - parent_checkpoint_id=( - (cp.parent_config or {}) - .get("configurable", {}) - .get("checkpoint_id") - ), - ) - for cp in boundaries - ] - return records, tip_id - async def _extract_interrupts( self, thread_id: str, tip_checkpoint_id: str ) -> list[dict[str, Any]]: @@ -672,4 +470,4 @@ async def _extract_interrupts( } } ) - return _interrupt_writes(cp_tuple) + return interrupt_records(cp_tuple.pending_writes if cp_tuple else None) diff --git a/src/server/services/history/replay/__init__.py b/src/server/services/history/replay/__init__.py index 4fe013fc2..624c40bf4 100644 --- a/src/server/services/history/replay/__init__.py +++ b/src/server/services/history/replay/__init__.py @@ -1,6 +1,6 @@ -"""Assemble checkpoint-sourced replay items for the replay endpoint. +"""Assemble checkpoint-sourced replay for the replay endpoint. -Produces the same ``{"event": type, "data": dict}`` items as the stored +Produces the same ``{"event": type, "data": dict}`` stream as the stored ``sse_events`` path, sourcing the transcript from checkpoints (via ``CheckpointHistoryReader`` + the pure projector) and merging the non-derivable remainder from persisted events: @@ -23,67 +23,136 @@ event inlines resolved data files that are deliberately kept out of the checkpointer. - Sandbox image paths in projected text resolve through ``image_capture`` - ui records; a turn whose images cannot be resolved falls back to its stored - events wholesale (turn-level granularity keeps ordering coherent). - -Package layout: this root orchestrates; items builds table-sourced events, -stored_merge anchors persisted payloads, task_lane + segment_claim project -background-task namespaces, widgets inlines offloaded payloads on the way out. + ui records, a subagent run's through its replay facts, and an older turn's + through the URLs its stored events were captured to (``legacy``). +- What only the live stream measured or returned (reasoning durations, a + subagent run's returned steering) comes from the run's replay facts + (``facts``), and from the stored events for a run settled before those. +- A row backfilled with its legacy facts (``legacy``) takes everything the + stored events gave it from those facts instead, and its stored events are + never read. + +A settled turn is projected once: its slice (what the turn added to the +checkpoint) and its projected lines are stored with its response row, and a +read streams the lines while their key still matches the inputs. Only a turn +whose key moved is projected again, from its stored slice when the branch +still agrees with it, and from checkpoints otherwise. + +A turn that cannot be projected replays from its stored events (or as a +stub) on the read paths that ask for it, so one bad turn costs only itself. + +Package layout: + +- this root pairs checkpoint turns with persisted ones, pages them, and + serves the stored lines whose key still matches; +- ``cold`` projects and stores the rest, a batch at a time; +- ``turn`` projects one turn: ``items`` builds its table-sourced events, + ``run_lane`` its background-task runs, ``stored_merge`` resolves its + persisted payloads into legacy facts, ``legacy`` applies them, and + ``facts`` applies what runs recorded on their rows; +- ``lines`` holds the stored wire format, and ``widgets`` inlines offloaded + payloads on the way out; +- ``keys`` decides whether stored lines still stand, ``stored_events`` + replays a turn from its stored events, ``refresh`` projects turns as they + settle, and ``errors`` says why a read gives up. + +Only the modules that decide what lines hold are digested into the lines +key (``keys.projection_sources``). """ from __future__ import annotations +import asyncio import logging +from dataclasses import dataclass from typing import Any -from ptc_agent.agent.middleware.image_capture import ( - IMAGE_MD_RE, - is_sandbox_image_path, -) -from src.server.services.runs.sse_producer import resolve_token_threshold -from src.server.services.history import projection_cache -from src.server.services.history import projector +from src.server.database import conversation as conversation_db +from src.server.database.conversation import turn_slices as slices_db +from src.server.database.conversation.replay_rows import ThreadRows +from src.server.database.pool import AppDataPoolTimeout +from src.server.services.history import slices from src.server.services.history import task_status -from src.server.services.history.projector import ( - history_events_to_sse, - messages_to_history_events, +from src.server.services.history.reader import ( + CheckpointHistoryReader, + TurnAnchor, + pair_turns, ) -from src.server.services.history.reader import CheckpointHistoryReader -from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound +from src.server.services.history.replay import cold from src.server.services.history.replay import items -from src.server.services.history.replay import stopped -from src.server.services.history.replay import stored_merge -from src.server.services.history.replay import task_lane +from src.server.services.history.replay import lines as replay_lines from src.server.services.history.replay import widgets +from src.server.services.history.replay.errors import ( + CheckpointReplayUnavailable, + ClaimsPassNeeded, +) +from src.server.services.history.replay.keys import lines_epoch, lines_key +from src.server.services.history.replay.lines import Line, line_of +from src.server.services.history.replay.stored_events import ( + build_sse_replay_items, + with_stored_events, +) +from src.server.services.history.replay.turn import Derive, Inputs +from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound + +__all__ = [ + "CheckpointReplayUnavailable", + "ClaimsPassNeeded", + "Derive", + "ReplayPage", + "build_checkpoint_replay_items", + "build_replay_page", + "build_sse_replay_items", + "ThreadRows", + "finish_lines", + "lines_epoch", + "load_thread_inputs", + "pair_anchors", + "project_turns", + "read_replay_page", + "stored_events_page", + "unserved_turns", + "with_stored_events", +] logger = logging.getLogger(__name__) -class CheckpointReplayUnavailable(Exception): - """Checkpoint history cannot faithfully cover this thread's replay.""" +@dataclass +class ReplayPage: + lines: list[Line] + # The oldest turn in this page; the next page asks for turns before it. + first_turn_index: Any | None + has_more: bool -IMAGE_CAPTURE_UI_NAME = "image_capture" +async def load_thread_inputs(thread_id: str) -> tuple[ThreadRows, str] | None: + """A thread's replay rows and the commit pointer to read its branch + from, or None when the thread is unknown or never committed a turn.""" + rows = await conversation_db.get_replay_thread_data(thread_id) + if rows is None: + return None + tip = rows.thread.get("latest_checkpoint_id") + return (rows, tip) if tip else None -async def build_checkpoint_replay_items( - thread_id: str, - queries: list[dict[str, Any]], - responses_by_turn: dict[Any, dict[str, Any]], +async def build_replay_page( + rows: ThreadRows, branch_tip_checkpoint_id: str | None = None, - last_n_turns: int | None = None, - usages: list[dict[str, Any]] | None = None, - provenance: list[dict[str, Any]] | None = None, -) -> list[dict[str, Any]]: - """Build the replay item list from checkpoints. - - With ``last_n_turns`` set, only the most recent N turns are materialized - (windowed initial load — latency bounded by the window, not thread length); - otherwise the full thread is built. + *, + before_turn: Any | None = None, + limit: int | None = None, + turn_fallback: bool = False, + cache: bool = True, +) -> ReplayPage: + """Replay lines for one page of the thread: with ``limit``, the newest N + checkpointed turns before ``before_turn`` (all of them without it), plus + the persisted turns with no committed boundary among them. Turns pair to persisted query rows by stamped ``turn_index`` metadata - (ordinal-anchored for pre-stamping threads); a persisted turn with no + (the next persisted turn for a resume or a pre-stamping thread, see + ``pair_turns``); a persisted turn with no committed boundary — the in-flight active turn — replays as its user_message stub only. Raises ``CheckpointReplayUnavailable`` when coverage cannot be established (missing checkpoints, inconsistent @@ -93,532 +162,407 @@ async def build_checkpoint_replay_items( completed response with no boundary, so pairing raises and they stay on the sse path. - Settled turns serve from the per-turn projection cache when every entry - is present (no state materialization); any miss rebuilds from checkpoints - and backfills the cache. Widget ``data_ref`` resolution always runs on - the way out — entries store the unresolved ref. + Lines are wire-ready but not final: task cards and offloaded widget data + are stamped on the way out (``finish_lines``). ``turn_fallback``: see + ``cold.project_stale``; ``cache``: see ``cold.Policy``. """ reader = CheckpointHistoryReader.get_instance() - turn_indexes = sorted( - { - q.get("turn_index") - for q in queries - if isinstance(q, dict) and q.get("turn_index") is not None - } + branch = await _branch(reader, rows, branch_tip_checkpoint_id) + if branch is None: + raise CheckpointReplayUnavailable("no checkpoint turns found") + inputs, anchored, tip_id = branch + page, has_more = _select_page(anchored, inputs, before_turn, limit) + + lines = _turn_lines( + reader, + inputs, + anchored, + [(ti, a) for ti, a in page if a is not None], + cold.Policy(turn_fallback=turn_fallback, cache=cache), ) - queries_by_turn: dict[Any, list[dict[str, Any]]] = {} - for q in queries: - if isinstance(q, dict): - queries_by_turn.setdefault(q.get("turn_index"), []).append(q) - - # Both paths key the projection cache on these rows, so index once. - usage_by_response = items._usage_rows_by_response(usages) - provenance_by_response = items._rows_by_response(provenance, many=True) - - out: list[dict[str, Any]] | None = None - if projection_cache.cache_active(): - out = await _assemble_from_cache( - reader, - thread_id, - queries_by_turn, - responses_by_turn, - turn_indexes, - branch_tip_checkpoint_id, - last_n_turns, - usage_by_response, - provenance_by_response, - ) - if out is None: - out = await _build_and_backfill( - reader, - thread_id, - queries_by_turn, - responses_by_turn, - turn_indexes, - branch_tip_checkpoint_id, - last_n_turns, - usage_by_response, - provenance_by_response, - ) - await widgets._resolve_widget_data_refs(out) - return out - - -def _turn_fingerprint( - response: dict[str, Any] | None, - usage_by_response: dict[str, Any], - provenance_by_response: dict[str, Any], -) -> str: - """Cache-key fingerprint for one turn's table-sourced replay inputs. - - Degrades to a constant on failure rather than propagating: the cache is - an optimization, so an unhashable row must cost a rebuild, never the - replay itself. Reads and writes agree on the fallback, and tails are - unique per turn, so it cannot collide across turns. - """ - response_id = str(response.get("conversation_response_id")) if response else None - try: - return projection_cache.turn_fingerprint( - response, - provenance_by_response.get(response_id) or [], - usage_by_response.get(response_id), + tip_interrupts: list[dict[str, Any]] = [] + if before_turn is None: + # The newest page ends with the tip's pending interrupts, a checkpoint + # read the lines never wait on. Both settle before either error is + # raised, the lines' first, so a failed page leaves no read behind. + results = await asyncio.gather( + lines, + reader.aget_tip_interrupts(rows.thread_id, tip_id), + return_exceptions=True, ) - except Exception: - logger.warning( - f"[REPLAY] turn fingerprint failed for response={response_id}", - exc_info=True, - ) - return "unfingerprinted" - - -async def _assemble_from_cache( - reader: CheckpointHistoryReader, - thread_id: str, - queries_by_turn: dict[Any, list[dict[str, Any]]], - responses_by_turn: dict[Any, dict[str, Any]], - turn_indexes: list[Any], - branch_tip_checkpoint_id: str | None, - last_n_turns: int | None, - usage_by_response: dict[str, Any], - provenance_by_response: dict[str, Any], -) -> list[dict[str, Any]] | None: - """Concatenate cached per-turn entries — a light boundary walk plus one - raw tip read, no state materialization. Returns None on any miss (the - caller rebuilds and backfills). Pairing guards raise the same - ``CheckpointReplayUnavailable`` signals as the full build.""" - try: - anchors, tip_id = await reader.aget_turn_anchors( - thread_id, branch_tip_checkpoint_id - ) - except CheckpointBranchTipNotFound as e: - raise CheckpointReplayUnavailable(str(e)) from e - if not anchors or tip_id is None or any( - a.tail_checkpoint_id is None for a in anchors - ): - return None - if last_n_turns is not None: - anchors = anchors[-max(1, min(last_n_turns, len(anchors))) :] - - pairs = _pair_turns_to_queries( - anchors, turn_indexes, responses_by_turn, windowed=last_n_turns is not None - ) - cached = await projection_cache.get_cached_turns( - thread_id, - [ - ( - a.tail_checkpoint_id, - _turn_fingerprint( - responses_by_turn.get(ti), - usage_by_response, - provenance_by_response, - ), - ) - for ti, a in pairs - if a is not None - ], - ) - if any(v is None for v in cached.values()): - return None - - out: list[dict[str, Any]] = [] - for turn_index, anchor in pairs: + for result in results: + if isinstance(result, BaseException): + raise result + lines_by_turn, tip_interrupts = results + else: + lines_by_turn = await lines + out: list[Line] = [] + for turn_index, anchor in page: if anchor is None: out.extend( - items._stub_turn_items( - thread_id, turn_index, queries_by_turn, responses_by_turn + line_of(item) + for item in items._stub_turn_items( + rows.thread_id, + turn_index, + inputs.queries_by_turn, + inputs.responses_by_turn, ) ) else: - out.extend(cached[anchor.tail_checkpoint_id]) - for interrupt in await reader.aget_tip_interrupts(thread_id, tip_id): - out.append(items._interrupt_item(thread_id, interrupt)) - return out + out.extend(lines_by_turn[turn_index]) + for interrupt in tip_interrupts: + out.append(line_of(items._interrupt_item(rows.thread_id, interrupt))) + return ReplayPage( + lines=out, + first_turn_index=page[0][0] if page else None, + has_more=has_more, + ) -async def _build_and_backfill( - reader: CheckpointHistoryReader, - thread_id: str, - queries_by_turn: dict[Any, list[dict[str, Any]]], - responses_by_turn: dict[Any, dict[str, Any]], - turn_indexes: list[Any], - branch_tip_checkpoint_id: str | None, - last_n_turns: int | None, - usage_by_response: dict[str, Any], - provenance_by_response: dict[str, Any], -) -> list[dict[str, Any]]: - """Materialize checkpoint state and project every requested turn, storing - each settled turn's finished segment in the projection cache.""" - try: - if last_n_turns is not None: - history = await reader.aget_recent_history( - thread_id, last_n_turns, branch_tip_checkpoint_id +async def read_replay_page( + rows: ThreadRows, + *, + source: str = "auto", + before_turn: Any | None = None, + limit: int | None = None, +) -> tuple[ReplayPage, str]: + """The page a thread replays as, and its source (``checkpoint`` or ``sse``). + + One assembly for every reader that shows a thread, so the owner and a + share viewer see the same turns. ``auto`` projects from checkpoints and + replays only a turn that cannot be projected from its stored events; the + whole page goes to storage only when the thread's turns cannot be paired + at all. ``checkpoint`` raises ``CheckpointReplayUnavailable`` rather than + fall back, and ``sse`` reads storage alone. + """ + if source in ("auto", "checkpoint"): + try: + if rows.thread.get("latest_checkpoint_id") is None: + # The commit pointer (stamped at turn persist) is the only tip + # checkpoint replay may read: without it the reader would walk + # the newest checkpoint, which mid-run is uncommitted state. + raise CheckpointReplayUnavailable( + "thread has no committed checkpoint pointer" + ) + page = await build_replay_page( + rows, + rows.thread.get("latest_checkpoint_id"), + before_turn=before_turn, + limit=limit, + turn_fallback=source == "auto", ) - else: - history = await reader.aget_thread_history( - thread_id, branch_tip_checkpoint_id + return page, "checkpoint" + except AppDataPoolTimeout: + # Storage is read through this same pool, so a fallback would wait + # on it again before failing the same way. A checkpointer pool + # timeout falls back like any other failure. + raise + except CheckpointReplayUnavailable as e: + if source == "checkpoint": + raise + logger.info( + "[REPLAY] Checkpoint replay unavailable for %s, falling back " + "to sse: %s", + rows.thread_id, + e, ) - except CheckpointBranchTipNotFound as e: - raise CheckpointReplayUnavailable(str(e)) from e + except Exception: + if source == "checkpoint": + raise + logger.warning( + "[REPLAY] Checkpoint replay failed for %s, falling back to sse", + rows.thread_id, + exc_info=True, + ) + return await stored_events_page(rows, before_turn, limit), "sse" - if not history.turns: - raise CheckpointReplayUnavailable("no checkpoint turns found") - pairs = _pair_turns_to_queries( - history.turns, - turn_indexes, - responses_by_turn, - windowed=last_n_turns is not None, + +async def stored_events_page( + rows: ThreadRows, + before_turn: Any | None = None, + limit: int | None = None, +) -> ReplayPage: + """A page replayed from stored events alone, over the same window the + checkpoint path pages by, counted in persisted turns.""" + turns = sorted( + { + q["turn_index"] + for q in rows.queries + if isinstance(q, dict) + and q.get("turn_index") is not None + and (before_turn is None or q["turn_index"] < before_turn) + } + ) + page_turns = set(turns[-limit:] if limit is not None else turns) + page_queries = [ + q + for q in rows.queries + if isinstance(q, dict) and q.get("turn_index") in page_turns + ] + responses = await with_stored_events( + rows.responses_by_turn, sorted(page_turns), verbatim=True + ) + return ReplayPage( + lines=[ + line_of(item) + for item in build_sse_replay_items(rows.thread_id, page_queries, responses) + ], + first_turn_index=min(page_turns) if page_turns else None, + has_more=len(page_turns) < len(turns), ) - out: list[dict[str, Any]] = [] - lane = task_lane.TaskLaneProjector(thread_id, windowed=last_n_turns is not None) - await lane.prepare(reader, pairs) - # Stores are deferred past trailing_items(): only then is it known which - # turns carry trailing salvage and must stay out of the cache. - cacheable: list[tuple[Any, str | None, str, list[dict[str, Any]]]] = [] - for turn_index, turn in pairs: - if turn is None: - out.extend( - items._stub_turn_items( - thread_id, turn_index, queries_by_turn, responses_by_turn - ) - ) - continue +async def build_checkpoint_replay_items( + rows: ThreadRows, + branch_tip_checkpoint_id: str | None = None, + last_n_turns: int | None = None, + *, + cache: bool = True, +) -> list[dict[str, Any]]: + """The newest ``last_n_turns`` (or every turn) as decoded replay items, + widget data inlined. Task cards are left unstamped.""" + page = await build_replay_page( + rows, branch_tip_checkpoint_id, limit=last_n_turns, cache=cache + ) + out = [line.item() for line in page.lines] + await widgets._resolve_widget_data_refs(out) + return out - response = responses_by_turn.get(turn_index) - response_id = ( - str(response.get("conversation_response_id")) if response else None - ) - stored_events = stored_merge._stored_events(response) - segment = [ - items._user_message_item(thread_id, q, response) - for q in queries_by_turn.get(turn_index, []) +async def finish_lines( + thread_id: str, lines: list[Line], *, status_only: bool = False +) -> None: + """Stamp what only read time knows onto flagged lines, in place: a task + card's status from current liveness, a widget's offloaded data. Every + other line streams as stored. ``status_only`` leaves a failed task's + reason off, as ``stamp_task_artifact_data`` does for a share viewer.""" + flagged = [i for i, line in enumerate(lines) if line.flags] + if not flagged: + return + decoded = [lines[i].item() for i in flagged] + await widgets._resolve_widget_data_refs( + [ + item + for i, item in zip(flagged, decoded) + if replay_lines.FLAG_WIDGET in lines[i].flags ] + ) + await task_status.stamp_replay_task_status( + thread_id, decoded, status_only=status_only + ) + for i, item in zip(flagged, decoded): + lines[i] = line_of(item) - turn_items = history_events_to_sse( - messages_to_history_events(turn.messages), thread_id=thread_id - ) - # Compaction signals (offload counts, the summarize event) live in - # private state keys, not the messages channel — re-emit them at the - # head of the turn they landed in (live they fire before the first - # post-compaction model call). Fallback notices follow the same - # head placement: live they fire before the succeeding model's chunks. - turn_items[:0] = projector.context_signal_items(thread_id, turn) + projector.model_fallback_items( - thread_id, turn - ) - task_items, turn_task_ids = lane.project_for_turn( - turn, turn_index, response_id - ) - turn_items.extend(task_items) - # Legacy path→URL records belong to the turn whose state delta contains - # them. A thread-global map is incorrect when a sandbox filename is - # reused later: last-write-wins would rewrite the older turn's image to - # the newer content-addressed object. - _apply_image_url_map( - turn_items, _collect_image_url_map(turn.new_ui_records) - ) - # Table-sourced synthesis rides ahead of the merge: a turn with stored - # events drops these copies and replays the stored ones instead (the - # _STORED_PREFERRED_EVENTS transition rule). - turn_items = items._insert_provenance_items( - turn_items, provenance_by_response.get(response_id) or [] - ) - turn_items.extend( - items._interrupt_item(thread_id, intr) for intr in turn.ending_interrupts - ) - credit_item = items._credit_usage_item( - thread_id, response, usage_by_response.get(response_id) - ) - if credit_item: - turn_items.append(credit_item) - - if _has_unresolved_sandbox_images(turn_items) and stored_events: - # Non-derivable image URLs live only in the stored events for - # this turn — replay the MAIN lane from storage wholesale. - turn_items = stored_merge._replay_main_lane_from_storage( - task_items, stored_events, lane.turn_lossy_lanes - ) - else: - turn_items = stored_merge._merge_stored_payloads( - turn_items, - stored_events, - stopped.resurrect_lanes(response, lane.turn_lossy_lanes), - ) - _fill_token_thresholds(turn_items) - # Terminal error: never in stored events (persisted before it is - # yielded live), so it appends after the merge on every turn. A user - # stop's close is read off the same row (see stopped.stop_close_item). - turn_items += items.terminal_items(thread_id, response, turn_items) - - # After the stored-events merge so both projected and stored-copy - # artifacts are covered, and before caching: the watermark is a fact - # about what this build claimed, and a cached turn's claims are final - # (a turn with a still-writing run is never cached). - task_status.stamp_projected_watermarks(turn_items, lane.claimed_watermarks) - for item in turn_items: - items._enrich(item, thread_id, turn_index, response_id) - segment.extend(turn_items) - out.extend(segment) - # A still-writing subagent transcript (tail mode) must not be frozen: - # its task-ns writes never move this turn's tail, so a partial entry - # would never be invalidated. Rebuild-per-read until the task's - # stream finalizes, then the next read caches the full transcript. - # - # Same discipline while a settled lane's archive is still owed: the - # collector races the refresh-at-finalize, and it never invalidates - # this cache. A lossy lane's capture-only rows and the fuller stored - # copy behind a projected eviction pointer exist only in that - # archive — caching before it lands would freeze the loss for the - # cache TTL. A stored transcript-class row clears the debt: those - # classes are written only by the atomic archive writers (collector, - # stop drain), never by the live root path. A lossy main lane owes - # nothing: its rows land in the finalize CAS that sets the status. - awaiting_archive = set(lane.turn_lossy_lanes) - for i in turn_items: - d = i.get("data") or {} - agent = str(d.get("agent", "")) - if ( - i.get("event") == "tool_call_result" - and agent.startswith("task:") - and isinstance(d.get("content"), str) - and d["content"].startswith(stored_merge._EVICTED_RESULT_PREFIX) - ): - awaiting_archive.add(agent) - if awaiting_archive: - awaiting_archive -= { - str((e.get("data") or {}).get("agent", "")) - for e in stored_events or [] - if stored_merge._valid_stored(e) and e["event"] in stored_merge._ARCHIVE_EVIDENCE_EVENTS - } - # Gate on the PRE-read probe (lane.prepare), not a probe here: a - # stream sealing between the namespace read and a post-projection - # probe would let this build cache pre-terminal state it read - # moments earlier — frozen for the cache TTL, since task-ns writes - # never move this turn's tail. A lane the projection could not - # reconcile vetoes outright: no stored evidence retires that debt. - if ( - not awaiting_archive - and not lane.turn_uncacheable_lanes - and not (turn_task_ids & lane.live_streams_at_read) - ): - cacheable.append( - ( - turn_index, - turn.tail_checkpoint_id, - _turn_fingerprint( - response, usage_by_response, provenance_by_response - ), - segment, - ) - ) +async def project_turns( + rows: ThreadRows, + branch_tip_checkpoint_id: str, + *, + turn_indexes: list[Any] | None = None, + last_n_turns: int | None = None, + derive: Derive | None = None, + reproject: bool = False, + claims_pass: bool = True, +) -> dict[Any, list[Line]]: + """Bring the given turns' stored lines up to date (or the newest + ``last_n_turns``), so the next read streams them without projecting. + Returns the lines of every selected turn with a checkpoint boundary. + + ``reproject`` projects every selected turn, current lines or not, as + ``derive`` does, which only sees the turns it projects (see ``Derive``). + Without ``claims_pass``, turns that need the thread's run launches placed + raise ``ClaimsPassNeeded`` rather than read the whole thread for them.""" + reader = CheckpointHistoryReader.get_instance() + branch = await _branch(reader, rows, branch_tip_checkpoint_id) + if branch is None: + return {} + inputs, anchored, _tip = branch + wanted = set(turn_indexes or ()) + selected = [ + (ti, a) + for k, (ti, a) in enumerate(anchored) + if ti in wanted + or (last_n_turns is not None and k >= len(anchored) - last_n_turns) + ] + return await _turn_lines( + reader, + inputs, + anchored, + selected, + cold.Policy(claims_pass=claims_pass, derive=derive), + reproject=reproject or derive is not None, + ) - out.extend(await lane.trailing_items()) - # Trailing salvage rides items but belongs to no turn's checkpoint range, - # so the all-cache-hit fast path (which never runs the task lane) would - # silently drop it. Keep the salvage-stamped turn uncacheable — skip its - # store and evict any entry from before the orphan appeared — so every - # read misses there and rebuilds until the salvage resolves. - salvaged = lane.salvaged_turn_indexes - fingerprints: dict[Any, str] = {} - for turn_index, tail_checkpoint_id, fingerprint, segment in cacheable: - fingerprints[turn_index] = fingerprint - if turn_index not in salvaged: - await projection_cache.store_turn( - thread_id, tail_checkpoint_id, fingerprint, segment - ) - if salvaged: - tails_by_turn = {ti: t.tail_checkpoint_id for ti, t in pairs if t is not None} - # A salvaged turn is usually uncacheable for other reasons too, so it - # may never have reached `cacheable` — derive its fingerprint directly - # rather than assuming an entry was recorded there. - await projection_cache.delete_turns( - thread_id, - [ - ( - tails_by_turn[ti], - fingerprints.get(ti) - or _turn_fingerprint( - responses_by_turn.get(ti), - usage_by_response, - provenance_by_response, - ), - ) - for ti in salvaged - if tails_by_turn.get(ti) - ], - ) - for interrupt in history.interrupts: - out.append(items._interrupt_item(thread_id, interrupt)) +async def unserved_turns( + rows: ThreadRows, + branch_tip_checkpoint_id: str, + *, + reader: CheckpointHistoryReader | None = None, +) -> list[Any] | None: + """The branch's turns a read would project rather than stream from their + stored slice and lines, or None when the branch has no turns. + + Checkpoint storage maintenance reads this to know a thread's history no + longer depends on its old checkpoint states. Raises + ``CheckpointReplayUnavailable`` when the turns cannot be paired. + """ + reader = reader or CheckpointHistoryReader.get_instance() + branch = await _branch(reader, rows, branch_tip_checkpoint_id) + if branch is None: + return None + inputs, anchored, _tip = branch + _, stale = await _stored_lines(inputs, anchored) + return [ti for ti, _ in stale] - return out +# ---------------------------------------------------------------- pairing -def _pair_turns_to_queries( - turns: list[Any], - turn_indexes: list[Any], - responses_by_turn: dict[Any, dict[str, Any]], - windowed: bool, -) -> list[tuple[Any, Any]]: - """Pair checkpoint turns with persisted turn_indexes, metadata-keyed. - - A ``TurnSlice`` pairs by its stamped ``turn_index`` metadata when present, - falling back to head-anchored ordinal position (pre-stamping threads; - resume boundaries never carry metadata). Returns ordered - ``(turn_index, TurnSlice | None)`` — a ``None`` slice is a persisted turn - with no committed boundary (the in-flight active turn, or a run that never - checkpointed), replayed as its user_message stub only. In windowed mode, - unpaired rows older than the window are dropped, not stubbed. - - Raises ``CheckpointReplayUnavailable`` on anything a projection could - silently mislabel: a stamped index missing from the rows, non-monotonic - pairing, or a *completed* response with no boundary (a completed turn - always persists its boundary pointer, so checkpoints can't cover it). + +def pair_anchors( + anchors: list[TurnAnchor], turn_indexes: list[Any] +) -> list[tuple[Any, TurnAnchor]]: + """Pair checkpoint turns with persisted turn_indexes (``pair_turns``). + + Raises on anything a projection could silently mislabel: a stamped index + missing from the rows, a turn left without one, or non-monotonic pairing. """ known = set(turn_indexes) - pairs: list[tuple[Any, Any]] = [] - for turn in turns: - if turn.turn_index is not None: - ti = turn.turn_index - if ti not in known: - raise CheckpointReplayUnavailable( - f"checkpoint turn_index {ti} has no persisted turn" - ) - elif turn.turn_ordinal < len(turn_indexes): - ti = turn_indexes[turn.turn_ordinal] - else: + pairs: list[tuple[Any, TurnAnchor]] = [] + for anchor, ti in zip(anchors, pair_turns(anchors, turn_indexes)): + if ti is None: raise CheckpointReplayUnavailable( "more checkpoint turns than persisted turns" ) - pairs.append((ti, turn)) - - paired_tis = [ti for ti, _ in pairs] - if paired_tis != sorted(set(paired_tis)): - raise CheckpointReplayUnavailable("turn pairing is not monotonic") - - window_start = paired_tis[0] - paired = set(paired_tis) - for ti in turn_indexes: - if ti in paired: - continue - if windowed and ti < window_start: - continue - if (responses_by_turn.get(ti) or {}).get("status") == "completed": + if ti not in known: raise CheckpointReplayUnavailable( - f"persisted turn {ti} completed but has no checkpoint boundary" + f"checkpoint turn_index {ti} has no persisted turn" ) - pairs.append((ti, None)) + pairs.append((ti, anchor)) - pairs.sort(key=lambda p: p[0]) + paired = [ti for ti, _ in pairs] + if paired != sorted(set(paired)): + raise CheckpointReplayUnavailable("turn pairing is not monotonic") return pairs -def build_sse_replay_items( - thread_id: str, - queries: list[dict[str, Any]], - responses_by_turn: dict[Any, dict[str, Any]], -) -> list[dict[str, Any]]: - """Replay items sourced verbatim from persisted ``sse_events`` (the fallback). - - Same ``{"event", "data"}`` shape as the checkpoint path, so the endpoint - emits either source through one loop. The terminal error event is - synthesized from the response row here too — it is yielded live *after* - the persist snapshot, so stored events never contain it. So is a user - stop's close when the archive lacks one. +def _select_page( + anchored: list[tuple[Any, TurnAnchor]], + inputs: Inputs, + before_turn: Any | None, + limit: int | None, +) -> tuple[list[tuple[Any, TurnAnchor | None]], bool]: + """The page's ``(turn_index, anchor | None)`` pairs in order, and whether + older turns remain. + + A ``None`` anchor is a persisted turn with no committed boundary (the + in-flight active turn, or a run that never checkpointed), replayed as its + user_message stub. One older than the page's oldest checkpointed turn + belongs to an earlier page. A *completed* response with no boundary + raises: a completed turn always persists its boundary pointer, so + checkpoints can't cover it. """ - out: list[dict[str, Any]] = [] - terminals_emitted: set[str] = set() - for query in queries: - if not isinstance(query, dict): - continue - turn_index = query.get("turn_index") - response = responses_by_turn.get(turn_index) - response_id = ( - str(response.get("conversation_response_id")) if response else None - ) - out.append(items._user_message_item(thread_id, query, response)) - turn_items = [ - {"event": event["event"], "data": dict(event["data"])} - for event in stored_merge._stored_events(response) - if stored_merge._valid_stored(event) - ] - if response_id and response_id not in terminals_emitted: - terminals = items.terminal_items(thread_id, response, turn_items) - if terminals: - terminals_emitted.add(response_id) - turn_items += terminals - for item in turn_items: - items._enrich(item, thread_id, turn_index, response_id) - out.extend(turn_items) - return out - + in_range = [ + (ti, a) for ti, a in anchored if before_turn is None or ti < before_turn + ] + has_more = limit is not None and len(in_range) > limit + chosen: list[tuple[Any, TurnAnchor | None]] = list( + in_range[-limit:] if has_more else in_range + ) + oldest = chosen[0][0] if has_more else None -def _collect_image_url_map(records: list[dict[str, Any]]) -> dict[str, str]: - url_map: dict[str, str] = {} - for record in records: - if not isinstance(record, dict) or record.get("name") != IMAGE_CAPTURE_UI_NAME: + paired = {ti for ti, _ in anchored} + for ti in inputs.turn_indexes: + if ti in paired: + continue + if oldest is not None and ti < oldest: + continue + if before_turn is not None and ti >= before_turn: continue - path_to_url = (record.get("props") or {}).get("path_to_url") - if isinstance(path_to_url, dict): - url_map.update( - {str(k): str(v) for k, v in path_to_url.items() if k and v} + if (inputs.responses_by_turn.get(ti) or {}).get("status") == "completed": + raise CheckpointReplayUnavailable( + f"persisted turn {ti} completed but has no checkpoint boundary" ) - return url_map + chosen.append((ti, None)) + chosen.sort(key=lambda p: p[0]) + return chosen, has_more -def _apply_image_url_map( - turn_items: list[dict[str, Any]], url_map: dict[str, str] -) -> list[dict[str, Any]]: - if not url_map: - return turn_items - - def replacer(match): - alt, path = match.group(1), match.group(2) - if path in url_map: - return f"![{alt}]({url_map[path]})" - return match.group(0) +async def _branch( + reader: CheckpointHistoryReader, + rows: ThreadRows, + branch_tip_checkpoint_id: str | None, +) -> tuple[Inputs, list[tuple[Any, TurnAnchor]], str] | None: + """The rows indexed, the branch's checkpoint turns paired with them, and + the tip the branch was read from; None when it has no turns, which each + entry point answers its own way.""" + try: + anchors, tip_id = await reader.aget_turn_anchors( + rows.thread_id, branch_tip_checkpoint_id + ) + except CheckpointBranchTipNotFound as e: + raise CheckpointReplayUnavailable(str(e)) from e + if not anchors or tip_id is None: + return None + inputs = Inputs.of(rows) + return inputs, pair_anchors(anchors, inputs.turn_indexes), tip_id - for item in turn_items: - if item.get("event") != "message_chunk": - continue - data = item.get("data", {}) - if data.get("content_type") != "text": - continue - content = data.get("content") - if content: - data["content"] = IMAGE_MD_RE.sub(replacer, content) - return turn_items +# ---------------------------------------------------------------- lines -def _has_unresolved_sandbox_images(turn_items: list[dict[str, Any]]) -> bool: - for item in turn_items: - if item.get("event") != "message_chunk": - continue - data = item.get("data", {}) - if data.get("content_type") != "text": - continue - content = data.get("content") or "" - for match in IMAGE_MD_RE.finditer(content): - if is_sandbox_image_path(match.group(2)): - return True - return False +async def _turn_lines( + reader: CheckpointHistoryReader, + inputs: Inputs, + anchored: list[tuple[Any, TurnAnchor]], + selected: list[tuple[Any, TurnAnchor]], + policy: cold.Policy, + *, + reproject: bool = False, +) -> dict[Any, list[Line]]: + """Lines for the selected turns: stored ones whose key still matches + (unless ``reproject`` or the policy reads no cache), the rest projected + (and stored when final).""" + if not selected: + return {} + if reproject or not policy.cache: + out, stale = {}, selected + else: + out, stale = await _stored_lines(inputs, selected) + if stale: + out.update(await cold.project_stale(reader, inputs, anchored, stale, policy)) + return out -def _fill_token_thresholds(turn_items: list[dict[str, Any]]) -> None: - """Stamp the UI-ring threshold on projected token_usage events. - The live handler adds it server-side (config, not graph state); replay - uses the same resolver so both wires carry the same value. - """ - for item in turn_items: - data = item["data"] +async def _stored_lines( + inputs: Inputs, + selected: list[tuple[Any, TurnAnchor]], +) -> tuple[dict[Any, list[Line]], list[tuple[Any, TurnAnchor]]]: + """The selected turns' stored lines a read serves as they are, and the + stale turns it has to project.""" + response_ids = [rid for ti, _ in selected if (rid := inputs.response_id(ti))] + stored = await slices_db.get_turn_lines(response_ids) + out: dict[Any, list[Line]] = {} + stale: list[tuple[Any, TurnAnchor]] = [] + for ti, anchor in selected: + row = stored.get(inputs.response_id(ti) or "") if ( - item["event"] == "context_window" - and data.get("action") == "token_usage" - and "threshold" not in data + slices.matches(row, anchor) + and row.lines is not None + and row.lines_key + == lines_key(inputs, ti, anchor, inputs.responses_by_turn.get(ti)) ): - data["threshold"] = resolve_token_threshold() - - + try: + out[ti] = replay_lines.decode(row.lines) + continue + except Exception: + # Projected again and stored over, rather than failing every + # read of this thread on one bad row. + logger.warning( + "[REPLAY] stored lines undecodable for %s turn %s", + inputs.thread_id, + ti, + exc_info=True, + ) + stale.append((ti, anchor)) + return out, stale diff --git a/src/server/services/history/replay/cold.py b/src/server/services/history/replay/cold.py new file mode 100644 index 000000000..1c311a9d7 --- /dev/null +++ b/src/server/services/history/replay/cold.py @@ -0,0 +1,514 @@ +"""Project the turns a read cannot serve from stored lines, and store them. + +A batch at a time, each batch's stores running while the next one loads; +the whole-thread claims pass runs here, once, when a batch needs it. What a +turn projects to is ``turn``'s; this module decides only which turns are +projected together, what is written, and what a failure costs, so it is not +digested into the lines epoch. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + +from src.server.database.conversation import turn_slices as slices_db +from src.server.services.history import slices +from src.server.services.history.reader import CheckpointHistoryReader, TurnAnchor +from src.server.services.history.replay import lines as replay_lines +from src.server.services.history.replay import run_lane +from src.server.services.history.replay.errors import ( + INFRASTRUCTURE, + ClaimsNeeded, + ClaimsPassNeeded, + ClaimsUnavailable, +) +from src.server.services.history.replay.keys import lines_key +from src.server.services.history.replay.lines import Line, line_of +from src.server.services.history.replay.stored_events import ( + fallback_lines, + with_stored_events, +) +from src.server.services.history.replay.turn import ( + Derive, + Inputs, + project_turn, + settled_at, +) +from src.server.services.history.slices import TurnSlice + +logger = logging.getLogger(__name__) + +# Turns projected per batch, which bounds the checkpoint states and stored +# events in hand at once. +PROJECT_BATCH = 16 + +# Turns by index with their anchors on the branch, in turn order. +_Turns = list[tuple[Any, TurnAnchor]] + + +@dataclass(frozen=True) +class Policy: + """What an entry point lets its read do about turns it must project.""" + + # A turn that cannot be projected replays from its stored events rather + # than failing the read. + turn_fallback: bool = False + # A batch may run the whole-thread claims pass it needs, rather than + # raise ClaimsPassNeeded. + claims_pass: bool = True + # The backfill's hook on each turn that derived its legacy facts from + # its stored events. + derive: Derive | None = None + # Read stored slices and lines, and store what the read projects. Off, + # every turn and run is cut from its checkpoints and nothing is written, + # so a stored row cannot hide what the code projects today. + cache: bool = True + + +@dataclass +class _Writes: + """Rows to store, a statement per kind.""" + + turns: list[slices_db.TurnRow] = field(default_factory=list) + updates: list[slices_db.TurnUpdate] = field(default_factory=list) + runs: list[slices_db.RunRow] = field(default_factory=list) + + async def store(self, thread_id: str) -> None: + """Best effort, kind by kind: a turn that fails to store is + projected again on the next read.""" + for kind, write, rows in ( + ("run slice", slices_db.upsert_run_slices, self.runs), + ("turn slice", slices_db.upsert_turn_slices, self.turns), + ("turn lines", slices_db.update_turns, self.updates), + ): + if not rows: + continue + try: + await write(rows) + except Exception: + logger.warning( + "[REPLAY] %s store failed for %s (%d rows)", + kind, + thread_id, + len(rows), + exc_info=True, + ) + + +class _Pending: + """Stores in flight while the next batch loads. Every one is awaited + before the request returns, so a later read in it never misses its own + writes, and the post-turn refresh and backfill still finish theirs. + Without ``store`` every write is dropped here, the one place they all + pass.""" + + def __init__(self, thread_id: str, *, store: bool) -> None: + self._thread_id = thread_id + self._store = store + self._tasks: list[asyncio.Task[None]] = [] + + def start(self, writes: _Writes) -> None: + if self._store and (writes.turns or writes.updates or writes.runs): + self._tasks.append(asyncio.create_task(writes.store(self._thread_id))) + + async def drain(self) -> None: + tasks, self._tasks = self._tasks, [] + if tasks: + await asyncio.gather(*tasks) + + +@dataclass +class _Loaded: + """A stale turn's slice, read but not yet projected.""" + + turn: TurnSlice + # Cut from checkpoints by this read, so its whole row is written. + fresh: bool + # The claims stored beside the slice it was read from. + claims: list[str | None] | None + + +class _ColdRead: + """One read's stale turns, projected a batch at a time. The batches + share the run lane, the stores in flight, and the whole-thread claims + pass, which runs at most once per read.""" + + def __init__( + self, + reader: CheckpointHistoryReader, + inputs: Inputs, + anchored: _Turns, + stale: _Turns, + policy: Policy, + ) -> None: + self._reader = reader + self._inputs = inputs + self._anchored = anchored + self._stale = [ti for ti, _ in stale] + self._policy = policy + self.lane = run_lane.RunLane( + reader, inputs.thread_id, inputs.run_facts, cache=policy.cache + ) + self.pending = _Pending(inputs.thread_id, store=policy.cache) + # Stale turns projected or replayed from storage. + self._done: set[Any] = set() + # Slices read for turns not yet done: a batch the claims pass + # interrupted, and the turns the pass cut from checkpoints, so each + # is read once and written once. + self._in_hand: dict[Any, _Loaded] = {} + # Launches of the turns read so far, so the claims pass reads only + # the rest. + self._launched: dict[Any, list[run_lane.Launch]] = {} + # Run claims of the whole thread, once the pass placed them. + self._placed: dict[Any, list[str | None]] | None = None + self._placement_failed = False + + async def project(self, batch: _Turns) -> dict[Any, list[Line]]: + """Project a batch and start its stores, running the claims pass + first when the batch needs it.""" + try: + return await self._project(batch) + except ClaimsNeeded as needed: + await self._place_claims(needed) + return await self._project(batch) + + async def recover(self, batch: _Turns, error: Exception) -> dict[Any, list[Line]]: + """The batch's lines after ``error``: from storage whole when its run + claims cannot be placed, otherwise a turn at a time, and from storage + for a turn that still fails.""" + if isinstance(error, ClaimsUnavailable): + logger.warning( + "[REPLAY] run claims of %s unavailable; turns %s..%s replay " + "from their stored events", + self._inputs.thread_id, + batch[0][0], + batch[-1][0], + exc_info=error, + ) + return await self._fall_back(batch) + if len(batch) == 1: + logger.warning( + "[REPLAY] turn %s of %s replays from its stored events", + batch[0][0], + self._inputs.thread_id, + exc_info=error, + ) + return await self._fall_back(batch) + logger.warning( + "[REPLAY] batch projection failed for %s; retrying per turn", + self._inputs.thread_id, + exc_info=error, + ) + out: dict[Any, list[Line]] = {} + for turn in batch: + try: + out.update(await self.project([turn])) + except INFRASTRUCTURE: + raise + except Exception as err: + logger.warning( + "[REPLAY] turn %s of %s replays from its stored events", + turn[0], + self._inputs.thread_id, + exc_info=err, + ) + out.update(await self._fall_back([turn])) + return out + + async def close(self) -> None: + """Store the runs a failed batch queued but never took (they settled + whatever became of the turns that launched them), then wait for + every store.""" + self.pending.start(_Writes(runs=self.lane.take_writes())) + await self.pending.drain() + + async def _project(self, batch: _Turns) -> dict[Any, list[Line]]: + inputs = self._inputs + loaded = await self._load(batch) + self._launched.update( + (ti, run_lane.launches_in(loaded[ti].turn)) for ti, _ in batch + ) + claims = { + ti: loaded[ti].claims if self._placed is None else self._placed.get(ti) + for ti, _ in batch + } + await self.lane.prepare( + [ + ( + loaded[ti].turn, + claims[ti], + settled_at(inputs.responses_by_turn.get(ti)), + ) + for ti, _ in batch + ] + ) + + # Only turns projected afresh need their row's stored events (and the + # version of the row those events came from). + responses = await with_stored_events( + inputs.responses_by_turn, [ti for ti, _ in batch] + ) + writes = _Writes() + out: dict[Any, list[Line]] = {} + for ti, anchor in batch: + response = responses.get(ti) + turn = loaded[ti].turn + segment, sealed = project_turn( + inputs, ti, turn, response, self.lane, self._policy.derive + ) + out[ti] = [line_of(item) for item in segment] + self._add_turn_write( + writes, + ti, + anchor, + turn, + response, + claims[ti] if self.lane.can_store else None, + out[ti] if sealed else None, + fresh=loaded[ti].fresh, + ) + # Taken once the batch projected: runs a failed batch queued stay + # queued for the next store. + writes.runs = self.lane.take_writes() + self.pending.start(writes) + self._finish(batch) + return out + + async def _load(self, batch: _Turns) -> dict[Any, _Loaded]: + missing = [(ti, a) for ti, a in batch if ti not in self._in_hand] + turns, claims, fresh = await load_turns( + self._reader, self._inputs, missing, cache=self._policy.cache + ) + self._in_hand.update( + (ti, _Loaded(turns[ti], ti in fresh, claims[ti])) for ti, _ in missing + ) + return {ti: self._in_hand[ti] for ti, _ in batch} + + async def _fall_back(self, batch: _Turns) -> dict[Any, list[Line]]: + """The batch from its stored events, never stored, so the next read + projects it again.""" + lines = await fallback_lines(self._inputs, [ti for ti, _ in batch]) + self._finish(batch) + return lines + + def _finish(self, batch: _Turns) -> None: + for ti, _ in batch: + self._in_hand.pop(ti, None) + self._done.add(ti) + + async def _place_claims(self, needed: ClaimsNeeded) -> None: + """Run the whole-thread claims pass, and store what it settled for + the turns this read will not project.""" + if not self._policy.claims_pass: + raise ClaimsPassNeeded(str(needed)) from needed + if self._placement_failed or self._placed is not None: + # The pass reads the whole thread, and again would read it the + # same way. + raise ClaimsUnavailable(str(needed)) from needed + try: + self._placed, extracted = await _place_launches( + self._reader, + self._inputs, + self.lane, + self._anchored, + self._launched, + cache=self._policy.cache, + ) + except INFRASTRUCTURE: + raise + except Exception as err: + self._placement_failed = True + raise ClaimsUnavailable(str(err)) from err + waiting = {ti for ti in self._stale if ti not in self._done} + self._in_hand.update( + (ti, _Loaded(turn, fresh=True, claims=None)) + for ti, turn in extracted.items() + if ti in waiting + ) + # Claims for turns already projected update the rows this read + # stored, so those stores land first. + await self.pending.drain() + self.pending.start(self._claims_writes(extracted, skip=waiting)) + + def _claims_writes( + self, extracted: dict[Any, TurnSlice], *, skip: set[Any] + ) -> _Writes: + """What the claims pass leaves to store for turns this read will not + project (``skip`` holds the ones it will): a slice cut afresh with + its claims, or the claims alone beside a stored slice. Pinning the + claims is what lets the pass run once per thread.""" + writes = _Writes(runs=self.lane.take_writes()) + for ti, anchor in self._anchored: + if ti in skip: + continue + claims = (self._placed or {}).get(ti) if self.lane.can_store else None + if ti not in extracted and claims is None: + continue + self._add_turn_write( + writes, + ti, + anchor, + extracted.get(ti), + self._inputs.responses_by_turn.get(ti), + claims, + None, + fresh=ti in extracted, + ) + return writes + + def _add_turn_write( + self, + writes: _Writes, + turn_index: Any, + anchor: TurnAnchor, + turn: TurnSlice | None, + response: dict[str, Any] | None, + claims: list[str | None] | None, + lines: list[Line] | None, + *, + fresh: bool, + ) -> None: + """Queue a turn's one write: its whole row when its slice was cut + afresh, otherwise whatever this read learned beside the stored slice. + Lines key off the row the projection read, which may be newer than + the one the page was keyed by.""" + if response is None or anchor.tail_checkpoint_id is None: + return + inputs = self._inputs + key = ( + lines_key(inputs, turn_index, anchor, response) + if lines is not None + else None + ) + text = replay_lines.encode(lines) if lines is not None else None + response_id = str(response["conversation_response_id"]) + if fresh and turn is not None: + codec, data = slices.encode(self._reader.serde, turn) + writes.turns.append( + slices_db.TurnRow( + response_id, + inputs.thread_id, + slices_db.StoredSlice( + input_checkpoint_id=anchor.input_checkpoint_id, + tail_checkpoint_id=anchor.tail_checkpoint_id, + slice_key=slices.SLICE_KEY, + codec=codec, + data=data, + lines_key=key, + lines=text, + claims=claims, + ), + ) + ) + elif text is not None or claims is not None: + writes.updates.append( + slices_db.TurnUpdate( + response_id, anchor.tail_checkpoint_id, key, text, claims + ) + ) + + +async def project_stale( + reader: CheckpointHistoryReader, + inputs: Inputs, + anchored: _Turns, + stale: _Turns, + policy: Policy, +) -> dict[Any, list[Line]]: + """Project and store stale turns a batch at a time: each checkpoint + state holds the thread's whole transcript up to its turn, and each + response row its stored events, so a thread read at once would hold + both for every turn. A batch's stores run while the next batch loads. + + With ``policy.turn_fallback``, a batch that fails is retried a turn at a + time, and a turn that still fails replays from its stored events (or as + a stub), never stored, so it is projected again on the next read. Only a + turn's own failure is retried: a batch whose run claims cannot be placed + replays from storage whole, and a database fault ends the read. + """ + read = _ColdRead(reader, inputs, anchored, stale, policy) + out: dict[Any, list[Line]] = {} + try: + for start in range(0, len(stale), PROJECT_BATCH): + batch = stale[start : start + PROJECT_BATCH] + try: + out.update(await read.project(batch)) + except INFRASTRUCTURE: + raise + except Exception as e: + if not policy.turn_fallback: + raise + out.update(await read.recover(batch, e)) + finally: + await read.close() + return out + + +async def load_turns( + reader: CheckpointHistoryReader, + inputs: Inputs, + batch: list[tuple[Any, TurnAnchor]], + *, + cache: bool = True, +) -> tuple[dict[Any, TurnSlice], dict[Any, list[str | None] | None], set[Any]]: + """``(turns, claims, fresh)``: the batch's turns by ``slices.load_turns``, + every one cut from checkpoints without ``cache``. Claims count only + beside the slice they were placed against.""" + stored = ( + await slices_db.get_turn_slices( + [rid for ti, _ in batch if (rid := inputs.response_id(ti))] + ) + if cache + else {} + ) + rows = { + ti: row for ti, _ in batch if (row := stored.get(inputs.response_id(ti) or "")) + } + turns, fresh = await slices.load_turns(reader, inputs.thread_id, batch, rows) + claims = {ti: None if ti in fresh else rows[ti].claims for ti, _ in batch} + return turns, claims, fresh + + +async def _place_launches( + reader: CheckpointHistoryReader, + inputs: Inputs, + lane: run_lane.RunLane, + anchored: list[tuple[Any, TurnAnchor]], + known: dict[Any, list[run_lane.Launch]], + *, + cache: bool, +) -> tuple[dict[Any, list[str | None]], dict[Any, TurnSlice]]: + """``(claims, extracted)``: every run launch of the thread claimed in + order (``assign_claims``), and the turns this pass had to cut from + checkpoints. + + Every turn's launches are needed, so the turns whose launches are not + ``known`` are read a batch at a time, stored slices first. Only their + launches are kept, except for the turns cut afresh: the caller writes + each of those once, with its claims. + """ + launches: list[tuple[list[run_lane.Launch], datetime | None]] = [] + extracted: dict[Any, TurnSlice] = {} + for start in range(0, len(anchored), PROJECT_BATCH): + batch = anchored[start : start + PROJECT_BATCH] + turns, _claims, fresh = await load_turns( + reader, + inputs, + [(ti, a) for ti, a in batch if ti not in known], + cache=cache, + ) + extracted.update((ti, turns[ti]) for ti in fresh) + for ti, _anchor in batch: + launches.append( + ( + known[ti] if ti in known else run_lane.launches_in(turns[ti]), + settled_at(inputs.responses_by_turn.get(ti)), + ) + ) + + placed = await lane.assign_claims(launches) + return {ti: c for (ti, _), c in zip(anchored, placed) if c}, extracted diff --git a/src/server/services/history/replay/errors.py b/src/server/services/history/replay/errors.py new file mode 100644 index 000000000..4db52047e --- /dev/null +++ b/src/server/services/history/replay/errors.py @@ -0,0 +1,29 @@ +"""Why a replay read cannot project, by what the caller does about it.""" + +from __future__ import annotations + +import psycopg + + +class CheckpointReplayUnavailable(Exception): + """Checkpoint history cannot faithfully cover this thread's replay.""" + + +class ClaimsPassNeeded(CheckpointReplayUnavailable): + """The turns can be projected only after a pass over the whole thread + places its run launches, which this caller does not wait for.""" + + +class ClaimsUnavailable(CheckpointReplayUnavailable): + """The whole-thread claims pass failed, or ran and still left a launch + unplaced: no turn of the read can be told apart by retrying it alone.""" + + +class ClaimsNeeded(Exception): + """A launch without a usable ledger stamp, in a turn whose claims were + never assigned: only a pass over the whole thread can place it.""" + + +# A database that refuses connections or queries fails every turn alike, so +# the read gives up rather than wait on it once per turn. +INFRASTRUCTURE = (psycopg.OperationalError, psycopg.InterfaceError) diff --git a/src/server/services/history/replay/facts.py b/src/server/services/history/replay/facts.py new file mode 100644 index 000000000..92115fce9 --- /dev/null +++ b/src/server/services/history/replay/facts.py @@ -0,0 +1,260 @@ +"""Replay facts: what a settled run shows on replay that its checkpoint does not keep. + +A run's row carries them in ``replay_facts``: how long the main agent thought +over each message it committed (``reasoning_ms``, on the turn's response row), +the URL each sandbox image a subagent run referenced was captured to +(``images``), and the steering a subagent run handed back unread +(``steering_returned``), both on the run's ``subagent_runs`` row. A row +settled before the column existed has none, and replay keeps reading what it +needs from the turn's stored events. + +Run facts are learned after the run settles, so they can change a turn whose +lines are already stored: they are read with the thread's other replay rows, +and the lines key of the turn that launched each run covers them. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Required, TypedDict + +from ptc_agent.agent.middleware.image_capture import IMAGE_MD_RE +from src.server.services.history.replay.lanes import MAIN_LANE, agent_lane + + +class ResponseFacts(TypedDict, total=False): + """A response row's ``replay_facts``. A backfilled row also holds its + legacy facts, under ``LEGACY_KEY`` (``legacy.LegacyFacts``).""" + + # Message id -> the ms of each thinking block the live stream timed. + reasoning_ms: dict[str, list[int]] + + +class ReturnedSteering(TypedDict, total=False): + content: Required[Any] + input_id: Required[str] + # ``run_ended``, or ``run_mismatch`` for steering meant for another run. + reason: Required[str] + # The newest message the run held when it returned a mismatch. + after: str | None + + +class SubagentRunFacts(TypedDict, total=False): + """A subagent run row's ``replay_facts``.""" + + # Sandbox image path -> the URL it was captured to. + images: dict[str, str] + steering_returned: list[ReturnedSteering] + + +@dataclass +class RunFacts: + """The thread's run facts, by run and by the turn that launched each.""" + + by_run: dict[str, SubagentRunFacts] = field(default_factory=dict) + by_turn: dict[Any, list[tuple[str, SubagentRunFacts]]] = field(default_factory=dict) + + @classmethod + def of(cls, rows: list[dict[str, Any]]) -> RunFacts: + out = cls() + for row in rows: + facts = row.get("replay_facts") + if not isinstance(facts, dict) or not row.get("task_run_id"): + continue + run_id = str(row["task_run_id"]) + out.by_run[run_id] = facts + out.by_turn.setdefault(row.get("turn_index"), []).append((run_id, facts)) + for runs in out.by_turn.values(): + runs.sort(key=lambda pair: pair[0]) + return out + + def for_run(self, run_id: str | None) -> SubagentRunFacts | None: + return self.by_run.get(str(run_id)) if run_id else None + + def for_turn(self, turn_index: Any) -> list[tuple[str, SubagentRunFacts]]: + return self.by_turn.get(turn_index, []) + + +# --------------------------------------------------------------- reasoning + + +def reasoning_durations( + response: dict[str, Any] | None, +) -> dict[str, list[int]] | None: + """The response's per-message reasoning durations, or None when the row + predates them (replay then carries them from the stored events).""" + facts: ResponseFacts | None = (response or {}).get("replay_facts") + if not isinstance(facts, dict): + return None + durations = facts.get("reasoning_ms") + return durations if isinstance(durations, dict) else None + + +def apply_reasoning_durations( + turn_items: list[dict[str, Any]], durations: dict[str, list[int]] +) -> None: + """Stamp each main-lane message's total thinking time on its reasoning + close, in place. + + The projector joins a message's thinking blocks into one row with one + close, so that close stands for every block the live stream timed, and + reports their sum. + """ + closes: dict[str, dict[str, Any]] = {} + for item in turn_items: + data = item.get("data") or {} + if ( + item.get("event") == "message_chunk" + and data.get("content_type") == "reasoning_signal" + and agent_lane(data.get("agent")) == MAIN_LANE + and data.get("id") + ): + closes[data["id"]] = data + for message_id, data in closes.items(): + spans = durations.get(message_id) + if data.get("content") != "complete" or not isinstance(spans, list): + continue + timed = [ms for ms in spans if isinstance(ms, int) and not isinstance(ms, bool)] + if timed: + data["elapsed_ms"] = sum(timed) + + +# --------------------------------------------------------------- run lanes + + +def apply_image_url_map( + turn_items: list[dict[str, Any]], url_map: dict[str, str] +) -> list[dict[str, Any]]: + """Resolve sandbox image paths in text chunks to their captured URLs.""" + if not url_map: + return turn_items + + def replacer(match): + alt, path = match.group(1), match.group(2) + if path in url_map: + return f"![{alt}]({url_map[path]})" + return match.group(0) + + for item in turn_items: + if item.get("event") != "message_chunk": + continue + data = item.get("data", {}) + if data.get("content_type") != "text": + continue + content = data.get("content") + if content: + data["content"] = IMAGE_MD_RE.sub(replacer, content) + return turn_items + + +def apply_run_facts( + run_items: list[dict[str, Any]], + facts: SubagentRunFacts | None, + *, + agent: str, + thread_id: str, + message_ids: list[str | None], +) -> list[dict[str, Any]]: + """A projected run's items with its facts applied: captured images + resolved, and returned steering placed where the run returned it. + + ``message_ids`` is the run's transcript in order, which places a return + made before a model call after the message the run had reached. + """ + if not facts: + return run_items + images = facts.get("images") + if isinstance(images, dict): + apply_image_url_map( + run_items, {str(k): str(v) for k, v in images.items() if k and v} + ) + returned = facts.get("steering_returned") + if isinstance(returned, list) and returned: + run_items = _place_returned_steering( + run_items, returned, agent=agent, thread_id=thread_id, message_ids=message_ids + ) + return run_items + + +def _returned_item( + entry: ReturnedSteering, *, agent: str, thread_id: str +) -> dict[str, Any]: + return { + "event": "steering_returned", + "data": { + "agent": agent, + "content": entry.get("content"), + "input_id": entry.get("input_id"), + "reason": entry.get("reason"), + "thread_id": thread_id, + }, + } + + +def _place_returned_steering( + run_items: list[dict[str, Any]], + returned: list[Any], + *, + agent: str, + thread_id: str, + message_ids: list[str | None], +) -> list[dict[str, Any]]: + """Insert returned steering into a run's items. + + A return at run end follows everything the run wrote. A return before a + model call (``after``: the newest message the run then held) follows the + items of that message, or of the nearest message before it that + projects any; with no message yet it leads the run. + """ + last_item: dict[str, int] = {} + for idx, item in enumerate(run_items): + message_id = (item.get("data") or {}).get("id") + if message_id: + last_item[message_id] = idx + position = {mid: k for k, mid in enumerate(message_ids) if mid} + + def anchor(after: Any) -> int: + if after is None: + return -1 + k = position.get(after) + if k is None: + return len(run_items) - 1 + for mid in reversed(message_ids[: k + 1]): + if mid in last_item: + return last_item[mid] + return -1 + + inserts: dict[int, list[dict[str, Any]]] = {} + tail: list[dict[str, Any]] = [] + seen: set[tuple[Any, Any]] = set() + for entry in returned: + if not isinstance(entry, dict): + continue + key = (entry.get("input_id"), entry.get("reason")) + if entry.get("input_id"): + if key in seen: + continue + seen.add(key) + item = _returned_item(entry, agent=agent, thread_id=thread_id) + if entry.get("reason") == "run_mismatch": + inserts.setdefault(anchor(entry.get("after")), []).append(item) + else: + tail.append(item) + + merged = list(inserts.get(-1, ())) + for idx, item in enumerate(run_items): + merged.append(item) + merged.extend(inserts.get(idx, ())) + merged.extend(tail) + return merged + + +def returned_keys(turn_items: list[dict[str, Any]]) -> set[tuple[Any, Any]]: + """(agent, input_id) of every returned-steering item a projection placed, + whose stored copies a merge must then skip.""" + return { + (data.get("agent"), data.get("input_id")) + for item in turn_items + if item.get("event") == "steering_returned" + and (data := item.get("data") or {}).get("input_id") + } diff --git a/src/server/services/history/replay/items.py b/src/server/services/history/replay/items.py index 4277a3d23..d10f6cfca 100644 --- a/src/server/services/history/replay/items.py +++ b/src/server/services/history/replay/items.py @@ -5,10 +5,10 @@ from datetime import timedelta from typing import Any -from src.server.database.runs import lifecycle as tl_db +from src.server.contracts.status import TERMINAL_STATUSES from src.server.database.provenance import provenance_row_to_event -from src.server.services.history.replay import stopped, stored_merge -from src.server.services.runs.sse_producer import build_credit_usage_data +from src.server.services.history.replay import stopped +from src.server.services.runs.credit_usage import build_credit_usage_data from src.server.utils.error_sanitization import ( sanitize_error_text as _sanitize_error_text, ) @@ -35,7 +35,7 @@ def _user_message_item( # carry it: its content hasn't rendered, and marking it would suppress the # attach that streams it — hence the positive terminal check (an unknown # or legacy status must not stamp). - if response is not None and response.get("status") in tl_db.TERMINAL_STATUSES: + if response is not None and response.get("status") in TERMINAL_STATUSES: payload["run_id"] = str(response.get("conversation_response_id")) completed_at = run_completed_at(response) if completed_at is not None: @@ -56,7 +56,7 @@ def run_completed_at(response: dict[str, Any] | None) -> str | None: ``user_message`` payload and needs the same answer. Two copies of the fallback chain would drift the moment one of them learned a new column. """ - if response is None or response.get("status") not in tl_db.TERMINAL_STATUSES: + if response is None or response.get("status") not in TERMINAL_STATUSES: return None settled = response.get("usage_settled_at") if settled is not None: @@ -260,11 +260,9 @@ def _stub_turn_items( never checkpointed. The user_message stub, plus the terminal error or stop close the response row records, is the whole replay. Never cached. - A stopped or failed turn replays its stored rows as a lane that projected - nothing, since nothing the user watched stream is on the committed - branch. A stop during bring-up has no such rows; a turn that streamed - lands here when the finalize's tip read failed and left the commit - pointer behind its boundary. + Output the turn streamed before it stopped or failed is not on the + committed branch, so it is not in the model's context either, and replay + leaves it out. """ response = responses_by_turn.get(turn_index) response_id = str(response.get("conversation_response_id")) if response else None @@ -272,13 +270,7 @@ def _stub_turn_items( _user_message_item(thread_id, q, response) for q in queries_by_turn.get(turn_index, []) ] - turn_items: list[dict[str, Any]] = [] - resurrect_lanes = stopped.resurrect_lanes(response, frozenset()) - if resurrect_lanes: - turn_items = stored_merge._merge_stored_payloads( - [], stored_merge._stored_events(response), resurrect_lanes - ) - turn_items += terminal_items(thread_id, response, turn_items) + turn_items = terminal_items(thread_id, response, []) for item in turn_items: _enrich(item, thread_id, turn_index, response_id) return items + turn_items diff --git a/src/server/services/history/replay/keys.py b/src/server/services/history/replay/keys.py new file mode 100644 index 000000000..76f498c8d --- /dev/null +++ b/src/server/services/history/replay/keys.py @@ -0,0 +1,138 @@ +"""Whether a turn's stored lines still stand for it: the lines key. + +A key is the deploy's ``lines_epoch()``, a dot, and a digest of the rows the +turn was projected from. Nothing here decides what a line holds, so this +module is not among the ``projection_sources()`` it digests. +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from src.server.services.history import slices +from src.server.services.runs.sse_producer import resolve_token_threshold + +if TYPE_CHECKING: + from src.server.services.history.reader import TurnAnchor + from src.server.services.history.replay.turn import Inputs + +# Bump to re-project every stored turn when its output changes for a reason +# the digested sources below do not show (a ptc_agent helper's output). +LINES_VERSION = 1 + +_HERE = Path(__file__).resolve().parent +_SRC = _HERE.parents[3] + +# The package modules that decide what a turn's lines hold. The others decide +# only whether and when lines are stored and read: the root (pairing, paging, +# serving stored lines), cold (batching and stores), keys, errors, refresh, +# and stored_events (the fallback, whose lines are never stored). +_PROJECTION_MODULES = ( + "turn", + "items", + "facts", + "legacy", + "stored_merge", + "run_lane", + "lines", + "widgets", + "stopped", + "lanes", +) + + +def projection_sources() -> list[Path]: + """The code that turns slices into lines. A test holds this list closed + over its ``src`` imports, so a helper a projection module starts calling + lands here or in that test's reasoned exclusions.""" + return sorted( + [ + *(_HERE / f"{name}.py" for name in _PROJECTION_MODULES), + _HERE.parent / "projector.py", + _HERE.parent / "task_status.py", + # Helpers that shape the projected text and payloads. The + # sanitizer is here so a scrubbing fix reaches lines stored before + # it, the credit payload so a field it hides stays hidden. + _SRC / "server/utils/error_sanitization.py", + _SRC / "server/utils/content_normalizer.py", + _SRC / "server/utils/text_phase.py", + _SRC / "server/database/provenance.py", + _SRC / "server/services/runs/credit_usage.py", + _SRC / "server/contracts/status.py", + _SRC / "llms/content_utils.py", + _SRC / "llms/token_counter.py", + ] + ) + + +def _projection_source_digest() -> str: + """Digest of the projection's code: stored lines carry no expiry, so a + deploy that changes the projection must not keep serving what the old + code projected.""" + digest = hashlib.blake2b(digest_size=8) + for path in projection_sources(): + digest.update(str(path.relative_to(_SRC)).encode()) + digest.update(path.read_bytes()) + return digest.hexdigest() + + +_PROJECTION_VERSION = ( + f"{LINES_VERSION}.{slices.SLICE_KEY}.{_projection_source_digest()}" +) + + +def lines_epoch() -> str: + """The part of every lines key a deploy decides: the projection code and + the config it reads. Every key starts with it and a dot, so the lines a + deploy made stale are found in SQL without recomputing each turn's key.""" + digest = hashlib.blake2b(digest_size=8) + digest.update(f"{_PROJECTION_VERSION}\x1e{resolve_token_threshold()}".encode()) + return digest.hexdigest() + + +def _row_generation(row: dict[str, Any] | None, id_column: str) -> Any: + """A row reduced to its identity plus its Postgres row version: any + rewrite lands a new tuple with a new ``xmin``, so the version stands in + for every column without serializing megabyte-scale ones. Rows read + without ``xmin`` hash whole.""" + if not row or "xmin" not in row: + return row + return {id_column: row.get(id_column), "xmin": row["xmin"]} + + +def lines_key( + inputs: Inputs, + turn_index: Any, + anchor: TurnAnchor, + response: dict[str, Any] | None, +) -> str: + """The deploy's ``lines_epoch()``, then a digest of every input a turn's + lines depend on beyond its slice and its runs' slices (which only store + once final). ``response`` is the row the lines were projected from, which + may be newer than the one the read started with. The facts of the runs + it launched are in it: they are recorded after a run settles. + + Each source is hashed whole rather than a chosen subset of fields: an + over-inclusive key only costs an extra projection when something + irrelevant changes, while an under-inclusive one serves stale lines + forever. + """ + response_id = str(response["conversation_response_id"]) if response else None + digest = hashlib.blake2b(digest_size=16) + for source in ( + _row_generation(response, "conversation_response_id"), + [ + _row_generation(q, "conversation_query_id") + for q in inputs.queries_by_turn.get(turn_index, []) + ], + inputs.provenance_by_response.get(response_id) or [], + inputs.usage_by_response.get(response_id), + anchor.ending_interrupts, + inputs.run_facts.for_turn(turn_index), + ): + digest.update(json.dumps(source, default=str, sort_keys=True).encode()) + digest.update(b"\x1e") + return f"{lines_epoch()}.{digest.hexdigest()}" diff --git a/src/server/services/history/replay/legacy.py b/src/server/services/history/replay/legacy.py new file mode 100644 index 000000000..5f07184bb --- /dev/null +++ b/src/server/services/history/replay/legacy.py @@ -0,0 +1,376 @@ +"""Legacy facts: what a turn's stored events add to its projection, resolved. + +A turn settled while the live stream was archived replays some rows only that +archive kept: widget data files, a tool result the checkpoint holds as an +eviction pointer, reasoning durations, signal rows (provenance, context +window, credit usage, steering, interrupts, errors) in their streamed +positions, and the URLs its sandbox images were captured to. +``stored_merge.derive_turn`` resolves those against the turn's projection +into facts, and ``apply_turn`` places them. + +A row not yet backfilled derives its facts from its stored events on every +projection. A backfilled row carries them in ``replay_facts.legacy``, so its +replay never reads the stored events. A row with no stored events is +backfilled with ``{"v": FACTS_VERSION}`` alone: facts known empty, which is +not the same as no facts. + +Backfilled facts outlive the projection code they were resolved against, so +they name projected rows by what the checkpoint gives them, never by where +they fall: a message by its id, a tool call or result by its call id, an +artifact by the id the projector gave it (see ``row_keys``). Each signal row +keeps every anchor that streamed since the previous one, so when a later +projection drops the newest it follows the nearest that remains. A stored +widget payload upgrades the projected widget it was paired with, by id; one +whose widget the projection did not emit then takes a widget no payload +names that falls between its neighbours' widgets, or is placed like a signal +row. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Required, TypedDict + +from src.server.database.replay_facts import LEGACY_KEY +from src.server.services.history.replay import facts as replay_facts +from src.server.services.history.replay import widgets +from src.server.services.history.replay.lanes import agent_lane + +FACTS_VERSION = 3 + + +class StoredRow(TypedDict): + event: str + data: dict[str, Any] + + +class StoredResult(TypedDict, total=False): + content: Required[str] + content_type: str + + +class EventSignal(TypedDict, total=False): + event: Required[str] + data: Required[dict[str, Any]] + at: int + + +class WidgetSignal(TypedDict, total=False): + widget: Required[int] + at: int + + +Signal = EventSignal | WidgetSignal + + +class Merge(TypedDict, total=False): + """The parts ``stored_merge.derive_merge`` documents. Row keys are + stored as lists: they round-trip through JSONB.""" + + widgets: list[list[Any]] + results: dict[str, StoredResult] + durations: list[list[Any]] + anchors: list[list[Any]] + signals: list[Signal] + + +class LegacyFacts(TypedDict, total=False): + """A response row's ``replay_facts[LEGACY_KEY]``.""" + + v: Required[int] + # Sandbox image path -> the URL its stored copy was captured to. + images: dict[str, str] + merge: Merge + verbatim: list[StoredRow] + + +# Stored events replayed verbatim (anchored to their original position): +# non-derivable payloads, plus resolved-interrupt cards and error markers — +# `interrupt` renders answered HITL cards on replay (a pending interrupt is +# also re-emitted from the checkpoint tip; the frontend dedups by +# interrupt_id), and `error` keeps wire parity with sse replay. +# `context_window` and `steering_delivered` are checkpoint-projected, but a +# turn with stored events prefers those verbatim (see STORED_PREFERRED_EVENTS). +PASSTHROUGH_EVENTS = ( + "context_window", + "provenance", + "steering_delivered", + # A run's replay facts place it on its run lane; a stored copy is kept + # only for a run settled before those existed. + "steering_returned", + "credit_usage", + "interrupt", + "error", + "model_fallback", +) + + +# Projected/synthesized event types the stored stream also carries in full. +# A turn with stored events drops its projected copies and replays the stored +# ones (richer historical payloads, proven anchoring); a turn without them +# keeps the projected copies. The terminal ``error`` event is NOT here — +# stored events never contain it (persisted before it is yielded), so it is +# synthesized from the response row on both paths unconditionally. +STORED_PREFERRED_EVENTS = ( + "steering_delivered", + "context_window", + "provenance", + "credit_usage", + "interrupt", + "model_fallback", +) + + +# Content types the projector emits: a message's anchorable chunks. Live-only +# accumulation chunks (content_type=None) and tool-only messages have none, so +# both streams enumerate the same messages when keyed on these. +ANCHORABLE_CONTENT_TYPES = frozenset({"reasoning_signal", "reasoning", "text"}) + + +def facts_of(response: dict[str, Any] | None) -> LegacyFacts | None: + """The row's legacy facts, or None when it has not been backfilled (or + was backfilled by a version this code does not read).""" + facts = (response or {}).get("replay_facts") + legacy = facts.get(LEGACY_KEY) if isinstance(facts, dict) else None + if isinstance(legacy, dict) and legacy.get("v") == FACTS_VERSION: + return legacy + return None + + +# How a row key names an anchorable message chunk, from its lane, its data +# and how often its id recurred in that lane after another row (``row_keys``). +MessageKey = Callable[[str, dict[str, Any], int], tuple | None] + + +def own_message_key(lane: str, data: dict[str, Any], occurrence: int) -> tuple: + """A message chunk by its id and content type, the key facts name it by. + The projector gives every message without an id the same placeholder, so + a recurring id also carries its occurrence: a message's own chunks are + contiguous.""" + key = ("message_chunk", lane, data.get("id"), data["content_type"]) + return key + (occurrence,) if occurrence else key + + +def row_keys( + turn_items: list[dict[str, Any]], message_key: MessageKey = own_message_key +) -> list[tuple | None]: + """The key of each row, or None for a row no key names: a tool call + round by its first call id, a result by its call id, an artifact by its + id, and a message chunk by ``message_key``.""" + keys: list[tuple | None] = [] + running: dict[str, Any] = {} + occurrences: dict[tuple[str, Any], int] = {} + for item in turn_items: + event, data = item["event"], item["data"] + lane = agent_lane(data.get("agent")) + key: tuple | None = None + if event == "message_chunk" and data.get("content_type") in ANCHORABLE_CONTENT_TYPES: + message_id = data.get("id") + if lane not in running or running[lane] != message_id: + running[lane] = message_id + occurrences[lane, message_id] = occurrences.get((lane, message_id), -1) + 1 + key = message_key(lane, data, occurrences[lane, message_id]) + else: + running.pop(lane, None) + if event == "tool_calls": + call_ids = _call_ids(data) + key = ("tool_calls", call_ids[0]) if call_ids else None + elif event == "tool_call_result" and data.get("tool_call_id"): + key = ("tool_call_result", data["tool_call_id"]) + elif event == "artifact" and data.get("artifact_id"): + key = ("artifact", data["artifact_id"]) + keys.append(key) + return keys + + +def key_index( + turn_items: list[dict[str, Any]], message_key: MessageKey = own_message_key +) -> dict[tuple, int]: + """Projected position of each row key. + + The last occurrence wins, so a message key covers its whole group (a + reasoning start and close share one key, and the close is the one a + duration lands on). The checkpoint holds a parallel round as one row, + the live archive as one row per call, so any of its calls names it. + """ + index: dict[tuple, int] = {} + for idx, key in enumerate(row_keys(turn_items, message_key)): + if key is None: + continue + index[key] = idx + if key[0] == "tool_calls": + for call_id in _call_ids(turn_items[idx]["data"]): + index["tool_calls", call_id] = idx + return index + + +def _call_ids(data: dict[str, Any]) -> list[str]: + return [tc["id"] for tc in data.get("tool_calls") or [] if tc.get("id")] + + +def apply_turn( + turn_items: list[dict[str, Any]], + task_items: list[dict[str, Any]], + facts: LegacyFacts, +) -> list[dict[str, Any]]: + """The turn's items with its legacy facts applied. + + A turn whose stored images no basename could settle carries its main + lane (and every row outside the lanes its runs projected) ``verbatim``; + its run lanes keep the projection, merged like any other turn's. + """ + replay_facts.apply_image_url_map(turn_items, facts.get("images") or {}) + merge = facts.get("merge") + if "verbatim" in facts: + rows = [ + {"event": row["event"], "data": dict(row["data"])} + for row in facts["verbatim"] + ] + return rows + apply(task_items, merge) + return apply(turn_items, merge) + + +def apply( + turn_items: list[dict[str, Any]], merge: Merge | None +) -> list[dict[str, Any]]: + """Place a turn's merge facts into its projection. + + ``merge`` is None for a turn with no stored events, whose projected + signal rows stand. Otherwise those are replaced by the stored ones, in + this order: widget payloads upgrade the widgets they were paired with, + evicted results are restored, reasoning durations land on their closes, + and each signal row follows the newest of its anchors that the + projection still holds. Returned steering a run's own facts already + placed is not placed twice. + """ + if merge is None: + return turn_items + placed_returns = replay_facts.returned_keys(turn_items) + + turn_items = [i for i in turn_items if i["event"] not in STORED_PREFERRED_EVENTS] + # Keyed before the widget upgrade: facts name a widget by its projected id. + index = key_index(turn_items) + + stored_widgets = merge.get("widgets") or [] + upgraded = _upgrade_widgets(turn_items, stored_widgets) + for k, idx in upgraded.items(): + index["widget", k] = idx + + results = merge.get("results") or {} + if results: + for item in turn_items: + data = item["data"] + if item["event"] != "tool_call_result": + continue + stored = results.get(data.get("tool_call_id")) + if stored is None: + continue + data["content"] = stored["content"] + if "content_type" in stored: + data["content_type"] = stored["content_type"] + + for key, elapsed_ms in merge.get("durations") or (): + idx = index.get(tuple(key)) + if idx is None: + continue + target = turn_items[idx]["data"] + if ( + target.get("content_type") == "reasoning_signal" + and target.get("content") == "complete" + ): + target["elapsed_ms"] = elapsed_ms + + last_lane_index: dict[str, int] = {} + for idx, item in enumerate(turn_items): + last_lane_index[agent_lane(item["data"].get("agent"))] = idx + # A main lane stopped inside its first model call projects nothing. + turn_end = len(turn_items) - 1 + + anchors = [tuple(key) for key in merge.get("anchors") or ()] + inserts_after: dict[int, list[dict[str, Any]]] = {} + anchor_idx = -1 # before the first projected item + previous_at = -1 + for signal in merge.get("signals") or (): + at = signal.get("at") + if at is not None: + # The anchors streamed since the previous signal, newest first. + # When none remains the signal follows the previous one. + for i in range(at, min(at, previous_at + 1) - 1, -1): + if (idx := index.get(anchors[i])) is not None: + anchor_idx = idx + break + previous_at = at + if "widget" in signal: + # A stored widget no projected widget took is placed like a + # signal row. + k = signal["widget"] + if k in upgraded or not 0 <= k < len(stored_widgets): + continue + row = {"event": "artifact", "data": dict(stored_widgets[k][1])} + else: + event, data = signal["event"], signal["data"] + if event == "error": + # A run's error is the last thing it streamed, so it follows + # everything its lane committed even when what streamed just + # before it (output no checkpoint kept) anchors nothing. + anchor_idx = max( + anchor_idx, + last_lane_index.get(agent_lane(data.get("agent")), turn_end), + ) + if ( + event == "steering_returned" + and (data.get("agent"), data.get("input_id")) in placed_returns + ): + continue + row = {"event": event, "data": dict(data)} + inserts_after.setdefault(anchor_idx, []).append(row) + + merged = list(inserts_after.get(-1, [])) + for idx, item in enumerate(turn_items): + merged.append(item) + merged.extend(inserts_after.get(idx, ())) + return merged + + +def _upgrade_widgets( + turn_items: list[dict[str, Any]], stored_widgets: list[list[Any]] +) -> dict[int, int]: + """Give each projected widget its stored payload; the position of each + payload taken. + + A payload goes to the projected widget it names. One that named none + (the projection did not emit its widget when the facts were resolved) + takes the first widget no payload names between the widgets its stored + neighbours took: both streams order widgets by tool execution, and the + live artifact id is random. + """ + projected = [ + idx + for idx, item in enumerate(turn_items) + if widgets._is_widget(item["event"], item["data"]) + ] + named = { + artifact_id: k + for k, (artifact_id, _data) in enumerate(stored_widgets) + if artifact_id is not None + } + taken: dict[int, int] = {} # stored payload -> projected position + unnamed: list[int] = [] + for idx in projected: + k = named.get(turn_items[idx]["data"].get("artifact_id")) + if k is None: + unnamed.append(idx) + elif k not in taken: + taken[k] = idx + for k, (artifact_id, _data) in enumerate(stored_widgets): + if artifact_id is not None: + continue + after = max((idx for j, idx in taken.items() if j < k), default=-1) + before = min((idx for j, idx in taken.items() if j > k), default=len(turn_items)) + idx = next((idx for idx in unnamed if after < idx < before), None) + if idx is not None: + unnamed.remove(idx) + taken[k] = idx + for k, idx in taken.items(): + turn_items[idx]["data"] = dict(stored_widgets[k][1]) + return taken diff --git a/src/server/services/history/replay/lines.py b/src/server/services/history/replay/lines.py new file mode 100644 index 000000000..5cbc9c0e4 --- /dev/null +++ b/src/server/services/history/replay/lines.py @@ -0,0 +1,65 @@ +"""Wire-ready replay lines: a turn's projection stored as the text it streams. + +Each line is ``event \\t flags \\t data-json``, where the JSON is exactly what +the SSE frame carries. Most lines stream without being decoded; ``flags`` +marks the few the endpoint must still stamp at read time. JSON escapes every +control character, so a line never contains a raw newline or tab. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +from src.server.services.history.replay.widgets import _is_widget +from src.server.services.history.task_status import _artifact_task_id + +# A task card: its status is stamped from liveness at read. +FLAG_TASK = "t" +# A widget whose data lives in object storage, inlined at read. +FLAG_WIDGET = "w" + + +@dataclass(slots=True) +class Line: + event: str + flags: str + data_json: str + + def item(self) -> dict[str, Any]: + return {"event": self.event, "data": json.loads(self.data_json)} + + +def _flags(event: str, data: dict[str, Any]) -> str: + flags = "" + if _artifact_task_id(data): + flags += FLAG_TASK + if _is_widget(event, data): + payload = data.get("payload") + if isinstance(payload, dict) and "data" not in payload and payload.get("data_ref"): + flags += FLAG_WIDGET + return flags + + +def line_of(item: dict[str, Any]) -> Line: + data = item["data"] + return Line( + item["event"], + _flags(item["event"], data), + json.dumps(data, ensure_ascii=False, default=str), + ) + + +def encode(lines: list[Line]) -> str: + return "\n".join(f"{ln.event}\t{ln.flags}\t{ln.data_json}" for ln in lines) + + +def decode(text: str) -> list[Line]: + # split("\n"), never splitlines(): JSON leaves U+2028 and friends raw, + # and splitlines() would break a line on them. + out: list[Line] = [] + for raw in text.split("\n") if text else (): + event, flags, data_json = raw.split("\t", 2) + out.append(Line(event, flags, data_json)) + return out diff --git a/src/server/services/history/replay/refresh.py b/src/server/services/history/replay/refresh.py new file mode 100644 index 000000000..c417c43af --- /dev/null +++ b/src/server/services/history/replay/refresh.py @@ -0,0 +1,109 @@ +"""Project turns as they settle, so a reader streams stored lines. + +Two triggers: a turn's run finalizing (the newest two turns, since the +previous turn's answered interrupts only become known at this turn's resume +boundary) and a background task run finalizing (the turn that launched it, +whose lines could not be stored while the run still wrote). Both are +optimizations: a read projects whatever these left undone. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Collection +from dataclasses import dataclass, field + +logger = logging.getLogger(__name__) + + +@dataclass +class _Wanted: + """A thread's next pass, as ``refresh_thread`` takes it.""" + + newest: bool = False + response_ids: set[str] = field(default_factory=set) + + +# Execution context only: coalesces this process's triggers per thread. Truth +# stays in the stored rows, so a trigger landing on another worker, or lost +# to a restart, only costs the next read a projection. +_runners: dict[str, asyncio.Task[None]] = {} +_pending: dict[str, _Wanted] = {} + + +def schedule_replay_refresh(thread_id: str, response_id: str | None = None) -> None: + """Fire-and-forget. At most one runner per thread; a trigger landing + while it runs is folded into its next pass.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: # no running loop (sync test contexts) + return + wanted = _pending.setdefault(thread_id, _Wanted()) + if response_id is None: + wanted.newest = True + else: + wanted.response_ids.add(response_id) + current = _runners.get(thread_id) + if current is not None and not current.done(): + return + task = loop.create_task(_run(thread_id)) + _runners[thread_id] = task + task.add_done_callback(lambda done, tid=thread_id: _runner_done(tid, done)) + + +def _runner_done(thread_id: str, task: asyncio.Task[None]) -> None: + """Drop a finished runner without clobbering a newer one.""" + if _runners.get(thread_id) is task: + _runners.pop(thread_id, None) + + +async def _run(thread_id: str) -> None: + while wanted := _pending.pop(thread_id, None): + await refresh_thread(thread_id, wanted.response_ids, newest=wanted.newest) + + +async def refresh_thread( + thread_id: str, + response_ids: Collection[str], + *, + newest: bool = False, + claims_pass: bool = True, +) -> None: + """Project the turns that launched ``response_ids``, and with ``newest`` + the newest two. Without ``claims_pass`` a caller that waits on this never + reads the whole thread: turns that need it are left to this process's + runner.""" + from src.server.services.history.replay import ( + CheckpointReplayUnavailable, + ClaimsPassNeeded, + load_thread_inputs, + project_turns, + ) + + try: + loaded = await load_thread_inputs(thread_id) + if loaded is None: + return + rows, branch_tip = loaded + launching = { + ti + for ti, r in rows.responses_by_turn.items() + if str(r.get("conversation_response_id")) in response_ids + } + await project_turns( + rows, + branch_tip, + turn_indexes=sorted(launching), + last_n_turns=2 if newest else None, + claims_pass=claims_pass, + ) + except ClaimsPassNeeded: + if newest: + schedule_replay_refresh(thread_id) + for response_id in response_ids: + schedule_replay_refresh(thread_id, response_id) + except CheckpointReplayUnavailable as e: + logger.debug(f"[REPLAY] refresh skipped for {thread_id}: {e}") + except Exception as e: + logger.warning(f"[REPLAY] refresh failed for {thread_id}: {e}") diff --git a/src/server/services/history/replay/run_lane.py b/src/server/services/history/replay/run_lane.py new file mode 100644 index 000000000..c98680272 --- /dev/null +++ b/src/server/services/history/replay/run_lane.py @@ -0,0 +1,777 @@ +"""Background-task runs on the replay stream, from stored run slices. + +A launch artifact (Task ``init``/``resume``) in a turn names the run it +started; that run's transcript is its run slice, the span its namespace wrote +between its own input boundary and the next run's. A ledgered launch finds +its run by the ``task_run_id`` stamped on that boundary. A launch from before +the ledger has no stamp and is matched by its prompt, which only a pass over +the whole thread in order can do (``assign_claims``); the result is stored +with the turn, so the pass runs once. + +A run still writing belongs to its live stream, not to replay, and a turn +projected while one of its runs could still change is never stored. Neither +is a run whose end rests only on the liveness probe, which fails open: its +slice serves the read that cut it and is cut again by the next. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from enum import Enum +from typing import Any + +from langchain_core.messages import ToolMessage + +from ptc_agent.agent.transcript.classify import is_run_boundary_message +from src.config.settings import get_redis_ttl_workflow_events +from src.server.database.conversation import turn_slices as slices_db +from src.server.database.runs import subagent_runs as sr_db +from src.server.services.history import projector +from src.server.services.history import slices +from src.server.services.history import task_streams +from src.server.services.history.replay import facts as replay_facts +from src.server.services.history.replay.errors import ( + CheckpointReplayUnavailable, + ClaimsNeeded, +) +from src.server.services.history.projector import ( + history_events_to_sse, + messages_to_history_events, +) +from src.server.services.history.reader import CheckpointHistoryReader, TaskRun +from src.server.services.history.slices import SpanDelta, TurnSlice +from src.server.services.history.task_status import ( + TERMINAL_TASK_STATUSES, + resolve_task_details, +) +from src.server.utils.content_normalizer import normalize_text_content + +logger = logging.getLogger(__name__) + +# Cap concurrent per-task checkpointer reads: each holds a pool connection, +# so an uncapped gather over a task-heavy thread can exhaust the shared +# checkpointer pool and stall live turns. +_TASK_READ_CONCURRENCY = 8 + + +@dataclass(frozen=True) +class Launch: + task_id: str + action: str + prompt: str | None + run_id: str | None + # Position among the turn's init/resume launches: the index into the + # turn's claims. None for a workflow launch, which claims nothing. + claim_index: int | None + + +def launches_in(turn: SpanDelta) -> list[Launch]: + """A turn's launch artifacts in order. ``update`` (steering) never + launches a run.""" + out: list[Launch] = [] + claim_index = 0 + for message in turn.messages: + if not isinstance(message, ToolMessage): + continue + artifact = (message.additional_kwargs or {}).get("task_artifact") + if not isinstance(artifact, dict) or not artifact.get("task_id"): + continue + action = artifact.get("action", "init") + if action not in ("init", "resume", "workflow"): + continue + prompt = artifact.get("prompt") + run_id = artifact.get("task_run_id") + out.append( + Launch( + task_id=str(artifact["task_id"]), + action=action, + prompt=prompt.strip() if isinstance(prompt, str) else None, + run_id=str(run_id) if run_id else None, + claim_index=None if action == "workflow" else claim_index, + ) + ) + if action != "workflow": + claim_index += 1 + return out + + +def _namespace_items( + messages: list[Any], *, thread_id: str, agent: str +) -> list[dict[str, Any]]: + """Task-namespace messages as SSE items tagged to the task's lane. + + Artifact events are dropped. The task lane honors only side-channel + artifacts that act on the live workspace (a file write refreshes open + panels, a preview server opens one), and a replay would act again on + writes and servers long past. + """ + return [ + item + for item in history_events_to_sse( + messages_to_history_events(messages, agent=agent), + thread_id=thread_id, + ) + if item.get("event") != "artifact" + ] + + +def _run_items(thread_id: str, run: SpanDelta, agent: str) -> list[dict[str, Any]]: + out = projector.context_signal_items(thread_id, run, agent=agent) + out.extend(projector.model_fallback_items(thread_id, run, agent=agent)) + out.extend(_namespace_items(run.messages, thread_id=thread_id, agent=agent)) + return out + + +def _opener_text(messages: list[Any]) -> str | None: + """Text of a run's input message: the prompt its launch carried.""" + for message in messages: + if is_run_boundary_message(message): + content = message.content + if isinstance(content, str): + return content.strip() + text, _ = normalize_text_content(content) + return text.strip() if text else None + return None + + +async def project_task_transcript( + thread_id: str, task_id: str +) -> list[dict[str, Any]]: + """Full SSE-shaped transcript of one task namespace, on demand. + + Serves drill-in views for tasks replay never projects as lanes: + workflow children have no Task-tool launch artifact in the main + transcript, so no turn claims their runs. Same conversion as the lane, + so the client reuses its reduction path. + + Empty unless the task has settled: a running task belongs to its live + stream, and a checkpoint read here would hand the caller a partial + transcript that never self-heals. An unresolvable status counts as + running. + """ + try: + details = await resolve_task_details(thread_id, [task_id]) + except Exception: + logger.warning( + "[REPLAY] task status probe failed for %s/%s", + thread_id, + task_id, + exc_info=True, + ) + return [] + if (details.get(task_id) or {}).get("status") not in TERMINAL_TASK_STATUSES: + return [] + + reader = CheckpointHistoryReader.get_instance() + history = await reader.aget_task_history(thread_id, task_id) + return _run_items(thread_id, history, f"task:{task_id}") + + +async def _gather_bounded(coros: list[Any]) -> list[Any]: + semaphore = asyncio.Semaphore(_TASK_READ_CONCURRENCY) + + async def _run(coro: Any) -> Any: + async with semaphore: + return await coro + + return await asyncio.gather( + *(_run(coro) for coro in coros), return_exceptions=True + ) + + +class Finality(Enum): + """Whether a run has finished writing its span.""" + + OPEN = "open" + FINAL = "final" + # Final only by the liveness probe, which fails open: an unledgered run + # then also needs its stream's seal before it or the turn that launched + # it may be stored. + FINAL_IF_SEALED = "final_if_sealed" + + +@dataclass(frozen=True) +class _Claim: + # The run's (task id, input checkpoint), or None if it never wrote a + # boundary. + key: tuple[str, str] | None + finality: Finality + # The ledger run its boundary was stamped with. + run_id: str | None = None + + +@dataclass +class _RunSlice: + """A run's slice in hand, cut by this read or read from the table.""" + + slice: SpanDelta + # Stored, or queued to be: only a settled run is, so it is final. + stored: bool = False + # Where the stored row says the run ended, and the ledger run it names: + # set only for a row read from the table. + tail: str | None = None + run_id: str | None = None + + +@dataclass(frozen=True) +class _Ledger: + """The thread's run ledger, as one read saw it.""" + + status: dict[str, str] + started: dict[str, float] + failure: dict[str, str] + # Tasks with a run the ledger still holds open. + open_tasks: frozenset[str] + + +# The ledger until it is read: no run in it, and nothing to keep a run open. +_UNREAD = _Ledger({}, {}, {}, frozenset()) + + +class RunLane: + """Resolve and project the runs a set of turns launched.""" + + def __init__( + self, + reader: CheckpointHistoryReader, + thread_id: str, + run_facts: replay_facts.RunFacts, + *, + cache: bool = True, + ): + self._reader = reader + self._thread_id = thread_id + # Without it no stored run slice is read: every run is cut from its + # checkpoints. + self._cache = cache + # The read's snapshot: the same rows the turns' lines keys covered. + self._facts = run_facts + self._serde = reader.serde + # None once a read of it failed. + self._ledger: _Ledger | None = _UNREAD + self._stream_states: dict[str, str] = {} + self._probed: set[str] = set() + self._live_tasks: set[str] = set() + self._walks: dict[str, list[TaskRun]] = {} + # (task_id, input checkpoint) -> run slice, stored or extracted. + self._runs: dict[tuple[str, str], _RunSlice] = {} + # Ledger run -> the key of the stored row that names it. + self._by_run_id: dict[str, tuple[str, str]] = {} + self._loaded_tasks: set[str] = set() + self._writes: list[slices_db.RunRow] = [] + self._workflow: dict[str, SpanDelta] = {} + # (turn input checkpoint, claim index) -> resolved claim. + self._claims: dict[tuple[str, int], _Claim] = {} + + # ------------------------------------------------------------ prepare + + @property + def can_store(self) -> bool: + """Whether what this lane resolved may be stored. Without the ledger + every run's end rests on the liveness probe, and claims and lines + resolved from it would outlive the outage.""" + return self._ledger is not None + + def take_writes(self) -> list[slices_db.RunRow]: + """Run slice rows queued since the last call, for the caller to + store alongside its turns.""" + writes, self._writes = self._writes, [] + return writes + + async def prepare( + self, + turns: list[tuple[TurnSlice, list[str | None] | None, datetime | None]], + ) -> None: + """Resolve every run these turns (each with its claims and the time + its row settled) launched, slicing each finished run nothing has + sliced yet and queueing the settled ones to store. + + Raises ``ClaimsNeeded`` when a turn holds a launch only a + whole-thread pass can place. + """ + launched = [ + (turn, claims, settled_at, launches_in(turn)) + for turn, claims, settled_at in turns + ] + all_launches = [ln for _t, _c, _s, lns in launched for ln in lns] + if not all_launches: + return + await self._load_evidence(all_launches) + + # Walk every task a launch cannot settle from stored rows alone. + walk: set[str] = set() + for _turn, claims, _settled, lns in launched: + for ln in lns: + if ln.action == "workflow" or self._in_progress(ln.run_id): + continue + if ln.run_id and self._stored_key(ln.run_id) is not None: + continue + if not ln.run_id and claims is not None: + claimed = _claimed(claims, ln) + if claimed is None or self._is_stored((ln.task_id, claimed)): + continue + walk.add(ln.task_id) + await self._walk(sorted(walk)) + + wanted: dict[str, dict[str, TaskRun]] = {} + settled: list[TaskRun] = [] + for turn, claims, settled_at, lns in launched: + for ln in lns: + if ln.action == "workflow": + continue + at = (turn.anchor.input_checkpoint_id, ln.claim_index) + if self._in_progress(ln.run_id): + # This exact run is still executing: its stream replays + # the run from its first event, so a projected copy + # would render twice. + self._claims[at] = _Claim(None, Finality.OPEN) + continue + found = self._find_run(ln, claims) + if found is None: + self._claims[at] = _Claim( + None, + Finality.FINAL + if self._ledgered(ln.run_id) + else Finality.FINAL_IF_SEALED, + ) + continue + key, run = found + if run is None: # known only by its stored row + self._claims[at] = _Claim(key, Finality.FINAL) + continue + finality = self._finality(run) + self._claims[at] = _Claim(key, finality, run.task_run_id) + if finality is not Finality.OPEN and key not in self._runs: + wanted.setdefault(run.task_id, {})[run.input_checkpoint_id] = run + if self._settled(run.task_id, finality, settled_at): + settled.append(run) + await self._extract(wanted) + for run in settled: + self._queue_store(run) + + async def _load_evidence(self, launches: list[Launch]) -> None: + """What resolving these launches reads before any run's namespace: + seal evidence first (the ledger, then the streams and liveness of + unledgered tasks), which only moves toward terminal, so a run seen + terminal here cannot write after the namespace reads that follow; + then the launched tasks' stored runs and workflow snapshots.""" + await self._load_ledger() + unledgered = sorted( + {ln.task_id for ln in launches if not self._ledgered(ln.run_id)} + ) + self._stream_states.update( + await task_streams.task_stream_states(self._thread_id, unledgered) + ) + await self._load_liveness(unledgered) + await self._load_stored({ln.task_id for ln in launches}) + await self._read_workflows( + sorted({ln.task_id for ln in launches if ln.action == "workflow"}) + ) + + async def _load_ledger(self) -> None: + if self._ledger is not _UNREAD: + return + try: + rows = await sr_db.list_runs_for_thread(self._thread_id) + except Exception: + self._ledger = None + logger.warning( + "[REPLAY] run-ledger read failed for %s", + self._thread_id, + exc_info=True, + ) + return + status: dict[str, str] = {} + started: dict[str, float] = {} + failure: dict[str, str] = {} + open_tasks: set[str] = set() + for r in rows: + run_id = str(r["task_run_id"]) + status[run_id] = str(r["status"]) + if r.get("started_at") is not None: + started[run_id] = r["started_at"].timestamp() * 1000.0 + error = r.get("failure") + if isinstance(error, dict) and error.get("error"): + failure[run_id] = str(error["error"]) + if r.get("task_id") and str(r["status"]) not in sr_db.TERMINAL_STATUSES: + open_tasks.add(str(r["task_id"])) + self._ledger = _Ledger(status, started, failure, frozenset(open_tasks)) + + async def _load_liveness(self, task_ids: list[str]) -> None: + """Same liveness truth that stamps card status: an unledgered task is + live only while its writer provably runs. On probe failure nothing is + marked live (a transient duplicate beats a missing transcript), which + is why such a run still needs its stream's seal to be stored.""" + task_ids = [t for t in task_ids if t not in self._probed] + if not task_ids: + return + self._probed.update(task_ids) + try: + details = await resolve_task_details(self._thread_id, task_ids) + except Exception: + logger.warning( + "[REPLAY] task liveness probe failed for %s", + self._thread_id, + exc_info=True, + ) + return + self._live_tasks.update( + task_id + for task_id, detail in details.items() + if (detail or {}).get("status") == "running" + ) + + async def _load_stored(self, task_ids: set[str]) -> None: + task_ids = task_ids - self._loaded_tasks + if not task_ids or not self._cache: + return + self._loaded_tasks |= task_ids + for row in await slices_db.get_run_slices(self._thread_id, sorted(task_ids)): + if row.slice.slice_key != slices.SLICE_KEY: + continue + run_slice = slices.decoded(self._serde, row.slice) + if run_slice is None: + continue + key = (row.task_id, row.slice.input_checkpoint_id) + self._runs[key] = _RunSlice( + run_slice, + stored=True, + tail=row.slice.tail_checkpoint_id, + run_id=row.task_run_id, + ) + if row.task_run_id: + self._by_run_id[row.task_run_id] = key + for task_id in task_ids & self._walks.keys(): + self._drop_moved(task_id) + + def _drop_moved(self, task_id: str) -> None: + """Forget stored slices the walk no longer ends where they ended: the + run wrote on after it was stored, so its slice is cut again.""" + for run in self._walks.get(task_id) or (): + key = (task_id, run.input_checkpoint_id) + found = self._runs.get(key) + if ( + found + and found.tail is not None + and found.tail != run.tail_checkpoint_id + ): + del self._runs[key] + + def _is_stored(self, key: tuple[str, str]) -> bool: + found = self._runs.get(key) + return found is not None and found.stored + + def _stored_key(self, run_id: str) -> tuple[str, str] | None: + """The key of the stored row naming this ledger run, while it holds.""" + key = self._by_run_id.get(run_id) + found = self._runs.get(key) if key else None + return key if found is not None and found.run_id == run_id else None + + async def _read_workflows(self, task_ids: list[str]) -> None: + """A workflow run's namespace holds no transcript, only the driver's + terminal ui snapshot (empty until the run settles), so it is read + whole rather than sliced.""" + task_ids = [t for t in task_ids if t not in self._workflow] + results = await _gather_bounded( + [self._reader.aget_task_history(self._thread_id, t) for t in task_ids] + ) + for task_id, history in zip(task_ids, results): + if isinstance(history, BaseException): + self._unavailable(task_id, history) + self._workflow[task_id] = history + + async def _walk(self, task_ids: list[str]) -> None: + task_ids = [t for t in task_ids if t not in self._walks] + results = await _gather_bounded( + [self._reader.aget_task_runs(self._thread_id, t) for t in task_ids] + ) + for task_id, runs in zip(task_ids, results): + if isinstance(runs, BaseException): + self._unavailable(task_id, runs) + self._walks[task_id] = runs + self._drop_moved(task_id) + + async def _extract(self, wanted: dict[str, dict[str, TaskRun]]) -> None: + """Slice finished runs nothing has sliced yet, for this read.""" + tasks = sorted(t for t, runs in wanted.items() if runs) + results = await _gather_bounded( + [ + self._reader.aget_run_slices( + self._thread_id, t, list(wanted[t].values()) + ) + for t in tasks + ] + ) + for task_id, extracted in zip(tasks, results): + if isinstance(extracted, BaseException): + self._unavailable(task_id, extracted) + for run, run_slice in zip(wanted[task_id].values(), extracted): + self._runs[(task_id, run.input_checkpoint_id)] = _RunSlice(run_slice) + + def _unavailable(self, task_id: str, error: BaseException) -> None: + logger.warning( + "[REPLAY] Failed to read subagent checkpoint state task:%s", + task_id, + exc_info=(type(error), error, error.__traceback__), + ) + # Silent continuation would produce a plausible-looking but + # incomplete transcript and bypass the endpoint's SSE fallback. + raise CheckpointReplayUnavailable( + f"subagent checkpoint state unavailable for task:{task_id}" + ) from error + + def _settled( + self, task_id: str, finality: Finality, settled_at: datetime | None + ) -> bool: + """A run's end may be stored: the ledger is readable and the run is + final by more than the liveness probe, so a stored slice can be + trusted as final ever after (``_finality``).""" + if not self.can_store or finality is Finality.OPEN: + return False + return finality is Finality.FINAL or self._stream_sealed(task_id, settled_at) + + def _queue_store(self, run: TaskRun) -> None: + """Queue a settled run's slice to store. Best effort: one that fails + to store is cut again by the next read that needs it.""" + found = self._runs.get((run.task_id, run.input_checkpoint_id)) + if found is None or found.stored: + return + codec, data = slices.encode(self._serde, found.slice) + self._writes.append( + slices_db.RunRow( + thread_id=self._thread_id, + task_id=run.task_id, + # A stamp the ledger no longer holds cannot be referenced. + task_run_id=run.task_run_id + if self._ledgered(run.task_run_id) + else None, + slice=slices_db.StoredSlice( + input_checkpoint_id=run.input_checkpoint_id, + tail_checkpoint_id=run.tail_checkpoint_id, + slice_key=slices.SLICE_KEY, + codec=codec, + data=data, + ), + ) + ) + found.stored = True + + # ------------------------------------------------------------ claims + + @property + def _book(self) -> _Ledger: + """The ledger as read: empty until then, and after a failed read.""" + return self._ledger or _UNREAD + + def _ledgered(self, run_id: str | None) -> bool: + return bool(run_id) and run_id in self._book.status + + def _in_progress(self, run_id: str | None) -> bool: + return ( + self._ledgered(run_id) + and self._book.status[run_id] not in sr_db.TERMINAL_STATUSES + ) + + def _finality(self, run: TaskRun) -> Finality: + if self._is_stored((run.task_id, run.input_checkpoint_id)): + return Finality.FINAL # only a settled run is ever stored + runs = self._walks.get(run.task_id) or [] + if runs and run.ordinal < runs[-1].ordinal: + return Finality.FINAL # a later run started, so this one ended + if self._ledgered(run.task_run_id): + if self._book.status[run.task_run_id] in sr_db.TERMINAL_STATUSES: + return Finality.FINAL + return Finality.OPEN + if run.task_id in self._live_tasks or run.task_id in self._book.open_tasks: + return Finality.OPEN + return Finality.FINAL_IF_SEALED + + def _find_run( + self, launch: Launch, claims: list[str | None] | None + ) -> tuple[tuple[str, str], TaskRun | None] | None: + """The key of the run a launch started, with the run as the walk + found it (None when only its stored row is known), or None if it + never wrote a boundary.""" + if launch.run_id: + stored = self._stored_key(launch.run_id) + if stored is not None: + return stored, None + runs = self._walks.get(launch.task_id) or [] + for run in runs: + if run.task_run_id == launch.run_id: + return (run.task_id, run.input_checkpoint_id), run + if all(run.task_run_id for run in runs): + return None + if claims is None: + raise ClaimsNeeded(launch.task_id) + claimed = _claimed(claims, launch) + if claimed is None: + return None + key = (launch.task_id, claimed) + for run in self._walks.get(launch.task_id) or (): + if run.input_checkpoint_id == claimed: + return key, run + if self._is_stored(key): + return key, None + # The claim was placed against a run the walk no longer has, and + # nothing stored it: its transcript cannot be read. + raise CheckpointReplayUnavailable( + f"claimed run task:{launch.task_id} is no longer on its walk" + ) + + async def assign_claims( + self, turns: list[tuple[list[Launch], datetime | None]] + ) -> list[list[str | None]]: + """Place every run launch of the thread, given each turn's launches + and the time its row settled, oldest turn first. + + One cursor per task walks its runs in order. A stamped launch claims + the run carrying its stamp; an unstamped one claims the next run whose + input matches its prompt (or, when either side has no text to compare, + the next run). A launch with no match claims nothing: its run never + wrote a boundary, and pairing it by position would hand it the next + run's transcript. Every final run is cut to compare its prompt; the + settled ones a launch claims are queued to store. + """ + launches = [[ln for ln in lns if ln.action != "workflow"] for lns, _ in turns] + task_ids = sorted({ln.task_id for lns in launches for ln in lns}) + await self._load_evidence([ln for lns in launches for ln in lns]) + await self._walk(task_ids) + await self._extract( + { + t: { + r.input_checkpoint_id: r + for r in self._walks[t] + if (t, r.input_checkpoint_id) not in self._runs + and self._finality(r) is not Finality.OPEN + } + for t in task_ids + } + ) + + cursors = dict.fromkeys(task_ids, 0) + out: list[list[str | None]] = [] + for lns, (_launches, settled_at) in zip(launches, turns): + claims: list[str | None] = [] + for ln in lns: + runs = self._walks[ln.task_id] + cursor = cursors[ln.task_id] + claimed: TaskRun | None = None + if ln.run_id: + claimed = next( + (r for r in runs[cursor:] if r.task_run_id == ln.run_id), None + ) + if claimed is None and cursor < len(runs): + if ln.prompt is None or self._opener(runs[cursor]) is None: + claimed = runs[cursor] + else: + claimed = next( + (r for r in runs[cursor:] if self._opener(r) == ln.prompt), + None, + ) + if claimed is not None: + cursors[ln.task_id] = claimed.ordinal + 1 + if self._settled( + claimed.task_id, self._finality(claimed), settled_at + ): + self._queue_store(claimed) + claims.append(claimed.input_checkpoint_id if claimed else None) + out.append(claims) + return out + + def _opener(self, run: TaskRun) -> str | None: + found = self._runs.get((run.task_id, run.input_checkpoint_id)) + return _opener_text(found.slice.messages) if found is not None else None + + # ------------------------------------------------------------ project + + def project( + self, turn: TurnSlice, settled_at: datetime | None + ) -> tuple[list[dict[str, Any]], bool, dict[str, float]]: + """``(items, sealed, watermarks)`` for the runs this turn launched. + + ``sealed``: every run is settled (``_settled``), so the projection + may be stored. + ``watermarks``: task id -> newest ledgered run start this projection + contains, the client's authority for which runs its history payload + already holds. + """ + out: list[dict[str, Any]] = [] + launches = launches_in(turn) + sealed = not launches or self.can_store + watermarks: dict[str, float] = {} + book = self._book + for ln in launches: + agent = f"task:{ln.task_id}" + if ln.action == "workflow": + # Frames are reconciled against the ledger row, the status + # authority; un-reconciled frames rebuild per read. + if self._ledgered(ln.run_id): + if self._in_progress(ln.run_id): + sealed = False + elif ln.run_id or not self._stream_sealed(ln.task_id, settled_at): + sealed = False + out.extend( + projector.workflow_run_items( + self._workflow.get(ln.task_id) or SpanDelta(), + task_id=ln.task_id, + ledger_status=book.status.get(ln.run_id) if ln.run_id else None, + ledger_failure=book.failure.get(ln.run_id) + if ln.run_id + else None, + ) + ) + continue + claim = self._claims.get((turn.anchor.input_checkpoint_id, ln.claim_index)) + if claim is None or claim.finality is Finality.OPEN: + sealed = False + continue + if claim.finality is Finality.FINAL_IF_SEALED and not self._stream_sealed( + ln.task_id, settled_at + ): + sealed = False + if claim.key is None: + continue + found = self._runs.get(claim.key) + if found is None: + sealed = False + continue + if ln.run_id and ln.run_id in book.started: + started = book.started[ln.run_id] + watermarks[ln.task_id] = max( + watermarks.get(ln.task_id, started), started + ) + out.extend( + replay_facts.apply_run_facts( + _run_items(self._thread_id, found.slice, agent), + self._facts.for_run(claim.run_id or ln.run_id), + agent=agent, + thread_id=self._thread_id, + message_ids=[getattr(m, "id", None) for m in found.slice.messages], + ) + ) + return out, sealed, watermarks + + def _stream_sealed(self, task_id: str, settled_at: datetime | None) -> bool: + """An unledgered run is final by its stream's end sentinel, or by a + missing stream once the turn settled long enough ago that a sealed + stream would have expired: a task still writing always has its + stream, which carries no TTL until it seals. Twice the TTL leaves + room for a task that outlived its turn.""" + if task_id in self._book.open_tasks: + return False + state = self._stream_states.get(task_id, task_streams.STREAM_UNKNOWN) + if state == task_streams.STREAM_SEALED: + return True + if state == task_streams.STREAM_MISSING and settled_at is not None: + age = datetime.now(timezone.utc) - settled_at + return age > timedelta(seconds=2 * get_redis_ttl_workflow_events()) + return False + + +def _claimed(claims: list[str | None], launch: Launch) -> str | None: + index = launch.claim_index + return claims[index] if index is not None and index < len(claims) else None diff --git a/src/server/services/history/replay/segment_claim.py b/src/server/services/history/replay/segment_claim.py deleted file mode 100644 index 4b585e1cc..000000000 --- a/src/server/services/history/replay/segment_claim.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Run-segment claims: split a task namespace into runs and attribute launches. - -The content-matching claim path is legacy quarantine — ledgered launches join -by stamped task_run_id, but stamp/segment misalignment and pre-ledger data -still fall back to opener-text matching. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any - -from ptc_agent.agent.transcript.classify import is_run_boundary_message -from src.server.utils.content_normalizer import normalize_text_content - - -def split_run_segments(messages: list[Any]) -> list[list[Any]]: - """Split a task namespace's transcript at run boundaries (the plain - HumanMessage each spawn/resume opens with). A defensive leading slice - with no boundary attaches to the first run.""" - segments: list[list[Any]] = [] - current: list[Any] = [] - saw_boundary = False - for message in messages: - if is_run_boundary_message(message): - if saw_boundary: - segments.append(current) - current = [] - saw_boundary = True - current.append(message) - if current: - segments.append(current) - return segments - - -@dataclass -class TaskRuns: - history: Any - segments: list[list[Any]] - cursor: int = 0 - attributed: bool = False - last_ctx: tuple[Any, str | None] | None = None - # A live writer's in-flight run is owned by its stream, not replay: the - # final launch of a live task claims no segment (tail mode commits the - # launch while the run still writes — projecting it would duplicate the - # epoch the stream replays from seq 1). - live: bool = False - remaining_launches: int = 0 - # Ledger join (M4): stamps[i] is the task_run_id the i-th run segment's - # input boundary carries (None = pre-ledger run); run_status maps - # task_run_id -> ledger status. Empty when the walk/ledger read failed — - # every claim then takes the legacy content-matching path. - stamps: list[str | None] = field(default_factory=list) - run_status: dict[str, str] = field(default_factory=dict) - - -def segment_opener_text(segment: list[Any]) -> str | None: - """Text of a segment's run-boundary HumanMessage (its launch input).""" - for message in segment: - if is_run_boundary_message(message): - content = message.content - if isinstance(content, str): - return content.strip() - text, _ = normalize_text_content(content) - return text.strip() if text else None - return None - - -def claim_segment_by_stamp(runs: TaskRuns, run_id: str) -> list[Any] | None: - """Exact ledger join: the segment whose input boundary carries this - launch's ``task_run_id``. - - Usable only when stamps align 1:1 with segments (both derive from the - same run boundaries; a mismatch means the walk saw a boundary shape - the splitter didn't) — misalignment or a missing stamp returns None - and the caller falls back to the legacy content-matching claim. - """ - if len(runs.stamps) != len(runs.segments): - return None - for idx in range(runs.cursor, len(runs.segments)): - if runs.stamps[idx] == run_id: - runs.cursor = idx + 1 - return runs.segments[idx] - return None - - -# dies with Phase 7: ledgered launches always join by stamp; only pre-ledger -# v1 task lanes need the content match below. -def claim_segment(runs: TaskRuns, prompt: str | None) -> list[Any] | None: - """The launch's run segment, or None if no segment belongs to it. - - Scans forward from the cursor for the segment whose boundary opener - matches the launch prompt (skipped-over segments belong to earlier, - already-projected turns and are never revisited). Verification needs - both sides: a prompt-less artifact or a boundary-less segment (legacy - data) can't be checked, so those claim positionally at the cursor — - only a present-but-different opener refuses the pairing. - """ - if runs.cursor >= len(runs.segments): - return None - candidate = runs.segments[runs.cursor] - if prompt is None or segment_opener_text(candidate) is None: - runs.cursor += 1 - return candidate - for idx in range(runs.cursor, len(runs.segments)): - if segment_opener_text(runs.segments[idx]) == prompt: - segment = runs.segments[idx] - runs.cursor = idx + 1 - return segment - return None diff --git a/src/server/services/history/replay/stopped.py b/src/server/services/history/replay/stopped.py index 0a5e59ab4..2c89481a5 100644 --- a/src/server/services/history/replay/stopped.py +++ b/src/server/services/history/replay/stopped.py @@ -1,6 +1,4 @@ -"""Stopped-turn knowledge: which lanes a lossy terminal leaves unfinished, and the close a user stop owes. - -Builds only on ``lanes``; ``stored_merge`` and ``items`` both build on it.""" +"""The close a user-stopped turn owes its last assistant message.""" from __future__ import annotations @@ -10,33 +8,6 @@ from src.server.services.history.replay.lanes import MAIN_LANE, agent_lane -# Run terminals where the live capture can exceed the checkpoint: a run that -# raised or was stopped mid-write streamed output whose message never -# committed. Read from a task run's ledger row for its lane, and from the -# turn's response row for the main lane (see ``resurrect_lanes``). -# completed/interrupted runs end on a committed boundary, and a completed -# run's archive may hold phantom partials from a mid-stream model retry; -# resurrecting those would double-render. -LOSSY_TERMINAL_STATUSES = frozenset({"error", "cancelled"}) - - -def resurrect_lanes( - response: dict[str, Any] | None, task_lanes: set[str] | frozenset[str] -) -> set[str]: - """The claimed task lanes whose run ended lossy, plus the main lane when - the turn itself did. - - A stopped or failed turn loses its in-flight step from the checkpoint - (the stop flush deliberately writes no in-flight messages, and a node - that raises commits nothing), yet the same finalize that records the - status archives every row the user watched stream. - """ - lanes = set(task_lanes) - if (response or {}).get("status") in LOSSY_TERMINAL_STATUSES: - lanes.add(MAIN_LANE) - return lanes - - def stop_close_item( thread_id: str, response: dict[str, Any] | None, @@ -53,9 +24,6 @@ def stop_close_item( has no message to name, so its close carries no ``id``: clients apply a close to the turn's bubble by ``turn_index``. System cancels are not stops and get no close. - - Public: the shared replay in ``server/app/public.py`` streams stored - events itself and needs the same close. """ if not is_user_stop(response): return None diff --git a/src/server/services/history/replay/stored_events.py b/src/server/services/history/replay/stored_events.py new file mode 100644 index 000000000..2f07dad01 --- /dev/null +++ b/src/server/services/history/replay/stored_events.py @@ -0,0 +1,112 @@ +"""Replay from a turn's stored ``sse_events``, and the re-read that brings +them to a turn projected afresh. + +A turn that cannot be projected replays from its stored events (or as a +stub), and its lines are never stored, so nothing here is digested into the +lines epoch. +""" + +from __future__ import annotations + +from typing import Any + +from src.server.database import conversation as conversation_db +from src.server.services.history.replay import items +from src.server.services.history.replay import legacy +from src.server.services.history.replay import stored_merge +from src.server.services.history.replay.lines import Line, line_of +from src.server.services.history.replay.turn import Inputs + + +async def with_stored_events( + responses_by_turn: dict[Any, dict[str, Any]], + turn_indexes: list[Any], + *, + verbatim: bool = False, +) -> dict[Any, dict[str, Any]]: + """``responses_by_turn`` with these turns' rows re-read whole. + + Replay rows arrive without ``sse_events`` and ``replay_facts``; only a + turn projected afresh (or replayed from storage) needs them. The re-read + replaces the row rather than adding the columns, so the events, the + facts and the ``xmin`` the lines key carries come from one tuple. + + A row whose legacy facts stand in for its stored events comes without + them, unless ``verbatim``: a replay of the stored events themselves (the + fallbacks, a parity check). + """ + wanted = { + str(row["conversation_response_id"]): ti + for ti in turn_indexes + if (row := responses_by_turn.get(ti)) and "sse_events" not in row + } + if not wanted: + return responses_by_turn + whole = await conversation_db.get_replay_responses( + list(wanted), legacy_facts=None if verbatim else legacy.FACTS_VERSION + ) + out = dict(responses_by_turn) + for response_id, ti in wanted.items(): + if response_id in whole: + out[ti] = whole[response_id] + return out + + +def build_sse_replay_items( + thread_id: str, + queries: list[dict[str, Any]], + responses_by_turn: dict[Any, dict[str, Any]], +) -> list[dict[str, Any]]: + """Replay items sourced verbatim from persisted ``sse_events`` (the fallback). + + Same ``{"event", "data"}`` shape as the checkpoint path, so the endpoint + emits either source through one loop. The terminal error event is + synthesized from the response row here too — it is yielded live *after* + the persist snapshot, so stored events never contain it. So is a user + stop's close when the archive lacks one. + """ + out: list[dict[str, Any]] = [] + terminals_emitted: set[str] = set() + for query in queries: + if not isinstance(query, dict): + continue + turn_index = query.get("turn_index") + response = responses_by_turn.get(turn_index) + response_id = ( + str(response.get("conversation_response_id")) if response else None + ) + out.append(items._user_message_item(thread_id, query, response)) + turn_items = [ + {"event": event["event"], "data": dict(event["data"])} + for event in stored_merge._stored_events(response) + if stored_merge._valid_stored(event) + ] + if response_id and response_id not in terminals_emitted: + terminals = items.terminal_items(thread_id, response, turn_items) + if terminals: + terminals_emitted.add(response_id) + turn_items += terminals + for item in turn_items: + items._enrich(item, thread_id, turn_index, response_id) + out.extend(turn_items) + return out + + +async def fallback_lines( + inputs: Inputs, turn_indexes: list[Any] +) -> dict[Any, list[Line]]: + """Each turn's replay from its stored events, or its stub without them.""" + responses = await with_stored_events( + inputs.responses_by_turn, turn_indexes, verbatim=True + ) + return { + ti: [ + line_of(item) + for item in build_sse_replay_items( + inputs.thread_id, + inputs.queries_by_turn.get(ti, []), + {ti: responses[ti]} if ti in responses else {}, + ) + ] + for ti in turn_indexes + } diff --git a/src/server/services/history/replay/stored_merge.py b/src/server/services/history/replay/stored_merge.py index 298c91c69..90a1c4c2a 100644 --- a/src/server/services/history/replay/stored_merge.py +++ b/src/server/services/history/replay/stored_merge.py @@ -1,59 +1,23 @@ -"""Stored-event merge: anchor stored payloads into a projected turn, restore losses.""" +"""Stored-event merge: resolve what a turn's stored events add to its projection. + +``derive_turn`` turns the stored events into the turn's legacy facts, which +``legacy.apply_turn`` places. A row the backfill has not reached derives them +here on every projection; a backfilled row carries them and never reads its +stored events (see ``legacy``). +""" from __future__ import annotations -import re -from collections.abc import Callable +from collections.abc import Callable, Iterator from typing import Any -from ptc_agent.agent.middleware.large_result_eviction import TOO_LARGE_TOOL_MSG -from src.server.services.history.replay import widgets -from src.server.services.history.replay.lanes import agent_lane - - -# Stable prefix of the pointer LargeResultEvictionMiddleware substitutes for an -# over-threshold tool result (derived from the canonical template so it tracks -# any edit there). The evicted full content is checkpointed as this pointer, so -# a projected result carrying it must be restored from the stored event. -_EVICTED_RESULT_PREFIX = TOO_LARGE_TOOL_MSG.split("{", 1)[0] - - -# Stored events replayed verbatim (anchored to their original position): -# non-derivable payloads, plus resolved-interrupt cards and error markers — -# `interrupt` renders answered HITL cards on replay (a pending interrupt is -# also re-emitted from the checkpoint tip; the frontend dedups by -# interrupt_id), and `error` keeps wire parity with sse replay. -# `context_window` and `steering_delivered` are checkpoint-projected, but a -# turn with stored events prefers those verbatim (see _merge_stored_payloads). -_PASSTHROUGH_EVENTS = ( - "context_window", - "provenance", - "steering_delivered", - # No projected twin: nothing was injected, so only the captured/stored - # copy can ever surface a returned steering input. - "steering_returned", - "credit_usage", - "interrupt", - "error", - "model_fallback", -) - - -# Projected/synthesized event types the stored stream also carries in full. -# While the sse dual-write is on, a turn with stored events drops its -# projected copies and replays the stored ones (richer historical payloads, -# proven anchoring); post-cutover turns have no stored events, so the -# projected path serves. The terminal ``error`` event is NOT here — stored -# events never contain it (persisted before it is yielded), so it is -# synthesized from the response row on both paths unconditionally. -_STORED_PREFERRED_EVENTS = ( - "steering_delivered", - "context_window", - "provenance", - "credit_usage", - "interrupt", - "model_fallback", +from ptc_agent.agent.middleware.image_capture import ( + IMAGE_MD_RE, + is_sandbox_image_path, ) +from src.server.services.history.replay import legacy +from src.server.services.history.replay import widgets +from src.server.services.history.replay.lanes import MAIN_LANE, agent_lane def _stored_events(response: dict[str, Any] | None) -> list[dict[str, Any]]: @@ -71,92 +35,31 @@ def _valid_stored(event: Any) -> bool: ) -def _replay_main_lane_from_storage( - task_items: list[dict[str, Any]], - stored_events: list[dict[str, Any]], - turn_lossy_lanes: set[str], -) -> list[dict[str, Any]]: - """Rebuild a turn's items with the main lane replayed verbatim from - storage while the task lane keeps the normal projection+merge contract.""" - # The task lane keeps the normal path's contract exactly, via the same - # merge: projection is the transcript authority for the agents it claimed - # (the checkpoint outranks every stored shape — collector full copy, - # user-stop partial snapshot, legacy interleave), while - # _merge_stored_payloads supplies what the checkpoint can't — - # stored-preferred signal rows anchored in position, evicted tool results - # restored from the fuller stored copy — and stored transcript rows serve - # only as anchors, never duplicates. Stored rows outside the claimed lanes - # replay verbatim: main rows, task-scoped custom artifacts (the projector - # never emits task artifacts), and unclaimed agents' rows (in_progress, - # cascade-truncated, unclaimed legacy — nothing projected, nothing to - # merge). Copy the nested ``data`` (like build_sse_replay_items): _enrich - # stamps into it, and the source dicts are the request's pristine - # ``sse_events`` rows. - projected_task_agents = { - agent - for i in task_items - if (agent := str((i.get("data") or {}).get("agent", ""))) - } - claimed_rows: list[dict[str, Any]] = [] - verbatim_rows: list[dict[str, Any]] = [] - for e in stored_events: - if not _valid_stored(e): - continue - agent = str((e.get("data") or {}).get("agent", "")) - if agent in projected_task_agents and e["event"] != "artifact": - claimed_rows.append(e) - else: - verbatim_rows.append({"event": e["event"], "data": dict(e["data"])}) - return verbatim_rows + _merge_stored_payloads( - task_items, claimed_rows, turn_lossy_lanes - ) - - -# Content types the projector emits — a message's anchorable chunks. Live-only -# accumulation chunks (content_type=None) and tool-only messages have none, so -# both streams enumerate the same messages when keyed on these. -_ANCHORABLE_CONTENT_TYPES = frozenset({"reasoning_signal", "reasoning", "text"}) +# The pointer LargeResultEvictionMiddleware checkpoints in place of an +# over-threshold tool result. Only a result evicted after it streamed is +# restored, which only older turns are (eviction now runs first), so the +# wording to recognize is the one those turns were recorded under: fixed +# here, not read from the agent-facing template, which may be reworded. +EVICTED_RESULT_PREFIX = "Tool result too large, the result of this tool call " -# Only these classes prove a lane's archive landed: the collector and the -# stop drain both replay the captured stream from its start (the opener is -# always first), while the live root path persists a task's custom -# artifacts — and nothing else — before any archive write, so an artifact -# row is not evidence. -_ARCHIVE_EVIDENCE_EVENTS = frozenset( - {"user_message", "message_chunk", "tool_calls", "tool_call_result"} -) +def is_evicted_pointer(content: Any) -> bool: + return isinstance(content, str) and content.startswith(EVICTED_RESULT_PREFIX) -# Image targets are rewritten between capture and checkpoint (sandbox path -# -> durable URL; the checkpoint rewrite sees whole messages while the -# archive rewrite scans row fragments), so copy-matching must not compare -# them verbatim. The rewrite preserves the file's basename (the storage key -# ends /{basename}), so normalizing targets to their basename keeps the two -# copies of one image equal while distinct images stay distinct — stripping -# the target entirely would collapse every same-alt image into one signature. -_IMAGE_MD_RE = re.compile(r"!\[([^\]]*)\]\(([^)]*)\)") - - -def _normalize_image_targets(text: str) -> str: - def _basename(match: re.Match) -> str: - target = match.group(2).split("?", 1)[0].split("#", 1)[0].rstrip("/") - return f"![{match.group(1)}]({target.rsplit('/', 1)[-1]})" - - return _IMAGE_MD_RE.sub(_basename, text) - - -def _message_lane_ordinals( +def message_lane_ordinals( out: list[dict[str, Any]], ) -> tuple[dict[str, int], dict[str, int]]: """Ordinal of each message within its lane, over anchorable-content messages, plus the per-lane message count. - Message ids differ between the streams (live chunks carry ``lc_run--…`` run - ids, checkpointed messages the provider id), so a chunk can't anchor by id. - But both streams enumerate a lane's messages in the same order, so the - lane-relative ordinal is a stable cross-stream identity — computed the same - way on each stream, no translation table needed. + Message ids differ between the streams (live chunks carry ``lc_run--…`` + run ids, checkpointed messages the provider id), so a stored chunk can't + find its projected message by id. But both streams enumerate a lane's + messages in the same order, so the lane-relative ordinal pairs them, + computed the same way on each stream. The pairing holds only against the + projection it was made with, which is why facts keep the projected row's + own key instead (``legacy.row_keys``). """ counters: dict[str, int] = {} ordinal_by_id: dict[str, int] = {} @@ -164,7 +67,7 @@ def _message_lane_ordinals( data = item["data"] if item["event"] != "message_chunk": continue - if data.get("content_type") not in _ANCHORABLE_CONTENT_TYPES: + if data.get("content_type") not in legacy.ANCHORABLE_CONTENT_TYPES: continue message_id = data.get("id") if not message_id or message_id in ordinal_by_id: @@ -175,163 +78,338 @@ def _message_lane_ordinals( return ordinal_by_id, counters -def _anchor_key( - event_type: str, data: dict[str, Any], ordinals: dict[str, int] -) -> tuple | None: - """Identity shared by a stored event and its projected counterpart. +def by_ordinal(ordinals: dict[str, int]) -> legacy.MessageKey: + """Key message chunks by lane ordinal, so a stored row and its projected + counterpart share a key. ``ordinals`` is ``message_lane_ordinals`` of + the stream the rows are from.""" - ``ordinals`` maps message id → lane ordinal for the same stream ``data`` - came from (see ``_message_lane_ordinals``). - """ - if event_type == "message_chunk": - content_type, message_id = data.get("content_type"), data.get("id") - if content_type not in _ANCHORABLE_CONTENT_TYPES or message_id not in ordinals: + def key(lane: str, data: dict[str, Any], _occurrence: int) -> tuple | None: + message_id = data.get("id") + if message_id not in ordinals: return None - return ("message_chunk", agent_lane(data.get("agent")), ordinals[message_id], content_type) - if event_type == "tool_calls": - tool_call_ids = tuple( - tc.get("id") for tc in data.get("tool_calls") or [] if tc.get("id") + return ("message_chunk", lane, ordinals[message_id], data["content_type"]) + + return key + + +# Image targets are rewritten between capture and checkpoint (sandbox path +# -> durable URL; the checkpoint rewrite sees whole messages while the +# archive rewrite scans row fragments), so copy-matching must not compare +# them verbatim. The rewrite preserves the file's basename (the storage key +# ends /{basename}), so normalizing targets to their basename keeps the two +# copies of one image equal while distinct images stay distinct — stripping +# the target entirely would collapse every same-alt image into one signature. +def _image_basename(target: str) -> str: + return target.split("?", 1)[0].split("#", 1)[0].rstrip("/").rsplit("/", 1)[-1] + + +def _normalize_image_targets(text: str) -> str: + return IMAGE_MD_RE.sub( + lambda match: f"![{match.group(1)}]({_image_basename(match.group(2))})", text + ) + + +def _image_targets( + rows: list[dict[str, Any]], +) -> Iterator[tuple[dict[str, Any], str]]: + """(row, target) of each markdown image in the rows' text chunks.""" + for row in rows: + data = row["data"] + content = data.get("content") + if ( + row["event"] == "message_chunk" + and data.get("content_type") == "text" + and isinstance(content, str) + ): + for match in IMAGE_MD_RE.finditer(content): + yield row, match.group(2) + + +def _stored_image_urls( + turn_items: list[dict[str, Any]], + task_items: list[dict[str, Any]], + stored_events: list[dict[str, Any]], +) -> tuple[dict[str, str], bool]: + """The stored image URL of each sandbox path the projection still shows, + and whether a path outside the run lanes has several. + + A legacy turn whose capture map never reached its checkpoint holds the + URLs only in its stored copy, where capture rewrote each path, one row + at a time, to a key ending in the file's basename. A path no stored URL + names was never captured (its stored copy shows the path too) and stays + as it is. + """ + urls: dict[str, set[str]] = {} + for _row, target in _image_targets(stored_events): + if not is_sandbox_image_path(target): + urls.setdefault(_image_basename(target), set()).add(target) + images: dict[str, str] = {} + ambiguous = False + run_rows = {id(item) for item in task_items} + for item, path in _image_targets(turn_items): + if not is_sandbox_image_path(path): + continue + found = urls.get(_image_basename(path), set()) + if len(found) == 1: + images[path] = next(iter(found)) + elif found and id(item) not in run_rows: + ambiguous = True + return images, ambiguous + + +def derive_turn( + turn_items: list[dict[str, Any]], + task_items: list[dict[str, Any]], + stored_events: list[dict[str, Any]], + *, + main_durations_known: bool = False, +) -> legacy.LegacyFacts: + """The turn's legacy facts, resolved against its projection. + + Sandbox image paths the projection still shows resolve to their stored + URLs (``images``). Only when several stored URLs share the basename of + a path outside the run lanes does the main lane replay from its stored + events ``verbatim``, where each image keeps the URL capture wrote there. + Rows outside the lanes the turn's runs projected go with it: the main + lane, task-scoped custom artifacts (the projector never emits task + artifacts), and unclaimed agents' rows (in_progress, cascade-truncated, + unclaimed legacy — nothing projected, nothing to merge). The run lanes + keep the projection, which is the transcript authority for the agents + it claimed, merged like any other turn. + """ + facts: legacy.LegacyFacts = {"v": legacy.FACTS_VERSION} + stored_events = [e for e in stored_events if _valid_stored(e)] + images, ambiguous = _stored_image_urls(turn_items, task_items, stored_events) + if images: + facts["images"] = images + if ambiguous: + projected_task_agents = { + agent + for i in task_items + if (agent := str((i.get("data") or {}).get("agent", ""))) + } + claimed: list[dict[str, Any]] = [] + verbatim: list[legacy.StoredRow] = [] + for e in stored_events: + agent = str(e["data"].get("agent", "")) + if agent in projected_task_agents and e["event"] != "artifact": + claimed.append(e) + else: + verbatim.append({"event": e["event"], "data": e["data"]}) + facts["verbatim"] = verbatim + merge = derive_merge( + task_items, claimed, main_durations_known=main_durations_known ) - return ("tool_calls", tool_call_ids[0]) if tool_call_ids else None - if event_type == "tool_call_result": - tool_call_id = data.get("tool_call_id") - return ("tool_call_result", tool_call_id) if tool_call_id else None - if event_type == "artifact": - artifact_id = data.get("artifact_id") - return ("artifact", artifact_id) if artifact_id else None - return None + else: + merge = derive_merge( + turn_items, stored_events, main_durations_known=main_durations_known + ) + if merge is not None: + facts["merge"] = merge + return facts -def _merge_stored_payloads( +def derive_merge( turn_items: list[dict[str, Any]], stored_events: list[dict[str, Any]], - resurrect_lanes: set[str] | frozenset[str] = frozenset(), -) -> list[dict[str, Any]]: - """Merge non-derivable stored events into a projected turn, in position. - - Widget payloads upgrade in place — the stored event carries the resolved - data files the checkpoint deliberately omits. Pairing is ordinal: the live - widget artifact_id is random and the stored event has no tool_call_id, but - both streams order widgets by tool execution. - - Passthrough events (context_window, provenance, …) insert after the - projected position of their nearest preceding stored anchor, reproducing - their original mid-turn placement instead of piling up at the end. - - ``resurrect_lanes``: lanes whose run died mid-write (``stopped.resurrect_lanes``). - Their stored rows beyond everything the checkpoint enumerates (see - ``_is_lost_transcript_row``) are real output the user saw live that no - checkpoint holds, replayed after the lane's last projected item (the - turn's end when nothing projected for the lane) instead of being dropped - as unanchored. + *, + main_durations_known: bool = False, +) -> legacy.Merge | None: + """What the stored events contribute to a projected turn, or None when it + has none (the projected signal rows then stand). + + Read-only: resolved against the projection as ``legacy.apply`` will see + it, its stored-preferred rows dropped and its widgets upgraded. + + Stored rows find their projected rows by lane ordinal (``by_ordinal``); + the facts then name those rows by their own keys (``legacy.row_keys``), + which a later projection keeps. + + - ``widgets``: ``[projected artifact id, payload]`` per stored widget, in + order. The stored event carries the resolved data files the + checkpoint deliberately omits. Pairing is ordinal here (the live + widget artifact id is random and the stored event has no tool call + id, but both streams order widgets by tool execution); the id is null + for a payload beyond the projected count. + - ``results``: the stored content of each result the checkpoint holds + only as an eviction pointer, when the live stream captured it whole + (older turns, where eviction ran after SSE emission). + - ``durations``: ``[row key, ms]`` per projected reasoning close. + - ``anchors``: the keys of the stored rows that found a projected row, + in stored order, and of each payload none took (``["widget", k]``). + - ``signals``: passthrough rows and stored widgets, in stored order. + ``at`` indexes the newest anchor streamed since the previous signal, + and the signal follows it, or the newest earlier one since then that + the projection still holds. Without ``at`` a signal follows the + previous one, so its original mid-turn placement survives instead of + piling up at the end. + + Stored transcript rows the checkpoint never committed (a step stopped or + failed mid-stream) anchor nothing and are dropped: replay shows what the + turn committed, which is what the model will see next. The main lane's + reasoning durations are not taken when ``main_durations_known``: the + turn's replay facts supplied them. """ stored = [e for e in stored_events if _valid_stored(e)] if not stored: - return turn_items - - turn_items = [ - i for i in turn_items if i["event"] not in _STORED_PREFERRED_EVENTS + return None + + # The projection as apply sees it once paired widgets are upgraded, and + # the key each of its rows has before that (what the facts name it by). + view = [ + {"event": i["event"], "data": i["data"]} + for i in turn_items + if i["event"] not in legacy.STORED_PREFERRED_EVENTS ] - - projected_widgets = [i for i in turn_items if widgets._is_widget(i["event"], i["data"])] + row_keys = legacy.row_keys(view) stored_widgets = [e for e in stored if widgets._is_widget(e["event"], e["data"])] + projected_widgets = [i for i in view if widgets._is_widget(i["event"], i["data"])] + paired_ids: list[Any] = [] for item, stored_event in zip(projected_widgets, stored_widgets): - item["data"] = dict(stored_event["data"]) - # Stored widgets beyond the projected count (projection missed the - # ToolMessage artifact) are inserted by anchor like passthrough events. - extra_widget_ids = {id(e) for e in stored_widgets[len(projected_widgets):]} - - _restore_evicted_results(turn_items, stored) + paired_ids.append(item["data"].get("artifact_id") or None) + item["data"] = stored_event["data"] + widget_ordinal = {id(e): k for k, e in enumerate(stored_widgets)} + + merge: legacy.Merge = {} + if stored_widgets: + merge["widgets"] = [ + [paired_ids[k] if k < len(paired_ids) else None, e["data"]] + for k, e in enumerate(stored_widgets) + ] + results = _evicted_results(view, stored) + if results: + merge["results"] = results - projected_ordinals, projected_lane_counts = _message_lane_ordinals(turn_items) - stored_ordinals, stored_lane_counts = _message_lane_ordinals(stored) + projected_ordinals, projected_lane_counts = message_lane_ordinals(view) + stored_ordinals, stored_lane_counts = message_lane_ordinals(stored) # A lane whose stored count disagrees with the checkpoint holds a phantom - # partial or lost output, so its ordinals do not line up across the streams - # and its messages are paired by content instead. + # partial or uncommitted output, so its ordinals do not line up across + # the streams and its messages are paired by content instead. shifted_lanes = { lane for lane, count in stored_lane_counts.items() if count != projected_lane_counts.get(lane, 0) } - paired_lanes = set(resurrect_lanes) | shifted_lanes committed_pairs = ( - _pair_committed_messages(turn_items, stored, paired_lanes) if paired_lanes else {} + _pair_committed_messages(view, stored, shifted_lanes) + if shifted_lanes + else {} + ) + index = legacy.key_index(view, by_ordinal(projected_ordinals)) + stored_keys = dict( + zip(map(id, stored), legacy.row_keys(stored, by_ordinal(stored_ordinals))) ) - committed_stored_ids = set(committed_pairs) - - index_by_key: dict[tuple, int] = {} - for idx, item in enumerate(turn_items): - key = _anchor_key(item["event"], item["data"], projected_ordinals) - if key: - # Last occurrence wins so an anchor covers its whole message group - # (e.g. both reasoning-signal items share one key). - index_by_key[key] = idx - if item["event"] == "tool_calls": - # The checkpoint holds a parallel round as one row, the live archive - # as one row per call: each call's row anchors on the round. - for tc in item["data"].get("tool_calls") or []: - if tc.get("id"): - index_by_key[("tool_calls", tc["id"])] = idx - - def duration_key(data: dict[str, Any]) -> tuple | None: + + def duration_key(event: dict[str, Any]) -> tuple | None: + data = event["data"] lane = agent_lane(data.get("agent")) if lane not in shifted_lanes: - return _anchor_key("message_chunk", data, stored_ordinals) + return stored_keys[id(event)] twin = committed_pairs.get(data.get("id")) if twin is None: return None # a phantom attempt: its thinking is not the row's return ("message_chunk", lane, projected_ordinals[twin], "reasoning_signal") - _carry_reasoning_durations(turn_items, stored, index_by_key, duration_key) - - last_lane_index: dict[str, int] = {} - for idx, item in enumerate(turn_items): - last_lane_index[agent_lane(item["data"].get("agent"))] = idx - # A main lane stopped inside its first model call projects nothing. - turn_end = len(turn_items) - 1 + durations = _reasoning_durations( + view, + [ + e + for e in stored + if not ( + main_durations_known + and agent_lane(e["data"].get("agent")) == MAIN_LANE + ) + ], + index, + duration_key, + ) + if durations: + merge["durations"] = [ + [list(row_keys[idx]), elapsed_ms] for idx, elapsed_ms in durations + ] - inserts_after: dict[int, list[dict[str, Any]]] = {} - resurrected_ids: set[str] = set() - anchor_idx = -1 # before the first projected item + anchors: list[list[Any]] = [] + signals: list[legacy.Signal] = [] + newest: int | None = None # the anchor streamed since the previous signal for event in stored: - key = _anchor_key(event["event"], event["data"], stored_ordinals) - if key is not None and key in index_by_key: - anchor_idx = index_by_key[key] + signal: legacy.Signal | None = None + if event["event"] in legacy.PASSTHROUGH_EVENTS: + signal = {"event": event["event"], "data": event["data"]} + elif id(event) in widget_ordinal: + # Placed only if no projected widget takes the payload. + signal = {"widget": widget_ordinal[id(event)]} + if signal is not None: + if newest is not None: + signal["at"] = newest + newest = None + signals.append(signal) + key = stored_keys[id(event)] + if key is not None and (idx := index.get(key)) is not None: + row_key = row_keys[idx] + elif id(event) in widget_ordinal: + # A payload no projected widget took: a widget that takes it + # later stands where it streamed. + row_key = ("widget", widget_ordinal[id(event)]) + else: continue - if _is_lost_transcript_row( - event, key, stored_ordinals, projected_lane_counts, - resurrect_lanes, committed_stored_ids, resurrected_ids, + if row_key is None: + continue + if not anchors or anchors[-1] != list(row_key): + anchors.append(list(row_key)) + newest = len(anchors) - 1 + if signals: + merge["signals"] = signals + referenced = [signal["at"] for signal in signals if "at" in signal] + if referenced: + merge["anchors"] = anchors[: max(referenced) + 1] + return merge + + +def _evicted_results( + view: list[dict[str, Any]], stored: list[dict[str, Any]] +) -> dict[str, legacy.StoredResult]: + """Full tool-result content the checkpoint holds only as a pointer. + + Large results are evicted to the sandbox filesystem before the + ToolMessage is checkpointed, so a projected ``tool_call_result`` may + carry only the "too large, saved to …" pointer. Only results replay + restores are kept: none once the stored result is itself the pointer + (eviction ran before SSE), so newer turns keep nothing here. + """ + stored_results = { + e["data"].get("tool_call_id"): e["data"] + for e in stored + if e["event"] == "tool_call_result" and e["data"].get("tool_call_id") + } + out: dict[str, legacy.StoredResult] = {} + if not stored_results: + return out + for item in view: + data = item["data"] + if item["event"] != "tool_call_result" or not is_evicted_pointer( + data.get("content") ): - # Anchor the resurrection point too, so following stored rows - # (the run's error, further lost rows) keep their relative order. - anchor_idx = max( - anchor_idx, - last_lane_index.get(agent_lane(event["data"].get("agent")), turn_end), - ) - inserts_after.setdefault(anchor_idx, []).append( - {"event": event["event"], "data": dict(event["data"])} - ) - if event["event"] == "message_chunk" and event["data"].get("id"): - resurrected_ids.add(event["data"]["id"]) continue - if event["event"] in _PASSTHROUGH_EVENTS or id(event) in extra_widget_ids: - inserts_after.setdefault(anchor_idx, []).append( - {"event": event["event"], "data": dict(event["data"])} - ) - - merged = list(inserts_after.get(-1, [])) - for idx, item in enumerate(turn_items): - merged.append(item) - merged.extend(inserts_after.get(idx, ())) - return merged + tool_call_id = data.get("tool_call_id") + stored_data = stored_results.get(tool_call_id) + if not stored_data: + continue + content = stored_data.get("content") + if isinstance(content, str) and not is_evicted_pointer(content): + result: legacy.StoredResult = {"content": content} + if "content_type" in stored_data: + result["content_type"] = stored_data["content_type"] + out[tool_call_id] = result + return out -def _carry_reasoning_durations( - turn_items: list[dict[str, Any]], +def _reasoning_durations( + view: list[dict[str, Any]], stored: list[dict[str, Any]], - index_by_key: dict[tuple, int], + index: dict[tuple, int], key_of: Callable[[dict[str, Any]], tuple | None], -) -> None: - """Carry stored reasoning ``elapsed_ms`` onto the projected close. +) -> list[tuple[int, int]]: + """``(projected position, ms)`` for each projected reasoning close the + stored stream timed. How long the model thought is measured on the live stream and exists nowhere in the checkpoint, so the projected close would replay without @@ -354,62 +432,19 @@ def _carry_reasoning_durations( or not isinstance(data.get("elapsed_ms"), int) ): continue - key = key_of(data) - if key is None or key not in index_by_key: + key = key_of(event) + if key is None or key not in index: continue totals[key] = totals.get(key, 0) + data["elapsed_ms"] - + out: list[tuple[int, int]] = [] for key, elapsed_ms in totals.items(): - target = turn_items[index_by_key[key]]["data"] - if target.get("content_type") == "reasoning_signal" and target.get("content") == "complete": - target["elapsed_ms"] = elapsed_ms - - -def _is_lost_transcript_row( - event: dict[str, Any], - key: tuple | None, - stored_ordinals: dict[str, int], - projected_lane_counts: dict[str, int], - resurrect_lanes: set[str] | frozenset[str], - committed_stored_ids: set[str], - resurrected_ids: set[str], -) -> bool: - """Stored transcript content the checkpoint never committed. - - Reachable only for lanes whose run died mid-write: a message whose lane - ordinal lies beyond every projected message, or a tool round with no - projected twin, was streamed live but never checkpointed. Anchored rows - never reach here (they are consumed as duplicates), and a misaligned - committed copy (``committed_stored_ids``) is the checkpointed message - wearing a shifted ordinal, not lost output. - - Of the empty-content chunks only the stop path's synthetic - ``finish_reason: "stopped"`` close passes, and only for a message this - merge already resurrected (``resurrected_ids``): it is the row that marks - that partial as stopped, and it always follows the message it closes. - """ - data = event["data"] - lane = agent_lane(data.get("agent")) - if lane not in resurrect_lanes: - return False - if event["event"] == "message_chunk": - content = data.get("content") - message_id = data.get("id") - if not content: - return ( - data.get("finish_reason") == "stopped" - and message_id in resurrected_ids - ) - return ( - isinstance(content, str) - and data.get("content_type") in _ANCHORABLE_CONTENT_TYPES - and message_id in stored_ordinals - and stored_ordinals[message_id] >= projected_lane_counts.get(lane, 0) - and message_id not in committed_stored_ids - ) - if event["event"] in ("tool_calls", "tool_call_result"): - return key is not None - return False + target = view[index[key]]["data"] + if ( + target.get("content_type") == "reasoning_signal" + and target.get("content") == "complete" + ): + out.append((index[key], elapsed_ms)) + return out def _pair_committed_messages( @@ -425,9 +460,9 @@ def _pair_committed_messages( match-maximizing in-order pairing (LCS) of stored messages against projected messages: a phantom can never consume a projected slot a real copy needs, and occurrences stay distinct — a stored message repeating - a projected text beyond the pairing is lost output, not a copy. Ties - break toward the earliest stored message, so the copy is the first - occurrence and later repeats resurrect. + a projected text beyond the pairing is uncommitted output, not a copy. + Ties break toward the earliest stored message, so the copy is the first + occurrence. """ projected = _lane_message_signatures(turn_items, lanes) committed: dict[str, str] = {} @@ -506,42 +541,3 @@ def _lane_message_signatures( ] for lane, ids in order.items() } - - -def _restore_evicted_results( - turn_items: list[dict[str, Any]], stored: list[dict[str, Any]] -) -> None: - """Restore full tool-result content the checkpoint holds only as a pointer. - - Large results are evicted to the sandbox filesystem before the ToolMessage is - checkpointed, so a projected ``tool_call_result`` may carry only the "too - large, saved to …" pointer. When the live stream captured the full content - (older turns, where eviction ran after SSE emission), it survives in the - stored event — restore it in place by tool_call_id. A no-op once the stored - result is itself the pointer (eviction ran before SSE), so newer turns are - untouched. - """ - stored_results = { - e["data"].get("tool_call_id"): e["data"] - for e in stored - if e["event"] == "tool_call_result" and e["data"].get("tool_call_id") - } - if not stored_results: - return - for item in turn_items: - if item["event"] != "tool_call_result": - continue - content = item["data"].get("content") - if not (isinstance(content, str) and content.startswith(_EVICTED_RESULT_PREFIX)): - continue - stored_data = stored_results.get(item["data"].get("tool_call_id")) - if not stored_data: - continue - stored_content = stored_data.get("content") - if isinstance(stored_content, str) and not stored_content.startswith( - _EVICTED_RESULT_PREFIX - ): - item["data"]["content"] = stored_content - item["data"]["content_type"] = stored_data.get( - "content_type", item["data"].get("content_type") - ) diff --git a/src/server/services/history/replay/task_lane.py b/src/server/services/history/replay/task_lane.py deleted file mode 100644 index e50708ee7..000000000 --- a/src/server/services/history/replay/task_lane.py +++ /dev/null @@ -1,461 +0,0 @@ -"""Per-run projection of background-task namespaces onto the replay stream.""" - -from __future__ import annotations - -import asyncio -import inspect -import logging -from typing import Any - -from langchain_core.messages import ToolMessage - -from src.server.database.runs import subagent_runs as sr_db -from src.server.services.history import projection_cache -from src.server.services.history import projector -from src.server.services.history import replay -from src.server.services.history.reader import CheckpointHistoryReader -from src.server.services.history.projector import ( - history_events_to_sse, - messages_to_history_events, -) -from src.server.services.history.replay import items -from src.server.services.history.replay import segment_claim -from src.server.services.history.replay import stopped -from src.server.services.history.task_status import ( - TERMINAL_TASK_STATUSES, - resolve_task_details, -) - -logger = logging.getLogger(__name__) - -# Cap concurrent per-task checkpointer reads: each holds a pool connection, -# so an uncapped gather over a task-heavy thread can exhaust the shared -# checkpointer pool and stall live turns. -_TASK_READ_CONCURRENCY = 8 - - -async def _gather_bounded(coros: list[Any]) -> list[Any]: - semaphore = asyncio.Semaphore(_TASK_READ_CONCURRENCY) - - async def _run(coro: Any) -> Any: - async with semaphore: - return await coro - - return await asyncio.gather( - *(_run(coro) for coro in coros), return_exceptions=True - ) - - -def _namespace_items( - messages: list[Any], *, thread_id: str, agent: str -) -> list[dict[str, Any]]: - """Task-namespace messages as SSE items tagged to the task's lane. - - Artifact events are dropped. The task lane honors only side-channel - artifacts that act on the live workspace (a file write refreshes open - panels, a preview server opens one), and a replay would act again on - writes and servers long past. - """ - return [ - item - for item in history_events_to_sse( - messages_to_history_events(messages, agent=agent), - thread_id=thread_id, - ) - if item.get("event") != "artifact" - ] - - -async def project_task_transcript( - thread_id: str, task_id: str -) -> list[dict[str, Any]]: - """Full SSE-shaped transcript of one task namespace, on demand. - - Serves drill-in views for tasks replay never projects as lanes — - workflow children have no Task-tool launch artifact in the main - transcript, so ``TaskLaneProjector`` cannot see them. Same conversion - as the lane projector, so the client reuses its reduction path. - - Empty unless the task has settled — the lane path likewise leaves a - running task's segment to its live stream, and a checkpoint read here - would hand the caller a partial transcript that never self-heals. An - unresolvable status counts as running. - """ - try: - details = await resolve_task_details(thread_id, [task_id]) - except Exception: - logger.warning( - "[REPLAY] task status probe failed for %s/%s", - thread_id, - task_id, - exc_info=True, - ) - return [] - if (details.get(task_id) or {}).get("status") not in TERMINAL_TASK_STATUSES: - return [] - - reader = CheckpointHistoryReader.get_instance() - history = await reader.aget_task_history(thread_id, task_id) - task_agent = f"task:{task_id}" - out = projector.context_signal_items(thread_id, history, agent=task_agent) - out.extend( - projector.model_fallback_items(thread_id, history, agent=task_agent) - ) - out.extend( - _namespace_items(history.messages, thread_id=thread_id, agent=task_agent) - ) - return out - - -class TaskLaneProjector: - """Per-run projection of background-task namespaces. - - A task namespace holds every run's transcript back-to-back; each run - opens at a plain HumanMessage (its spawn/resume input). Launch artifacts - (action ``init``/``resume``) in the main transcript attribute to run - segments in order. Ledgered launches join exactly: the artifact's - ``task_run_id`` matches the stamp the run's input boundary carries in - checkpoint metadata, and the run ledger decides projection (``in_progress`` - runs belong to their live stream). Pre-ledger launches verify by content: - a segment's boundary HumanMessage carries the launch prompt verbatim, so - a launch only claims a segment whose opener matches its prompt. A launch - with no matching segment projects nothing — either its boundary isn't - checkpointed yet (the live stream owns that run until the next rebuild) - or the run never wrote one (a failed/no-op launch); blind positional - pairing would hand it the NEXT run's transcript. ``update`` (steering) - artifacts never launch a run. - - Windowed builds may start after a task's init: the cursor then starts at - ``offset`` (the leading segments belong to out-of-window turns). Any - namespace read failure makes checkpoint replay unavailable so - ``source=auto`` can use the complete stored-SSE fallback. - """ - - def __init__(self, thread_id: str, *, windowed: bool): - self._thread_id = thread_id - self._windowed = windowed - self._tasks: dict[str, segment_claim.TaskRuns] = {} - self._run_started: dict[str, float] = {} - # task_run_id -> short failure text from the ledger row's failure - # JSONB — the error surfaced when a workflow snapshot is missing or - # its terminal frame is overridden by the ledger status. - self._run_failure: dict[str, str] = {} - # Tasks whose Redis stream had no end sentinel when prepare() probed - # — BEFORE the namespace reads. The caller's cache guard: a turn that - # launched any of these must not be cached this build, because the - # stream may have sealed between our read and any later probe (the - # read could then hold pre-terminal state that never self-heals). - self.live_streams_at_read: set[str] = set() - # Turn indexes whose stamps carry trailing salvage (populated by - # trailing_items) — those turns must not be cached, or the fast path - # would replay them without the salvage. - self.salvaged_turn_indexes: set[Any] = set() - # task_id -> max started_at (epoch ms) over ledgered runs whose - # segment THIS build claimed. Stamped onto the turn's task artifacts - # as ``projected_run_started_ms``: the client's authority for which - # runs its history payload already contains. Derived from the claim - # act itself — never from a separate ledger read, which can name a - # run the projection skipped (its skip decision and this watermark - # must share one snapshot). - self.claimed_watermarks: dict[str, float] = {} - # Lanes claimed in the CURRENT turn whose run died mid-write - # (stopped.LOSSY_TERMINAL_STATUSES): the stored copy may hold output the - # checkpoint never committed, so the merge may resurrect their - # trailing rows. Reset by each project_for_turn call. - self.turn_lossy_lanes: set[str] = set() - # Lanes the CURRENT turn projected from incomplete inputs. Cache - # control only — deliberately NOT the lossy set, whose other consumer - # resurrects stored rows past the checkpoint: an infrastructure - # failure must cost a rebuild, never change merge semantics. - self.turn_uncacheable_lanes: set[str] = set() - # Workflow status reconciliation needs the ledger; when its read - # failed this build must not be cacheable for workflow lanes (a - # sealed stream would otherwise freeze the snapshot's own status - # for the cache TTL). - self._ledger_read_failed = False - - @staticmethod - def _launches_in(turn: Any) -> list[tuple[str, str, str | None, str | None]]: - """Ordered ``(task_id, action, prompt, task_run_id)`` launch artifacts - in a turn. ``task_run_id`` is None on pre-ledger data. - - ``workflow`` launches (RunWorkflow) count too: the run task's - namespace holds no transcript, but membership in the launched set is - what keeps the launching turn uncacheable while the run's streams - are still live/unarchived. - """ - launches: list[tuple[str, str, str | None, str | None]] = [] - for message in turn.messages: - if not isinstance(message, ToolMessage): - continue - artifact = (message.additional_kwargs or {}).get("task_artifact") - if not isinstance(artifact, dict) or not artifact.get("task_id"): - continue - action = artifact.get("action", "init") - if action in ("init", "resume", "workflow"): - prompt = artifact.get("prompt") - run_id = artifact.get("task_run_id") - launches.append( - ( - str(artifact["task_id"]), - action, - prompt.strip() if isinstance(prompt, str) else None, - str(run_id) if run_id else None, - ) - ) - return launches - - async def prepare( - self, reader: CheckpointHistoryReader, pairs: list[tuple[Any, Any]] - ) -> None: - launch_actions: dict[str, list[str]] = {} - for _, turn in pairs: - if turn is None: - continue - for task_id, action, _prompt, _run_id in self._launches_in(turn): - launch_actions.setdefault(task_id, []).append(action) - if not launch_actions: - return - - task_ids = list(launch_actions) - # Stream-seal probe strictly before the namespace reads: a "sealed" - # verdict here proves the reads below see final state (sealing is - # monotonic), which is what makes the caller's cache store sound. - self.live_streams_at_read = await projection_cache.live_task_streams( - self._thread_id, set(task_ids) - ) - histories = await _gather_bounded( - [reader.aget_task_history(self._thread_id, tid) for tid in task_ids] - ) - stamps_by_task, run_status = await self._load_ledger(reader, task_ids) - for task_id, history in zip(task_ids, histories): - if isinstance(history, BaseException): - logger.warning( - "[REPLAY] Failed to read subagent checkpoint state task:%s", - task_id, - exc_info=(type(history), history, history.__traceback__), - ) - # Silent continuation would produce a plausible-looking but - # incomplete transcript and bypass the endpoint's SSE fallback. - raise replay.CheckpointReplayUnavailable( - f"subagent checkpoint state unavailable for task:{task_id}" - ) from history - segments = segment_claim.split_run_segments(history.messages) - actions = launch_actions[task_id] - # A window that opens on a resume is missing the older runs' - # launches; their segments are skipped, not re-attributed. A full - # build always sees the init, so its cursor starts at segment 0. - cursor = ( - max(0, len(segments) - len(actions)) - if self._windowed and actions[0] != "init" - else 0 - ) - self._tasks[task_id] = segment_claim.TaskRuns( - history=history, - segments=segments, - cursor=cursor, - remaining_launches=len(actions), - stamps=stamps_by_task.get(task_id, []), - run_status=run_status, - ) - - # Same liveness truth that stamps card status (advisory-lock probe): - # a task is live only while its writer provably runs, so an expired - # stream never demotes a settled run's transcript. On probe failure - # nothing is marked live — availability over precision (a transient - # duplicate beats a missing transcript). - try: - details = await resolve_task_details( - self._thread_id, list(self._tasks) - ) - except Exception: - logger.warning( - "[REPLAY] task liveness probe failed for %s", - self._thread_id, - exc_info=True, - ) - details = {} - for task_id, runs in self._tasks.items(): - runs.live = (details.get(task_id) or {}).get("status") == "running" - - async def _load_ledger( - self, reader: CheckpointHistoryReader, task_ids: list[str] - ) -> tuple[dict[str, list[str | None]], dict[str, str]]: - """Boundary stamps per task + a thread-wide run_id -> status map. - - Both reads are best-effort: any failure (and readers without the - stamp walk — test fakes) yields empty results, which routes every - claim through the legacy content-matching path. - """ - stamps_by_task: dict[str, list[str | None]] = {} - run_status: dict[str, str] = {} - # iscoroutinefunction, not truthiness: test fakes are spec'd mocks - # whose auto-created attribute is sync — calling it would hand - # asyncio.gather a non-awaitable. (AsyncMock passes the check.) - stamp_walk = getattr(reader, "aget_task_run_stamps", None) - if stamp_walk is not None and inspect.iscoroutinefunction(stamp_walk): - results = await _gather_bounded( - [stamp_walk(self._thread_id, tid) for tid in task_ids] - ) - for task_id, stamps in zip(task_ids, results): - if isinstance(stamps, BaseException): - logger.warning( - "[REPLAY] task run-stamp walk failed for task:%s", - task_id, - exc_info=(type(stamps), stamps, stamps.__traceback__), - ) - else: - stamps_by_task[task_id] = stamps - try: - runs = await sr_db.list_runs_for_thread(self._thread_id) - run_status = { - str(r["task_run_id"]): str(r["status"]) for r in runs - } - self._run_started = { - str(r["task_run_id"]): r["started_at"].timestamp() * 1000.0 - for r in runs - if r.get("started_at") is not None - } - self._run_failure = { - str(r["task_run_id"]): str((r.get("failure") or {}).get("error")) - for r in runs - if isinstance(r.get("failure"), dict) - and (r.get("failure") or {}).get("error") - } - except Exception: - self._ledger_read_failed = True - logger.warning( - "[REPLAY] run-ledger read failed for %s", - self._thread_id, - exc_info=True, - ) - return stamps_by_task, run_status - - def project_for_turn( - self, turn: Any, turn_index: Any, response_id: str | None - ) -> tuple[list[dict[str, Any]], set[str]]: - """Items for every run launched in this turn, plus the launched task - ids (the caller's cache guard: a turn that launched a still-writing - run must not be cached).""" - out: list[dict[str, Any]] = [] - launched: set[str] = set() - self.turn_lossy_lanes = set() - self.turn_uncacheable_lanes = set() - for task_id, action, prompt, run_id in self._launches_in(turn): - runs = self._tasks.get(task_id) - if runs is None: - continue - launched.add(task_id) - task_agent = f"task:{task_id}" - runs.last_ctx = (turn_index, response_id) - runs.remaining_launches -= 1 - if action == "workflow": - # A workflow run's namespace holds no transcript, only the - # driver's terminal ui snapshot (empty until the run settles; - # the live stream owns it until then). Emit its lifecycle - # frames — reconciled against the ledger row, the status - # authority — and skip segment claiming entirely. - status = runs.run_status.get(run_id) if run_id else None - if self._ledger_read_failed: - # Un-reconciled frames must rebuild per read until the - # ledger is readable again. - self.turn_uncacheable_lanes.add(task_agent) - out.extend( - projector.workflow_run_items( - runs.history, - task_id=task_id, - ledger_status=status, - ledger_failure=( - self._run_failure.get(run_id) if run_id else None - ), - ) - ) - continue - status = runs.run_status.get(run_id) if run_id else None - if status == "in_progress": - # The ledger says this exact run is still executing: its - # stream replays the epoch (opener included) from seq 1, so - # claiming the segment here would render it twice. The turn - # stays uncached (launched set + live stream), so the settled - # rebuild projects it normally. - continue - if status is None and runs.live and runs.remaining_launches <= 0: - # Legacy gate (no ledger row for this launch): without a - # per-run status, only the final launch of a live task can be - # the in-flight run. - continue - segment = ( - segment_claim.claim_segment_by_stamp(runs, run_id) - if run_id and status is not None - else None - ) - if segment is None: - segment = segment_claim.claim_segment(runs, prompt) - if segment is None: - continue - started = self._run_started.get(run_id) if run_id else None - if started is not None: - prev = self.claimed_watermarks.get(task_id) - self.claimed_watermarks[task_id] = ( - started if prev is None else max(prev, started) - ) - if status in stopped.LOSSY_TERMINAL_STATUSES: - self.turn_lossy_lanes.add(task_agent) - if not runs.attributed: - # Namespace-scoped signals (compaction, model fallback) are - # not per-run; they ride with the first projected run. - runs.attributed = True - out.extend( - projector.context_signal_items( - self._thread_id, runs.history, agent=task_agent - ) - ) - out.extend( - projector.model_fallback_items( - self._thread_id, runs.history, agent=task_agent - ) - ) - out.extend( - _namespace_items( - segment, thread_id=self._thread_id, agent=task_agent - ) - ) - return out, launched - - async def trailing_items(self) -> list[dict[str, Any]]: - """Segments beyond the last in-window launch, for settled runs only. - - Covers a launch whose turn never committed (e.g. the launching turn - errored before persist) — salvaged under the last known launch's - stamps. Per-segment ledger gate: a stamped segment projects only when - its run row exists and is terminal — a missing row is a - cascade-truncated run (its launching turn was deleted; resurrecting - it would re-attach deleted work), an ``in_progress`` row belongs to - the live stream. Unstamped (pre-ledger) segments keep the legacy - whole-task liveness gate.""" - out: list[dict[str, Any]] = [] - for task_id, runs in self._tasks.items(): - if runs.cursor >= len(runs.segments) or runs.last_ctx is None: - continue - aligned = len(runs.stamps) == len(runs.segments) - task_agent = f"task:{task_id}" - turn_index, response_id = runs.last_ctx - salvaged_any = False - for idx in range(runs.cursor, len(runs.segments)): - stamp = runs.stamps[idx] if aligned else None - if stamp is not None: - if runs.run_status.get(stamp) not in sr_db.TERMINAL_STATUSES: - continue - elif runs.live: - continue - salvaged_any = True - for item in _namespace_items( - runs.segments[idx], - thread_id=self._thread_id, - agent=task_agent, - ): - items._enrich(item, self._thread_id, turn_index, response_id) - out.append(item) - if salvaged_any: - self.salvaged_turn_indexes.add(turn_index) - return out diff --git a/src/server/services/history/replay/turn.py b/src/server/services/history/replay/turn.py new file mode 100644 index 000000000..08c0de5ed --- /dev/null +++ b/src/server/services/history/replay/turn.py @@ -0,0 +1,203 @@ +"""One turn's replay items: what its slice, its runs and its rows project to. + +Everything here decides what a turn's stored lines hold, so the module is +digested into the lines epoch (``keys.projection_sources``). +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Protocol + +from src.server.database.conversation.replay_rows import ThreadRows +from src.server.services.history import projector +from src.server.services.history import task_status +from src.server.services.history.projector import ( + history_events_to_sse, + messages_to_history_events, +) +from src.server.services.history.replay import facts as replay_facts +from src.server.services.history.replay import items +from src.server.services.history.replay import legacy +from src.server.services.history.replay import run_lane +from src.server.services.history.replay import stored_merge +from src.server.services.history.slices import TurnSlice +from src.server.services.runs.sse_producer import resolve_token_threshold + +# The ui record image capture persists a turn's path -> URL map in. +IMAGE_CAPTURE_UI_NAME = "image_capture" + + +class Derive(Protocol): + """A caller's hook on each turn whose legacy facts were derived from its + stored events (the backfill). Called before the facts are placed, with + the turn's items still unplaced and whether every run it launched has + finished writing; what it returns is called with the placed items.""" + + def __call__( + self, + turn_index: Any, + response: dict[str, Any], + facts: legacy.LegacyFacts, + unplaced: tuple[list[dict[str, Any]], list[dict[str, Any]]], + sealed: bool, + ) -> Callable[[list[dict[str, Any]]], None]: ... + + +@dataclass(frozen=True) +class Inputs: + """The table rows a turn's projection reads, indexed once per request.""" + + thread_id: str + turn_indexes: list[Any] + queries_by_turn: dict[Any, list[dict[str, Any]]] + responses_by_turn: dict[Any, dict[str, Any]] + usage_by_response: dict[str, Any] + provenance_by_response: dict[str, Any] + run_facts: replay_facts.RunFacts + + @classmethod + def of(cls, rows: ThreadRows) -> Inputs: + queries_by_turn: dict[Any, list[dict[str, Any]]] = {} + for q in rows.queries: + if isinstance(q, dict): + queries_by_turn.setdefault(q.get("turn_index"), []).append(q) + return cls( + thread_id=rows.thread_id, + turn_indexes=sorted(ti for ti in queries_by_turn if ti is not None), + queries_by_turn=queries_by_turn, + responses_by_turn=rows.responses_by_turn, + usage_by_response=items._usage_rows_by_response(rows.usages), + provenance_by_response=items._rows_by_response(rows.provenance, many=True), + run_facts=replay_facts.RunFacts.of(rows.run_facts), + ) + + def response_id(self, turn_index: Any) -> str | None: + response = self.responses_by_turn.get(turn_index) + return str(response["conversation_response_id"]) if response else None + + +def project_turn( + inputs: Inputs, + turn_index: Any, + turn: TurnSlice, + response: dict[str, Any] | None, + lane: run_lane.RunLane, + derive: Derive | None = None, +) -> tuple[list[dict[str, Any]], bool]: + """One turn's replay items, and whether they are final (every run it + launched has finished writing).""" + thread_id = inputs.thread_id + response_id = str(response.get("conversation_response_id")) if response else None + + segment = [ + items._user_message_item(thread_id, q, response) + for q in inputs.queries_by_turn.get(turn_index, []) + ] + + turn_items = history_events_to_sse( + messages_to_history_events(turn.messages), thread_id=thread_id + ) + # Compaction signals (offload counts, the summarize event) live in + # private state keys, not the messages channel — re-emit them at the + # head of the turn they landed in (live they fire before the first + # post-compaction model call). Fallback notices follow the same + # head placement: live they fire before the succeeding model's chunks. + turn_items[:0] = projector.context_signal_items( + thread_id, turn + ) + projector.model_fallback_items(thread_id, turn) + durations = replay_facts.reasoning_durations(response) + if durations is not None: + replay_facts.apply_reasoning_durations(turn_items, durations) + task_items, sealed, watermarks = lane.project(turn, settled_at(response)) + turn_items.extend(task_items) + # Legacy path→URL records belong to the turn whose state delta contains + # them. A thread-global map is incorrect when a sandbox filename is + # reused later: last-write-wins would rewrite the older turn's image to + # the newer content-addressed object. + replay_facts.apply_image_url_map( + turn_items, _collect_image_url_map(turn.new_ui_records) + ) + # Table-sourced synthesis rides ahead of the merge: a turn with stored + # events drops these copies and replays the stored ones instead (the + # STORED_PREFERRED_EVENTS transition rule). + turn_items = items._insert_provenance_items( + turn_items, inputs.provenance_by_response.get(response_id) or [] + ) + turn_items.extend( + items._interrupt_item(thread_id, intr) for intr in turn.anchor.ending_interrupts + ) + credit_item = items._credit_usage_item( + thread_id, response, inputs.usage_by_response.get(response_id) + ) + if credit_item: + turn_items.append(credit_item) + + facts = legacy.facts_of(response) + placed: Callable[[list[dict[str, Any]]], None] | None = None + if facts is None: + facts = stored_merge.derive_turn( + turn_items, + task_items, + stored_merge._stored_events(response), + main_durations_known=durations is not None, + ) + if derive is not None and response: + placed = derive( + turn_index, response, facts, (turn_items, task_items), sealed + ) + turn_items = legacy.apply_turn(turn_items, task_items, facts) + if placed is not None: + placed(turn_items) + + _fill_token_thresholds(turn_items) + # Terminal error: never in stored events (persisted before it is + # yielded live), so it appends after the merge on every turn. A user + # stop's close is read off the same row (see stopped.stop_close_item). + turn_items += items.terminal_items(thread_id, response, turn_items) + + # After the stored-events merge so both projected and stored-copy + # artifacts are covered: the watermark is a fact about which runs this + # turn's lines contain. + task_status.stamp_projected_watermarks(turn_items, watermarks) + for item in turn_items: + items._enrich(item, thread_id, turn_index, response_id) + segment.extend(turn_items) + return segment, sealed + + +def settled_at(response: dict[str, Any] | None) -> datetime | None: + """When a settled turn's row last settled, as far as the row says.""" + if not response: + return None + value = response.get("usage_settled_at") or response.get("created_at") + return value if isinstance(value, datetime) else None + + +def _collect_image_url_map(records: list[dict[str, Any]]) -> dict[str, str]: + url_map: dict[str, str] = {} + for record in records: + if not isinstance(record, dict) or record.get("name") != IMAGE_CAPTURE_UI_NAME: + continue + path_to_url = (record.get("props") or {}).get("path_to_url") + if isinstance(path_to_url, dict): + url_map.update({str(k): str(v) for k, v in path_to_url.items() if k and v}) + return url_map + + +def _fill_token_thresholds(turn_items: list[dict[str, Any]]) -> None: + """Stamp the UI-ring threshold on projected token_usage events. + + The live handler adds it server-side (config, not graph state); replay + uses the same resolver so both wires carry the same value. + """ + for item in turn_items: + data = item["data"] + if ( + item["event"] == "context_window" + and data.get("action") == "token_usage" + and "threshold" not in data + ): + data["threshold"] = resolve_token_threshold() diff --git a/src/server/services/history/slices.py b/src/server/services/history/slices.py new file mode 100644 index 000000000..26118de68 --- /dev/null +++ b/src/server/services/history/slices.py @@ -0,0 +1,317 @@ +"""Turn and run slices: what one turn or one subagent run added, stored once. + +A slice is a pure function of two checkpoints (the boundary it starts at and +the state it ends at), so it is derived data: a row that disagrees with the +branch walk is rebuilt, and deleting rows loses nothing. Slices encode with +the checkpointer's own serializer, so whatever a checkpoint can hold, a slice +can too, under the same type allowlist. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +from collections import Counter +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from langchain_core.messages import AnyMessage + +from ptc_agent.agent.transcript.classify import ( + is_run_boundary_message, + is_summary_message, +) + +if TYPE_CHECKING: + from src.server.database.conversation.turn_slices import StoredSlice + from src.server.services.history.reader import CheckpointHistoryReader, TurnAnchor + +logger = logging.getLogger(__name__) + +# Bump when what a slice holds changes in a way the sources keyed below do +# not show (the walk that picks a slice's two checkpoints, say): every stored +# slice then rebuilds from its checkpoints on the next read. +SLICE_VERSION = 1 + +_SRC = Path(__file__).resolve().parents[3] + +# One namespace's checkpointed state values. +State = Mapping[str, Any] + + +@dataclass +class SpanDelta: + """What a namespace added between two of its checkpoints: a turn on the + main branch, a task run in its own namespace, or a whole namespace.""" + + messages: list[AnyMessage] = field(default_factory=list) + # The ``_summarization_event`` that landed in the span, or None. + # Compaction's summary message lives in this state key, never in the + # messages channel, so replay re-emits the summarize signal from here. + new_summarization_event: dict[str, Any] | None = None + # Growth of the compaction offload sets: replay re-emits the offload + # signals from these counts (one aggregated event per kind). + newly_offloaded_args: int = 0 + newly_offloaded_reads: int = 0 + # ``ui``-channel records that landed in the span (an id diff), such as + # the model_fallback notices middleware pushes. + new_ui_records: list[dict[str, Any]] = field(default_factory=list) + + +@dataclass(kw_only=True) +class TurnSlice(SpanDelta): + """One conversational turn on the current branch: what it added, and + where the walk placed it.""" + + anchor: TurnAnchor + + +def _new_summarization_event(start: State, end: State) -> dict[str, Any] | None: + """The span's freshly landed ``_summarization_event``, or None. + + Events are identified by their summary message id (uuid-stamped at build); + an end-state event matching the start state predates this span. + """ + + def _identity(values: State) -> tuple[Any, Any] | None: + event = values.get("_summarization_event") + if not isinstance(event, dict): + return None + summary_message = event.get("summary_message") + return (getattr(summary_message, "id", None), event.get("cutoff_index")) + + end_event = end.get("_summarization_event") + if isinstance(end_event, dict) and _identity(end) != _identity(start): + return end_event + return None + + +def _set_growth(start: State, end: State, key: str) -> int: + """How many ids ``key``'s set gained between the two states.""" + return len(set(end.get(key) or set()) - set(start.get(key) or set())) + + +def _new_ui_records(start: State, end: State) -> list[dict[str, Any]]: + """``ui``-channel records the end state has that the start state lacks.""" + start_ids = {r.get("id") for r in (start.get("ui") or []) if isinstance(r, dict)} + return [ + r + for r in (end.get("ui") or []) + if isinstance(r, dict) and r.get("id") not in start_ids + ] + + +def _new_messages( + start_values: State, end_values: State, *, has_input: bool = True +) -> list[AnyMessage]: + """Messages the end state has that the start state lacks, by id. + + Id-diff, not count-diff, so compaction and REMOVE_ALL are safe. A + message saved without an id (a full-list write from before DeltaChannel) + loads id-less from some checkpoints and with an id from others, so it + matches by content instead, each start message at most once. A message + with a fresh id matches only an id-less one: a repeated "continue" after + compaction dropped the first is new. + + Carried messages all precede the span's own input (``has_input``: every + turn but a HITL resume, every task run), so from that input on nothing + matches by content: a message there repeating one compaction removed is + new, however its id loaded. Nor do they come from before where the end + starts in the start (``_carried_from``): a window trim drops the head of + the list, possibly inside a resume, where no input bounds the matching. + """ + start = start_values.get("messages") or [] + end = end_values.get("messages") or [] + start_ids = {m.id for m in start if m.id is not None} + carried = start[_carried_from(start, end):] + idless = Counter(_content_key(m) for m in carried if m.id is None) + # Start messages whose id the end lost: an id-less end message may be one. + vanished: Counter[tuple[str, str]] | None = None + carried_until = _input_position(end, start_ids) if has_input else None + if carried_until is None: + carried_until = len(end) + new: list[AnyMessage] = [] + for i, m in enumerate(end): + if m.id is not None and m.id in start_ids: + continue + if i >= carried_until: + new.append(m) + continue + pools = [idless] + if m.id is None: + if vanished is None: + end_ids = {e.id for e in end if e.id is not None} + vanished = Counter( + _content_key(s) + for s in carried + if s.id is not None and s.id not in end_ids + ) + pools.append(vanished) + if any(pools): + key = _content_key(m) + pool = next((p for p in pools if p[key] > 0), None) + if pool is not None: + pool[key] -= 1 + continue + new.append(m) + return new + + +def _carried_from(start: list[AnyMessage], end: list[AnyMessage]) -> int: + """Where ``end`` starts in ``start``: the first end message the start + holds by id, counted back over the end messages before it, which a trim + kept but stamped fresh ids on. 0 when the end holds none of them.""" + positions = {m.id: i for i, m in enumerate(start) if m.id is not None} + for offset, message in enumerate(end): + at = positions.get(message.id) if message.id is not None else None + if at is not None: + return max(0, at - offset) + return 0 + + +def _input_position(end: list[AnyMessage], start_ids: set[str]) -> int | None: + """Where the span's own input sits in ``end``: the last real input the + start lacks by id. A span opens on one input and its end state is read + before the next one lands, so no real input follows it. None when + compaction took the input with everything before it.""" + for i in range(len(end) - 1, -1, -1): + m = end[i] + if ( + (m.id is None or m.id not in start_ids) + and is_run_boundary_message(m) + and not is_summary_message(m) + ): + return i + return None + + +def _content_key(message: AnyMessage) -> tuple[str, str]: + content = message.content + if not isinstance(content, str): + content = json.dumps(content, sort_keys=True, default=str) + return message.type, content + + +def delta(start: State, end: State, *, has_input: bool = True) -> SpanDelta: + """What ``end`` holds that ``start`` lacks, both one namespace's state + values. ``has_input``: the span opens on an input message of its own + (every turn but a HITL resume, every task run).""" + return SpanDelta( + messages=_new_messages(start, end, has_input=has_input), + new_summarization_event=_new_summarization_event(start, end), + newly_offloaded_args=_set_growth(start, end, "_offloaded_tool_call_ids"), + newly_offloaded_reads=_set_growth(start, end, "_offloaded_read_result_ids"), + new_ui_records=_new_ui_records(start, end), + ) + + +def turn_slice(anchor: TurnAnchor, start: State, end: State) -> TurnSlice: + span = delta(start, end, has_input=not anchor.is_resume) + return TurnSlice(anchor=anchor, **vars(span)) + + +def encode(serde: Any, span: SpanDelta) -> tuple[str, bytes]: + return serde.dumps_typed( + { + "messages": list(span.messages), + "new_summarization_event": span.new_summarization_event, + "newly_offloaded_args": span.newly_offloaded_args, + "newly_offloaded_reads": span.newly_offloaded_reads, + "new_ui_records": list(span.new_ui_records), + } + ) + + +def decode(serde: Any, codec: str, data: bytes) -> SpanDelta: + payload = serde.loads_typed((codec, data)) + return SpanDelta( + messages=list(payload.get("messages") or []), + new_summarization_event=payload.get("new_summarization_event"), + newly_offloaded_args=int(payload.get("newly_offloaded_args") or 0), + newly_offloaded_reads=int(payload.get("newly_offloaded_reads") or 0), + new_ui_records=list(payload.get("new_ui_records") or []), + ) + + +def decoded(serde: Any, row: StoredSlice) -> SpanDelta | None: + """The stored row's span, or None when it no longer decodes: such a row + is stale like any other, cut again and stored over rather than failing + every read of its thread.""" + try: + return decode(serde, row.codec, row.data) + except Exception: + logger.warning( + "stored slice at %s undecodable", row.input_checkpoint_id, exc_info=True + ) + return None + + +def matches(row: StoredSlice | None, anchor: TurnAnchor) -> bool: + """A stored slice is the turn's only while the branch still ends it where + it ended when stored, and the code that cut it is the code running.""" + return ( + row is not None + and row.input_checkpoint_id == anchor.input_checkpoint_id + and row.tail_checkpoint_id == anchor.tail_checkpoint_id + and row.slice_key == SLICE_KEY + ) + + +def stored_spans( + serde: Any, + anchored: Mapping[Any, TurnAnchor], + rows: Mapping[Any, StoredSlice], +) -> dict[Any, SpanDelta]: + """The span of each turn whose stored row ``matches`` its anchor and + decodes, keyed as both are: what a read may serve without a cut.""" + out: dict[Any, SpanDelta] = {} + for key, anchor in anchored.items(): + row = rows.get(key) + if matches(row, anchor) and (span := decoded(serde, row)) is not None: + out[key] = span + return out + + +async def load_turns( + reader: CheckpointHistoryReader, + thread_id: str, + anchored: list[tuple[Any, TurnAnchor]], + rows: Mapping[Any, StoredSlice], +) -> tuple[dict[Any, TurnSlice], set[Any]]: + """``(turns, fresh)``, keyed as ``anchored`` and ``rows`` are: each turn + from its stored row by ``stored_spans``, placed on the branch its anchor + names, and cut from the checkpoints otherwise (``fresh``).""" + held = stored_spans(reader.serde, dict(anchored), rows) + turns = { + key: TurnSlice(anchor=anchor, **vars(held[key])) + for key, anchor in anchored + if key in held + } + cut = [(key, anchor) for key, anchor in anchored if key not in held] + if cut: + extracted = await reader.aget_turn_slices(thread_id, [a for _, a in cut]) + turns.update((key, turn) for (key, _), turn in zip(cut, extracted)) + return turns, {key for key, _ in cut} + + +def slice_sources() -> list[Path]: + """The code that decides a slice's bytes: this module, which cuts and + encodes it, and the transcript classifier it finds a span's input by. + Digested whole, and a test holds the list closed over their imports, so + a fix to any of it reaches slices stored before it.""" + return [Path(__file__).resolve(), _SRC / "ptc_agent/agent/transcript/classify.py"] + + +def _slice_key() -> str: + digest = hashlib.blake2b(digest_size=8) + for path in slice_sources(): + digest.update(str(path.relative_to(_SRC)).encode()) + digest.update(path.read_bytes()) + return f"{SLICE_VERSION}.{digest.hexdigest()}" + + +# Stored with every slice; a row under another key is rebuilt. +SLICE_KEY = _slice_key() diff --git a/src/server/services/history/task_status.py b/src/server/services/history/task_status.py index cb05b4db8..6f22af658 100644 --- a/src/server/services/history/task_status.py +++ b/src/server/services/history/task_status.py @@ -172,7 +172,9 @@ def stamp_task_artifact_data( return {**data, "payload": payload} -async def stamp_replay_task_status(thread_id: str, items: list[dict]) -> None: +async def stamp_replay_task_status( + thread_id: str, items: list[dict], *, status_only: bool = False +) -> None: """Stamp every task artifact in an assembled replay item list, in place by positional replacement. Best-effort: a failure leaves the items unstamped (the client's live reconciliation still applies).""" @@ -182,7 +184,9 @@ async def stamp_replay_task_status(thread_id: str, items: list[dict]) -> None: return for i, item in enumerate(items): if isinstance(item, dict): - stamped = stamp_task_artifact_data(item.get("data"), details) + stamped = stamp_task_artifact_data( + item.get("data"), details, status_only=status_only + ) if stamped is not item.get("data"): items[i] = {**item, "data": stamped} except Exception: diff --git a/src/server/services/history/task_streams.py b/src/server/services/history/task_streams.py new file mode 100644 index 000000000..78ceca8ec --- /dev/null +++ b/src/server/services/history/task_streams.py @@ -0,0 +1,80 @@ +"""Probe the per-task Redis streams for whether a subagent run sealed.""" + +import asyncio +import json +import logging +from typing import Any + +from ptc_agent.agent.middleware.background_subagent.redis_stream import ( + SUBAGENT_STREAM_END_EVENT, + task_stream_key, +) +from src.utils.cache.redis_cache import get_cache_client + +logger = logging.getLogger(__name__) + + +STREAM_SEALED = "sealed" +STREAM_OPEN = "open" +STREAM_MISSING = "missing" +STREAM_UNKNOWN = "unknown" + + +async def task_stream_states(thread_id: str, task_ids: set[str]) -> dict[str, str]: + """Each task's per-task Redis stream: sealed, open, missing or unknown. + + Only an end sentinel proves a stream sealed. A missing stream is either + one that expired after sealing or one that was never written, which this + probe cannot tell apart; the caller decides with the launching turn's + age (``RunLane``). A failed probe is unknown. + + Probe BEFORE reading task-namespace state, and gate caching on this + pre-read verdict: sealing is monotonic, so "sealed before the read" + proves the read saw final state, while a post-read probe can pass for a + stream that sealed after the read returned pre-terminal data, freezing + it into stored lines. + """ + if not task_ids: + return {} + client = get_cache_client().client + if client is None: + return dict.fromkeys(task_ids, STREAM_UNKNOWN) + + async def _state(task_id: str) -> str: + try: + entries = await client.xrevrange( + task_stream_key(thread_id, task_id), count=1 + ) + except Exception as e: + logger.warning( + f"[TaskStreams] task stream check failed for " + f"{thread_id}/{task_id}: {e}" + ) + return STREAM_UNKNOWN + if not entries: + return STREAM_MISSING + raw = entries[0][1].get(b"event") or entries[0][1].get("event") + if isinstance(raw, bytes): + raw = raw.decode("utf-8", errors="replace") + return STREAM_SEALED if _is_stream_end_sentinel(raw) else STREAM_OPEN + + ordered = list(task_ids) + states = await asyncio.gather(*(_state(task_id) for task_id in ordered)) + return dict(zip(ordered, states)) + + +def _is_stream_end_sentinel(raw: Any) -> bool: + """Match forwarder.finalize()'s ``{"event": "subagent_stream_end"}`` + sentinel (a JSON dict without ``seq`` — same test as the reconnect + consumer's payload classifier).""" + if not isinstance(raw, str) or not raw.startswith("{"): + return False + try: + record = json.loads(raw) + except json.JSONDecodeError: + return False + return ( + isinstance(record, dict) + and record.get("event") == SUBAGENT_STREAM_END_EVENT + and "seq" not in record + ) diff --git a/src/server/services/history/window.py b/src/server/services/history/window.py new file mode 100644 index 000000000..77d50f83a --- /dev/null +++ b/src/server/services/history/window.py @@ -0,0 +1,160 @@ +"""The messages a window trimmed, read back from the turn slices. + +The main agent's checkpoint keeps only the transcript runs it can still see +(``compaction.window``). The messages it drops stay readable in two places: +the checkpoints written before the trim, and the turn slices replay and the +transcripts read. The agent asks here before each trim, so a thread is +trimmed only once its stored slices, read in branch order, start with exactly +the messages the trim drops, and replay of the trimmed turns never falls back +to decoding the checkpoints that still hold them. + +The check is by identity (``Window``), not by counting turns or runs: a turn +can open no run (an input killed before it ran) or two (two user messages in +one request), a request can open with an assistant message that belongs to +the run before it, and a thread compacted by the old summarizer holds a +summary where its slices hold the messages it replaced. A count misreads +each; a thread whose slices do not match is simply not trimmed. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Mapping +from typing import Any + +from langchain_core.messages import AnyMessage + +from ptc_agent.agent.transcript import Window +from src.server.database.conversation import turn_slices as slices_db +from src.server.database.conversation.turn_slices import StoredSlice +from src.server.database.conversation.threads_read import get_thread_checkpoint_id +from src.server.services.history import slices +from src.server.services.history.reader import CheckpointHistoryReader, TurnAnchor +from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound + +logger = logging.getLogger(__name__) + +#: Turns whose stored slices one query reads. +_BATCH = 32 + + +async def _anchors(thread_id: str, tip: str | None) -> list[TurnAnchor] | None: + try: + anchors, _ = await CheckpointHistoryReader.get_instance().aget_turn_anchors( + thread_id, tip + ) + except CheckpointBranchTipNotFound: + logger.debug("window: tip %s of %s is not on the graph", tip, thread_id) + return None + return anchors + + +def _decoded( + serde: Any, + rows: Mapping[tuple[str, str | None], StoredSlice], + turns: list[TurnAnchor], + need: int, +) -> list[list[AnyMessage] | None]: + """Each turn's messages from its stored row, in order, until ``need`` are + in hand; None for a turn whose row replay would not serve either + (``slices.matches``, ``slices.decoded``). CPU only.""" + out: list[list[AnyMessage] | None] = [] + have = 0 + for anchor in turns: + if have >= need: + break + row = rows.get((anchor.input_checkpoint_id, anchor.tail_checkpoint_id)) + span = slices.decoded(serde, row) if slices.matches(row, anchor) else None + messages = span.messages if span is not None else None + out.append(messages) + have += len(messages or ()) + return out + + +async def _head( + thread_id: str, turns: list[TurnAnchor], need: int, *, extract: bool +) -> list[AnyMessage] | None: + """The first ``need`` messages ``turns`` added, in order: from their + stored slices, and with ``extract`` from the checkpoints for a turn no + stored slice serves. None when the turns hold fewer, or without + ``extract`` once one has no stored slice to read.""" + reader = CheckpointHistoryReader.get_instance() + out: list[AnyMessage] = [] + for start in range(0, len(turns), _BATCH): + if len(out) >= need: + break + batch = turns[start : start + _BATCH] + rows = { + (row.input_checkpoint_id, row.tail_checkpoint_id): row + for row in await slices_db.get_slices_at( + thread_id, slices.SLICE_KEY, [a.input_checkpoint_id for a in batch] + ) + } + by_turn = await asyncio.to_thread( + _decoded, reader.serde, rows, batch, need - len(out) + ) + missing = [i for i, messages in enumerate(by_turn) if messages is None] + if missing and not extract: + return None + if missing: + extracted = await reader.aget_turn_slices( + thread_id, [batch[i] for i in missing] + ) + for i, turn in zip(missing, extracted): + by_turn[i] = turn.messages + out.extend(m for messages in by_turn for m in messages or ()) + return out[:need] if len(out) >= need else None + + +async def _read_head( + thread_id: str, tip: str | None, head: Window, *, extract: bool +) -> list[AnyMessage] | None: + """``head``'s messages on the branch ending at ``tip``, read as ``_head`` + reads them, or None when that branch's turns do not start with exactly + them.""" + if not head.messages: + return [] + anchors = await _anchors(thread_id, tip) + messages = await _head(thread_id, anchors or [], head.messages, extract=extract) + if messages is not None and await asyncio.to_thread(head.holds, messages): + return messages + # Only a read back is owed the messages; a check that misses leaves the + # thread untrimmed, as it was. + logger.log( + logging.WARNING if extract else logging.INFO, + "window: %s's turns at %s do not start with its %d trimmed messages", + thread_id, + tip, + head.messages, + ) + return None + + +async def runs_held(thread_id: str, head: Window) -> bool: + """Whether the stored slices of the committed branch's turns, read in + order, start with exactly ``head``'s messages: every slice up to there is + one replay would serve, and together they hash to its digest. + + One indexed read settles the common miss (a thread not yet backfilled + holds no slices); only a thread that may be covered pays for the branch + walk, which reads checkpoint metadata, never state, and then for decoding + the slices a trim would leave as those messages' only copy. + """ + if not head.messages: + return True + if not await slices_db.has_turn_slices(thread_id, slices.SLICE_KEY): + return False + tip = await get_thread_checkpoint_id(thread_id) + return bool(tip) and await _read_head(thread_id, tip, head, extract=False) is not None + + +async def earlier_runs( + thread_id: str, tip: str | None, head: Window +) -> list[AnyMessage] | None: + """``head``'s messages, the ones the window trimmed, on the branch ending + at ``tip`` (the newest checkpoint without one): from the turns' stored + slices, and from the checkpoints for a turn whose slice replay would not + serve. None when that branch's turns do not start with exactly them. + """ + return await _read_head(thread_id, tip, head, extract=True) diff --git a/src/server/services/livefs/mount.py b/src/server/services/livefs/mount.py index eb62556f6..316cdf415 100644 --- a/src/server/services/livefs/mount.py +++ b/src/server/services/livefs/mount.py @@ -35,7 +35,7 @@ from langchain_core.messages import AnyMessage -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.core.paths import SandboxLayout from ptc_agent.core.sandbox import livefs_mount from ptc_agent.core.sandbox.livefs_mount import CallContext, MountOutcome @@ -593,9 +593,9 @@ async def _gone(self, output: str) -> bool: return await livefs_mount.answers(self._sandbox) is False async def save_transcript( - self, target: TranscriptTarget, messages: Sequence[AnyMessage] + self, target: TranscriptTarget, messages: Sequence[AnyMessage], *, window: Window ) -> bool: - return await transcripts.save_live(target, messages) + return await transcripts.save_live(target, messages, window=window) async def restarted(computer_id: Any, sandbox_id: str) -> None: diff --git a/src/server/services/persistence/image_capture.py b/src/server/services/persistence/image_capture.py index 09fd7d206..dfdc29ede 100644 --- a/src/server/services/persistence/image_capture.py +++ b/src/server/services/persistence/image_capture.py @@ -23,23 +23,8 @@ logger = logging.getLogger(__name__) -async def capture_and_rewrite_images( - sse_events: list[dict], - sandbox, - thread_id: str = "", - *, - project: ProjectContext | None = None, - workspace_id: str | None = None, -) -> int: - """Scan SSE events for sandbox image paths, upload to storage, rewrite in-place. - - Returns number of images captured. No-op if storage is disabled. - Non-fatal: logs warnings on failure, never raises. - """ - if not is_storage_enabled() or not sse_events: - return 0 - - # Collect all unique sandbox image paths from text message_chunks +def sandbox_image_paths(sse_events: list[dict]) -> set[str]: + """Every sandbox image path the events' text references.""" image_paths: set[str] = set() for evt in sse_events: if evt.get("event") != "message_chunk": @@ -48,12 +33,34 @@ async def capture_and_rewrite_images( if data.get("content_type") != "text": continue content = data.get("content", "") + if not isinstance(content, str): + continue for match in IMAGE_MD_RE.finditer(content): if is_sandbox_image_path(match.group(2)): image_paths.add(match.group(2)) + return image_paths + + +async def capture_images( + sse_events: list[dict], + sandbox, + thread_id: str = "", + *, + project: ProjectContext | None = None, + workspace_id: str | None = None, +) -> dict[str, str]: + """Upload the sandbox images the events reference and rewrite them in + place, returning each captured path's URL. + + Empty if storage is disabled. Non-fatal: logs warnings on failure, + never raises. + """ + if not is_storage_enabled() or not sse_events: + return {} + image_paths = sandbox_image_paths(sse_events) if not image_paths: - return 0 + return {} if project is None and workspace_id: path_to_url = await _capture_from_held_folder( @@ -63,6 +70,43 @@ async def capture_and_rewrite_images( path_to_url = await capture_sandbox_images( sandbox, image_paths, thread_id, project=project ) + if not path_to_url: + return {} + + for evt in sse_events: + if evt.get("event") != "message_chunk": + continue + data = evt.get("data", {}) + if data.get("content_type") != "text": + continue + content = data.get("content", "") + if content: + data["content"] = rewrite_image_paths(content, path_to_url) + return path_to_url + + +async def capture_and_rewrite_images( + sse_events: list[dict], + sandbox, + thread_id: str = "", + *, + project: ProjectContext | None = None, + workspace_id: str | None = None, +) -> int: + """``capture_images``, plus the captured map recorded on the thread's + checkpoint for replay. Returns the number of images captured. + + For the main turn only, whose checkpoint is still its own when this + runs. A subagent's archive lands after its turn settled, so its map goes + to the run's replay facts instead (``subagent_archive``). + """ + path_to_url = await capture_images( + sse_events, + sandbox, + thread_id, + project=project, + workspace_id=workspace_id, + ) if not path_to_url: return 0 @@ -74,24 +118,14 @@ async def capture_and_rewrite_images( if thread_id: try: from src.server.services.history.reader import CheckpointHistoryReader + from src.server.services.history.replay.turn import IMAGE_CAPTURE_UI_NAME await CheckpointHistoryReader.get_instance().append_ui_record( - thread_id, "image_capture", {"path_to_url": path_to_url} + thread_id, IMAGE_CAPTURE_UI_NAME, {"path_to_url": path_to_url} ) except Exception as e: logger.warning(f"[IMAGE_CAPTURE] Failed to persist ui record: {e}") - # Rewrite image paths in SSE events in-place - for evt in sse_events: - if evt.get("event") != "message_chunk": - continue - data = evt.get("data", {}) - if data.get("content_type") != "text": - continue - content = data.get("content", "") - if content: - data["content"] = rewrite_image_paths(content, path_to_url) - return len(path_to_url) diff --git a/src/server/services/runs/coordinator.py b/src/server/services/runs/coordinator.py index 7da90fa06..ffe1952ca 100644 --- a/src/server/services/runs/coordinator.py +++ b/src/server/services/runs/coordinator.py @@ -14,7 +14,7 @@ from contextlib import nullcontext from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Any, Awaitable, Callable, Coroutine, Dict, List, Optional +from typing import Any, Awaitable, Callable, Coroutine, Dict, Optional from uuid import uuid4 from src.server.database.runs import lifecycle as tl_db @@ -24,6 +24,7 @@ FinalizeResult, ForkSpec, QuerySpec, + RunOutcome, RunSlotBusyError, ) from src.server.utils.error_sanitization import sanitize_error_text @@ -81,21 +82,6 @@ async def protected_finalize(coro: Coroutine, label: str): return await asyncio.shield(spawn_protected(coro, f"turn-finalize-{label}")) -@dataclass -class RunOutcome: - """Everything finalize needs, resolved in-band by the run's own executor.""" - - status: str # completed | interrupted | error | cancelled - interrupt_reason: Optional[str] = None - metadata: Dict[str, Any] = field(default_factory=dict) - warnings: Optional[List[str]] = None - errors: Optional[List[str]] = None - execution_time: Optional[float] = None - sse_events: Optional[List[Dict[str, Any]]] = None - per_call_records: Optional[list] = None - tool_usage: Optional[Dict[str, int]] = None - - @dataclass class RunHandle: """Identity of one run (= one conversation_responses row), START to finalize.""" @@ -334,13 +320,7 @@ async def finalize_run( result = await tl_db.finalize_run( run_id=handle.run_id, thread_id=handle.thread_id, - status=outcome.status, - interrupt_reason=outcome.interrupt_reason, - metadata=outcome.metadata, - warnings=outcome.warnings, - errors=outcome.errors, - execution_time=outcome.execution_time, - sse_events=outcome.sse_events, + outcome=outcome, checkpoint_id=checkpoint_id, usage_writer=usage_writer, conn=guard.conn if guard is not None else None, @@ -432,13 +412,7 @@ async def finalize_detached_run( result = await tl_db.finalize_run( run_id=run_id, thread_id=thread_id, - status=outcome.status, - interrupt_reason=outcome.interrupt_reason, - metadata=outcome.metadata, - warnings=outcome.warnings, - errors=outcome.errors, - execution_time=outcome.execution_time, - sse_events=outcome.sse_events, + outcome=outcome, checkpoint_id=checkpoint_id, ) if result.applied: @@ -654,12 +628,12 @@ def post_finalize_tail( self, thread_id: str, run: Optional[Dict[str, Any]] ) -> None: """The uniform after-a-won-CAS tail: every site that applies a - terminal transition schedules the projection refresh and the + terminal transition schedules the replay refresh and the transcript export and nudges the drainer, so they never drift apart per call site. The export runs at every terminal status, since each stamps the checkpoint and another thread may read this one's transcript.""" - self._schedule_projection_refresh(thread_id) + self._schedule_replay_refresh(thread_id) self._schedule_transcript_export(thread_id, run) self._schedule_archive_prune(thread_id, run) self._nudge_hook_drainer() @@ -674,16 +648,16 @@ def _nudge_hook_drainer(self) -> None: except Exception: logger.warning("[RunCoordinator] hook drainer nudge failed", exc_info=True) - def _schedule_projection_refresh(self, thread_id: str) -> None: + def _schedule_replay_refresh(self, thread_id: str) -> None: try: - from src.server.services.history.projection_cache import ( - schedule_projection_refresh, + from src.server.services.history.replay.refresh import ( + schedule_replay_refresh, ) - schedule_projection_refresh(thread_id) + schedule_replay_refresh(thread_id) except Exception: logger.warning( - f"[RunCoordinator] projection refresh scheduling failed for " + f"[RunCoordinator] replay refresh scheduling failed for " f"{thread_id}", exc_info=True, ) diff --git a/src/server/services/runs/credit_usage.py b/src/server/services/runs/credit_usage.py new file mode 100644 index 000000000..3c5c76bb1 --- /dev/null +++ b/src/server/services/runs/credit_usage.py @@ -0,0 +1,43 @@ +"""The ``credit_usage`` wire payload, shared by the live stream and replay. + +Its own module because replay stores what it builds: the replay projection +digests this file, so a change to the payload's shape re-projects every +stored turn instead of serving the old shape from storage. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Dict, Optional + + +def build_credit_usage_data( + thread_id: str, + token_usage: dict, + total_credits: float, + timestamp: Optional[str] = None, +) -> Dict[str, Any]: + """Aggregate per-model usage into the ``credit_usage`` wire payload. + + Intentionally omits USD costs and model names (hidden from the client). + Shared by the live handler and table-sourced replay so both wires carry + the same shape. + """ + total_input_tokens = 0 + total_output_tokens = 0 + total_tokens = 0 + for usage in (token_usage or {}).get("by_model", {}).values(): + total_input_tokens += usage.get("input_tokens", 0) + total_output_tokens += usage.get("output_tokens", 0) + total_tokens += usage.get("total_tokens", 0) + + return { + "thread_id": thread_id, + "tokens": { + "input_tokens": total_input_tokens, + "output_tokens": total_output_tokens, + "total_tokens": total_tokens, + }, + "total_credits": round(total_credits, 2), + "timestamp": timestamp or datetime.now().isoformat(), + } diff --git a/src/server/services/runs/event_archive.py b/src/server/services/runs/event_archive.py new file mode 100644 index 000000000..050ddbc65 --- /dev/null +++ b/src/server/services/runs/event_archive.py @@ -0,0 +1,117 @@ +"""The merged event archive a run persists as ``sse_events``. + +Token-level frames merge per message and per tool call, so the archive grows +with the turn rather than with its token count. Recovery salvages a dead +run's Redis stream through the same accumulator, so both archives merge alike. +""" + +import copy +from typing import Any, Dict, List + +from src.config.settings import get_merged_chunk_max_bytes + +MERGED_STREAM_CHUNK_MAX_BYTES_DEFAULT = get_merged_chunk_max_bytes() + + +class StreamEventAccumulator: + """Accumulates and merges token-level SSE events for persistence.""" + + def __init__(self, max_merged_bytes: int = MERGED_STREAM_CHUNK_MAX_BYTES_DEFAULT): + self._max_merged_bytes = max_merged_bytes + self._events: List[Dict[str, Any]] = [] + + def get_events(self) -> List[Dict[str, Any]]: + return copy.deepcopy(self._events) + + def add(self, event_type: str, data: Dict[str, Any]) -> None: + if not isinstance(data, dict): + return + + incoming = copy.deepcopy(data) + + if not self._events: + self._events.append({"event": event_type, "data": incoming}) + return + + prev = self._events[-1] + if prev.get("event") != event_type: + self._events.append({"event": event_type, "data": incoming}) + return + + if event_type == "message_chunk" and self._try_merge_message_chunk(prev, incoming): + return + + if event_type == "tool_call_chunks" and self._try_merge_tool_call_chunks(prev, incoming): + return + + self._events.append({"event": event_type, "data": incoming}) + + def _try_merge_message_chunk(self, prev_event: Dict[str, Any], incoming: Dict[str, Any]) -> bool: + prev_data = prev_event.get("data") + if not isinstance(prev_data, dict): + return False + + if incoming.get("content_type") == "reasoning_signal": + return False + if prev_data.get("content_type") == "reasoning_signal": + return False + + merge_keys = ("thread_id", "agent", "id", "role", "content_type", "phase") + if any(prev_data.get(k) != incoming.get(k) for k in merge_keys): + return False + + prev_content = prev_data.get("content") or "" + incoming_content = incoming.get("content") or "" + incoming_finish = incoming.get("finish_reason") + + if incoming_content: + if len(prev_content.encode("utf-8")) + len(incoming_content.encode("utf-8")) > self._max_merged_bytes: + return False + prev_data["content"] = f"{prev_content}{incoming_content}" + + if incoming_finish is not None: + prev_data["finish_reason"] = incoming_finish + + return bool(incoming_content) or (incoming_finish is not None) + + def _try_merge_tool_call_chunks(self, prev_event: Dict[str, Any], incoming: Dict[str, Any]) -> bool: + prev_data = prev_event.get("data") + if not isinstance(prev_data, dict): + return False + + merge_keys = ("thread_id", "agent", "id") + if any(prev_data.get(k) != incoming.get(k) for k in merge_keys): + return False + + prev_chunks = prev_data.get("tool_call_chunks") + incoming_chunks = incoming.get("tool_call_chunks") + if not (isinstance(prev_chunks, list) and isinstance(incoming_chunks, list)): + return False + if len(prev_chunks) != 1 or len(incoming_chunks) != 1: + return False + + prev_chunk = prev_chunks[0] + incoming_chunk = incoming_chunks[0] + if not (isinstance(prev_chunk, dict) and isinstance(incoming_chunk, dict)): + return False + + prev_call_id = prev_chunk.get("id") + incoming_call_id = incoming_chunk.get("id") + if prev_call_id is not None or incoming_call_id is not None: + if prev_call_id != incoming_call_id: + return False + else: + if prev_chunk.get("index") != incoming_chunk.get("index"): + return False + + prev_args = prev_chunk.get("args") or "" + incoming_args = incoming_chunk.get("args") or "" + if not isinstance(prev_args, str) or not isinstance(incoming_args, str): + return False + + if incoming_args: + if len(prev_args.encode("utf-8")) + len(incoming_args.encode("utf-8")) > self._max_merged_bytes: + return False + prev_chunk["args"] = f"{prev_args}{incoming_args}" + + return bool(incoming_args) diff --git a/src/server/services/runs/executor.py b/src/server/services/runs/executor.py index 199c5f9b6..e52e183e3 100644 --- a/src/server/services/runs/executor.py +++ b/src/server/services/runs/executor.py @@ -44,7 +44,6 @@ teardown, ) from src.server.services.runs.finalization import FinalizeVerdict -from src.server.services.runs.sse_producer import model_call_failure from src.server.services.runs.stream_writer import TransportLostError from src.server.services.runs.teardown import StopTeardown from src.server.dependencies.usage_limits import release_burst_slot @@ -946,8 +945,6 @@ async def _finalize_run( transaction carrying usage rows with it, so a persist failure can no longer be swallowed into a zombie-ACTIVE turn. """ - from src.server.services.runs.coordinator import RunOutcome - key = (thread_id, run_id) async with self.task_lock: task_info = self.executions.get(key) @@ -964,7 +961,13 @@ async def _finalize_run( workspace_id = metadata.get("workspace_id") user_id = metadata.get("user_id") - status, phase, interrupt_reason, error = await finalization.classify_outcome( + ( + status, + phase, + interrupt_reason, + pause_message, + error, + ) = await finalization.classify_outcome( kind, handler=handler, graph=graph, thread_id=thread_id, error=error ) # The stop drain's events the finalize archives, copied: a teardown a @@ -972,35 +975,21 @@ async def _finalize_run( # lanes it adds then are archived once their writers settle instead. carried_events = list(stop.events) if stop and status == "cancelled" else [] - ( - execution_time, - per_call_records, - tool_usage, - sse_events, - persist_metadata, - ) = await finalization.assemble_finalize_artifacts( + outcome = await finalization.assemble_finalize_artifacts( key, metadata=metadata, status=status, phase=phase, + interrupt_reason=interrupt_reason, + pause_message=pause_message, + error=error, handler=handler, cancelled_by_user=cancelled_by_user, workspace_id=workspace_id, user_id=user_id, stop_events=carried_events, + exc=exc, ) - if exc is not None: - persist_metadata = { - **persist_metadata, - **model_call_failure( - exc, - getattr( - getattr(handler, "agent_config", None), - "credential_source", - None, - ), - ), - } # ---- the single terminal transition ---- # verdict.applied gates every terminal business effect below: losers @@ -1011,16 +1000,6 @@ async def _finalize_run( # needs-input wake, watch clear) ride the finalize transaction # as durable outbox jobs — finalize_run derives them from the # row's START-stamped metadata, so no explicit factory here. - outcome = RunOutcome( - status=status, - interrupt_reason=interrupt_reason, - metadata=persist_metadata, - errors=[error] if error else None, - execution_time=execution_time, - sse_events=sse_events, - per_call_records=per_call_records, - tool_usage=tool_usage, - ) # I2 tail mode: subagent writers that survive the turn checkpoint # through the run's pinned session, so its release must wait for # them (N(root) drops at finalize either way). @@ -1072,7 +1051,7 @@ async def tail_drain() -> None: carried_events=carried_events, task_info=task_info, completion_callback=completion_callback, - execution_time=execution_time, + execution_time=outcome.execution_time, interrupt_reason=interrupt_reason, error=error, metadata=metadata, diff --git a/src/server/services/runs/finalization.py b/src/server/services/runs/finalization.py index 0995cf0f4..a612067a9 100644 --- a/src/server/services/runs/finalization.py +++ b/src/server/services/runs/finalization.py @@ -13,7 +13,13 @@ from typing import Literal, Optional from src.config.settings import get_checkpoint_flush_timeout -from src.server.contracts.status import is_user_stop +from src.server.contracts.status import ( + INTERRUPT_REASON_CREDIT_PAUSE, + classify_interrupts, + is_user_stop, +) +from src.server.database.runs.lifecycle import ProducerRecord, RunOutcome +from src.server.services.runs.stream_errors import model_call_failure from src.server.utils.persistence_utils import ( calculate_execution_time, get_sse_events_from_handler, @@ -53,12 +59,13 @@ async def classify_outcome( graph, thread_id: str, error: Optional[str], -) -> tuple[str, str, Optional[str], Optional[str]]: +) -> tuple[str, str, Optional[str], Optional[str], Optional[str]]: """In-band outcome classification -> (status, phase, interrupt_reason, - error). Interrupted-vs-completed comes from the streaming handler's - durability barrier; the timeout-prone aget_state probe survives only - for handler-less runs.""" + pause_message, error). Interrupted-vs-completed comes from the streaming + handler's durability barrier; the timeout-prone aget_state probe survives + only for handler-less runs.""" interrupt_reason: Optional[str] = None + pause_message: Optional[str] = None if kind == "cancelled": status, phase = "cancelled", "cancellation" elif kind == "failed": @@ -67,9 +74,10 @@ async def classify_outcome( if handler.interrupt_verified: status, phase = "interrupted", "interrupt" # Already classified by the handler's durability barrier - # (classify_interrupt_reason over the buffered payloads) — - # never respelled here. + # (classify_interrupts over the buffered payloads) — never + # respelled here. interrupt_reason = handler.interrupt_reason + pause_message = handler.pause_message else: # I8: a pause that never reached the checkpointer must not # advertise resumability. @@ -89,12 +97,8 @@ async def classify_outcome( timeout=get_checkpoint_flush_timeout(), ) if snapshot and snapshot.next: - from src.server.contracts.status import ( - classify_interrupt_reason, - ) - status, phase = "interrupted", "interrupt" - interrupt_reason = classify_interrupt_reason( + interrupt_reason, pause_message = classify_interrupts( intr for task in (snapshot.tasks or ()) for intr in (getattr(task, "interrupts", ()) or ()) @@ -105,7 +109,7 @@ async def classify_outcome( f"({thread_id}, ...)", exc_info=True, ) - return status, phase, interrupt_reason, error + return status, phase, interrupt_reason, pause_message, error async def assemble_finalize_artifacts( @@ -114,14 +118,19 @@ async def assemble_finalize_artifacts( metadata: dict, status: str, phase: str, + interrupt_reason: Optional[str], + pause_message: Optional[str], + error: Optional[str], handler, cancelled_by_user: bool, workspace_id: Optional[str], user_id: Optional[str], stop_events: Optional[list[dict]] = None, -) -> tuple: - """Build everything the finalize CAS archives: usage records, the - (possibly stop-reconciled) sse_events, and the persist metadata. + exc: Optional[BaseException] = None, +) -> RunOutcome: + """Build everything the finalize CAS writes: usage records, the + (possibly stop-reconciled) sse_events, the persist metadata and the + producer's record. ``stop_events`` are the subagent events a stop teardown drained from the tasks its kill evicted, which the caller passes only for a cancelled run. @@ -165,6 +174,21 @@ async def assemble_finalize_artifacts( persist_metadata[extra] = metadata[extra] if status == "cancelled": persist_metadata["cancelled_by_user"] = cancelled_by_user + if exc is not None: + persist_metadata.update( + model_call_failure( + exc, + getattr( + getattr(handler, "agent_config", None), "credential_source", None + ), + ) + ) + # The pause's own words, which settling the run's automation relays, so + # that reader never has to open the turn's event archive for them. A + # pause without words is stamped null: the interrupts that classified it + # are the ones the archive holds, so it has none either. + if interrupt_reason == INTERRUPT_REASON_CREDIT_PAUSE: + persist_metadata["credit_pause_message"] = pause_message # Steering inputs archive on the owning response (v4 identity model: # steering = no run, no turn). Replaces the old backfill that # fabricated query rows for orphan turn indexes. @@ -191,7 +215,30 @@ async def assemble_finalize_artifacts( f"[Finalize] artifact hook failed for {key}", exc_info=True, ) - return execution_time, per_call_records, tool_usage, sse_events, persist_metadata + return RunOutcome( + status=status, + interrupt_reason=interrupt_reason, + metadata=persist_metadata, + errors=[error] if error else None, + execution_time=execution_time, + sse_events=sse_events, + per_call_records=per_call_records, + tool_usage=tool_usage, + record=_producer_record(key, handler, stop_events), + ) + + +def _producer_record(key: tuple, handler, stop_events) -> Optional[ProducerRecord]: + """The producer's record, or None when it cannot be read: a raise here + would leave the row in_progress, where a missing record only costs the + row its replay facts.""" + if handler is None: + return None + try: + return handler.record(stop_events) + except Exception: + logger.warning(f"[Finalize] producer record unavailable for {key}", exc_info=True) + return None async def drive_finalize_cas( diff --git a/src/server/services/runs/reasoning_facts.py b/src/server/services/runs/reasoning_facts.py new file mode 100644 index 000000000..fa15a94d4 --- /dev/null +++ b/src/server/services/runs/reasoning_facts.py @@ -0,0 +1,61 @@ +"""How long the main agent thought, kept for replay. + +The checkpoint keeps a model call's reasoning but not its duration. A call +streams under a chunk id the checkpoint never sees (``lc_run--…``), so its +closes wait until the node's update names the message it committed, and a +call that commits nothing (an attempt a retry replaced) is dropped with them. +""" + +from typing import Optional + +from langchain_core.messages import AIMessage + + +class ReasoningDurations: + """Reasoning closes attributed to the main agent's committed messages. + + Only the newest main-agent chunk's closes are ever committed, so a close + from a subagent, a tool's inner model or a stop just waits under its own + id until the next commit drops it. + """ + + def __init__(self) -> None: + self._committed: dict[str, list[int]] = {} + self._pending: dict[str, list[int]] = {} + # (langgraph_node, chunk id) of the newest main-agent model chunk. + self._last_root_chunk: Optional[tuple[str, str]] = None + + def saw_root_chunk(self, node: str, chunk_id: str) -> None: + self._last_root_chunk = (node, chunk_id) + + def closed(self, message_id: str, elapsed_ms: int) -> None: + self._pending.setdefault(message_id, []).append(elapsed_ms) + + def node_updated(self, update: dict) -> None: + """Attribute the waiting closes to the message the node committed: + the one its newest model call streamed. Closes from an earlier call + in the same node belonged to an attempt it replaced.""" + if self._last_root_chunk is None: + return + node, chunk_id = self._last_root_chunk + if node not in update: + return + pending, self._pending = self._pending, {} + self._last_root_chunk = None + closes = pending.get(chunk_id) + output = update.get(node) + if not closes or not isinstance(output, dict): + return + messages = output.get("messages") + messages = getattr(messages, "value", messages) # Overwrite + if not isinstance(messages, (list, tuple)): + messages = [messages] + committed = next( + (m for m in reversed(messages) if isinstance(m, AIMessage) and m.id), + None, + ) + if committed is not None: + self._committed.setdefault(committed.id, []).extend(closes) + + def as_facts(self) -> dict[str, list[int]]: + return {k: list(v) for k, v in self._committed.items()} diff --git a/src/server/services/runs/recovery.py b/src/server/services/runs/recovery.py index bbbf68ef6..75d896e8e 100644 --- a/src/server/services/runs/recovery.py +++ b/src/server/services/runs/recovery.py @@ -24,6 +24,7 @@ from src.config.settings import get_recovery_scan_interval from src.server.database.runs import credit_ledger +from src.server.database.runs import subagent_repair from src.server.database.runs import subagent_runs as sr_db from src.server.database.runs import lifecycle as tl_db @@ -396,11 +397,12 @@ async def _stamp_task_retention( async def _recover_run( self, run: Dict[str, Any], run_id: str, thread_id: str ) -> bool: - status, interrupt_reason, errors, checkpoint_id = await self._classify( - run, run_id, thread_id + status, interrupt_reason, errors, checkpoint_id, stamped = ( + await self._classify(run, run_id, thread_id) ) sse_events, quality = await self._salvage_stream(thread_id, run_id) + from src.server.database.runs.lifecycle import ProducerRecord from src.server.services.runs.coordinator import RunCoordinator, RunOutcome # The funnel owns the entire post-CAS tail — projection refresh, @@ -414,9 +416,22 @@ async def _recover_run( RunOutcome( status=status, interrupt_reason=interrupt_reason, - metadata={"recovery": "scanner", "recovery_quality": quality}, + metadata={ + "recovery": "scanner", + "recovery_quality": quality, + **stamped, + }, errors=errors, sse_events=sse_events, + # The salvage is the only record of the dead stream. Nothing + # timed its thinking, so replay reads that from the events. + record=ProducerRecord( + provenance=[ + e + for e in sse_events or () + if isinstance(e, dict) and e.get("event") == "provenance" + ] + ), ), checkpoint_id=checkpoint_id, error_frame={ @@ -438,8 +453,11 @@ async def _recover_run( async def _classify( self, run: Dict[str, Any], run_id: str, thread_id: str - ) -> Tuple[str, Optional[str], Optional[List[str]], Optional[str]]: - """(status, interrupt_reason, errors, checkpoint_id) for a dead run. + ) -> Tuple[ + str, Optional[str], Optional[List[str]], Optional[str], Dict[str, Any] + ]: + """(status, interrupt_reason, errors, checkpoint_id, metadata) for a + dead run. `interrupted` demands a pending ``__interrupt__`` on a checkpoint CREATED BY this run (CheckpointMetadata.run_id, stamped from the @@ -480,25 +498,30 @@ async def _classify( if run.get("cancel_requested_at"): # The CAS adopts cancelled from the durable intent regardless; # requesting it just keeps the log honest. - return "cancelled", None, None, tip_id + return "cancelled", None, None, tip_id, {} if pending_interrupts: from src.server.contracts.status import ( - classify_interrupt_reason, + INTERRUPT_REASON_CREDIT_PAUSE, + classify_interrupts, ) - return ( - "interrupted", - classify_interrupt_reason(pending_interrupts), - None, - tip_id, + reason, pause_message = classify_interrupts(pending_interrupts) + # Stamped as the owner's finalize stamps it, so settling the + # run's automation never needs the salvaged archive for it. + stamped = ( + {"credit_pause_message": pause_message} + if reason == INTERRUPT_REASON_CREDIT_PAUSE + else {} ) + return "interrupted", reason, None, tip_id, stamped return ( "error", None, ["worker_lost: no live executor holds this run's writer fence"], tip_id, + {}, ) # --------------------------------------------------------------- salvage @@ -513,7 +536,7 @@ async def _salvage_stream( persistence so chunk merging matches owner-persisted turns. """ try: - from src.server.services.runs.sse_producer import StreamEventAccumulator + from src.server.services.runs.event_archive import StreamEventAccumulator from src.server.services.runs.stream_writer import stream_key from src.utils.cache.redis_cache import get_cache_client @@ -614,7 +637,7 @@ async def heal_task_chains(self) -> Dict[str, int]: every cycle costs nothing on a healthy ledger. """ try: - healed = await sr_db.repair_dangling_task_chains() + healed = await subagent_repair.repair_dangling_task_chains() if healed["rewound"] or healed["deleted"]: logger.warning( f"[RecoveryScanner] healed task chains: " diff --git a/src/server/services/runs/sse_producer.py b/src/server/services/runs/sse_producer.py index d135041bf..b94cc66d1 100644 --- a/src/server/services/runs/sse_producer.py +++ b/src/server/services/runs/sse_producer.py @@ -9,9 +9,7 @@ import copy import json import logging -import re import time -from datetime import datetime from typing import Any, AsyncGenerator, Dict, List, Optional, Set, Tuple, cast import json_repair @@ -23,13 +21,21 @@ is_thinking_status_signal, ) from src.server.utils.pg_sanitize import finite_json_dumps +from src.server.contracts.status import classify_interrupts +from src.server.database.provenance import MAIN_PROVENANCE_LANE +from src.server.database.runs.lifecycle import ProducerRecord +from src.server.services.runs.credit_usage import build_credit_usage_data +from src.server.services.runs.event_archive import ( + MERGED_STREAM_CHUNK_MAX_BYTES_DEFAULT, + StreamEventAccumulator, +) +from src.server.services.runs.reasoning_facts import ReasoningDurations +from src.server.services.runs.stream_errors import error_event_data from src.server.utils.text_phase import TextPhaseTracker from src.llms.content_utils import extract_reasoning_summary_index -from src.config import settings as app_settings from src.config.settings import ( get_workflow_timeout, is_sse_event_log_enabled, - get_merged_chunk_max_bytes, ) from opentelemetry.trace import Status, StatusCode from src.observability.tracing import tracer as _otel_tracer @@ -45,8 +51,6 @@ WORKFLOW_TIMEOUT = get_workflow_timeout() # seconds SSE_EVENT_LOG_ENABLED = is_sse_event_log_enabled() -MERGED_STREAM_CHUNK_MAX_BYTES_DEFAULT = get_merged_chunk_max_bytes() - DEFAULT_TOKEN_THRESHOLD = 120000 @@ -66,38 +70,6 @@ def resolve_token_threshold(agent_config=None) -> int: return cfg.compaction.token_threshold -def build_credit_usage_data( - thread_id: str, - token_usage: dict, - total_credits: float, - timestamp: Optional[str] = None, -) -> Dict[str, Any]: - """Aggregate per-model usage into the ``credit_usage`` wire payload. - - Intentionally omits USD costs and model names (hidden from the client). - Shared by the live handler and table-sourced replay so both wires carry - the same shape. - """ - total_input_tokens = 0 - total_output_tokens = 0 - total_tokens = 0 - for usage in (token_usage or {}).get("by_model", {}).values(): - total_input_tokens += usage.get("input_tokens", 0) - total_output_tokens += usage.get("output_tokens", 0) - total_tokens += usage.get("total_tokens", 0) - - return { - "thread_id": thread_id, - "tokens": { - "input_tokens": total_input_tokens, - "output_tokens": total_output_tokens, - "total_tokens": total_tokens, - }, - "total_credits": round(total_credits, 2), - "timestamp": timestamp or datetime.now().isoformat(), - } - - def _parse_tool_args( raw: str, ) -> Tuple[Optional[Dict[str, Any]], str, str]: @@ -138,275 +110,6 @@ def _parse_tool_args( return None, err_repr, err_window -# --------------------------------------------------------------------------- -# Stream error classification -# --------------------------------------------------------------------------- -# -# Chat-stream failures fall into two buckets and the user-facing remedy is -# different for each: -# -# - ``upstream`` — the LLM provider we called returned an error (their -# server 500'd, the key is rejected, rate-limited, etc). -# User should check their key / plan / provider status. -# - ``internal`` — our own pipeline failed (middleware bug, our DB, a -# schema mismatch in the payload we built). User can't -# do anything; we should log loudly and show a generic -# retry message. -# -# Classification is a module-prefix check with exception-chain walking — any -# exception in the chain sourced from a known provider SDK flips the whole -# failure to ``upstream``. - -# Provider SDK and LangChain-wrapper module prefixes. Any exception in the -# cause chain whose ``__module__`` matches one of these flips the whole -# failure to ``upstream``. Keep in sync with the SDKs wired up in -# ``src/llms/llm.py`` — missing a prefix means the user sees "our service -# failed" for what's really a provider error. -# -# ``httpx`` is in this list as a last-resort catch: a bare httpx exception -# that reaches the stream error handler has almost always come from the -# LangChain call path (SDKs raise via httpx). If our own service calls -# (credit checks, workspace manager) ever start raising bare httpx errors to -# the stream path we should wrap them in a distinct exception type before -# they bubble; classification is a UI hint, not a diagnostic source of truth. -_UPSTREAM_MODULE_PREFIXES: tuple[str, ...] = ( - # Raw provider SDKs - "anthropic", - "openai", - "google.api_core", - "google.genai", - "google.generativeai", - "cohere", - "httpx", - # Our own wrappers around a provider stream. They quote the provider and - # are raised fresh rather than chained, so nothing below would match them. - "src.llms.extension", - # LangChain wrappers — their exceptions may not chain through the raw SDK - # when the wrapper normalizes errors, so match them directly. - "langchain_openai", - "langchain_anthropic", - "langchain_deepseek", - "langchain_qwq", - "langchain_google_genai", - "langchain_google_vertexai", - "langchain_mistralai", - "langchain_together", - "langchain_groq", - "groq", -) - -_STATUS_CODE_RE = re.compile(r"\b([45]\d{2})\b") - - -def _parse_status_from_message(text: str) -> Optional[int]: - match = _STATUS_CODE_RE.search(text) - return int(match.group(1)) if match else None - - -# Statuses meaning the provider refused the request we sent: bad shape, -# unsupported content, oversized payload. Never the caller's credential and -# never the provider's health, so the only honest hint is to move to a model -# that accepts the request. 401/403/404 are absent because they have their own -# branches; they are about access, not about the request body. -_REQUEST_REFUSED_STATUSES: frozenset[int] = frozenset({400, 405, 413, 422}) - -# Hints that ask the user to go fix a credential. Only truthful when the user -# is the one holding it. -_CREDENTIAL_HINTS: frozenset[str] = frozenset({"api_key", "model_access"}) - - -def user_owns_credential(credential_source: Any) -> bool: - """Whether the reader of the error holds the credential that ran the call. - - Both halves are load-bearing, because ``platform`` means opposite things in - the two host modes: in OSS the key is the operator's own ``.env`` entry and - the operator is the user, while on the hosted service the user has no key at - all and "check your API key" sends them to a page that is not the cause. - An unknown source fails closed, since a wrong credential hint is worse than - a missing one. - """ - if app_settings.HOST_MODE == "oss": - return True - return str(credential_source) in ("oauth", "byok") - - -def find_resilience_trace(exc: BaseException) -> Optional[Dict[str, Any]]: - """Find the attempt trace attached by ``ModelResilienceMiddleware``. - - The middleware sets ``__model_resilience__`` (see ``RESILIENCE_TRACE_ATTR`` - in ``src/ptc_agent/agent/middleware/model_resilience.py``) on the primary - model's exception before re-raising it. Walks the cause chain defensively - in case a wrapper exception ends up on top. - """ - seen: set[int] = set() - current: Optional[BaseException] = exc - while current is not None and id(current) not in seen: - seen.add(id(current)) - trace = getattr(current, "__model_resilience__", None) - if isinstance(trace, dict): - return trace - current = current.__cause__ or current.__context__ - return None - - -def model_call_failure( - exc: BaseException, credential_source: Any -) -> Dict[str, Any]: - """Ledger metadata for a run that failed on its model call; ``{}`` when - the failure was anything else. - - The resilience trace is the signal, not the provider SDK prefix - ``classify_stream_exception`` matches: the middleware attaches it only to - a model call's exception, while an httpx 401 from a tool or data API - would pass for a rejected model key. The primary's status is the one its - credential earned, and that credential is the run's own. - """ - trace = find_resilience_trace(exc) - attempted = trace.get("attempted_models") if trace else None - if not (isinstance(attempted, list) and attempted and isinstance(attempted[0], dict)): - return {} - status = attempted[0].get("status_code") - if not isinstance(status, int): - return {} - return { - "error_status_code": status, - "error_credential_owned": user_owns_credential(credential_source), - } - - -def classify_stream_exception(exc: BaseException) -> Dict[str, Any]: - """Classify a chat-stream exception as ``upstream`` or ``internal``. - - Walks ``__cause__`` / ``__context__`` so a wrapped provider error (e.g. - a LangChain exception caused by ``anthropic.InternalServerError``) is - still recognized as upstream. Returns a dict with ``kind``, - ``status_code`` (when carried on the exception or parseable from its - message), and ``provider_module`` (the matched SDK prefix, or None). - """ - seen: set[int] = set() - current: Optional[BaseException] = exc - while current is not None and id(current) not in seen: - seen.add(id(current)) - module = getattr(type(current), "__module__", "") or "" - for prefix in _UPSTREAM_MODULE_PREFIXES: - if module == prefix or module.startswith(prefix + "."): - status = getattr(current, "status_code", None) - if not isinstance(status, int): - status = _parse_status_from_message(str(current)) - return { - "kind": "upstream", - "status_code": status if isinstance(status, int) else None, - "provider_module": prefix, - } - current = current.__cause__ or current.__context__ - - return {"kind": "internal", "status_code": None, "provider_module": None} - - -class StreamEventAccumulator: - """Accumulates and merges token-level SSE events for persistence.""" - - def __init__(self, max_merged_bytes: int = MERGED_STREAM_CHUNK_MAX_BYTES_DEFAULT): - self._max_merged_bytes = max_merged_bytes - self._events: List[Dict[str, Any]] = [] - - def get_events(self) -> List[Dict[str, Any]]: - return copy.deepcopy(self._events) - - def add(self, event_type: str, data: Dict[str, Any]) -> None: - if not isinstance(data, dict): - return - - incoming = copy.deepcopy(data) - - if not self._events: - self._events.append({"event": event_type, "data": incoming}) - return - - prev = self._events[-1] - if prev.get("event") != event_type: - self._events.append({"event": event_type, "data": incoming}) - return - - if event_type == "message_chunk" and self._try_merge_message_chunk(prev, incoming): - return - - if event_type == "tool_call_chunks" and self._try_merge_tool_call_chunks(prev, incoming): - return - - self._events.append({"event": event_type, "data": incoming}) - - def _try_merge_message_chunk(self, prev_event: Dict[str, Any], incoming: Dict[str, Any]) -> bool: - prev_data = prev_event.get("data") - if not isinstance(prev_data, dict): - return False - - if incoming.get("content_type") == "reasoning_signal": - return False - if prev_data.get("content_type") == "reasoning_signal": - return False - - merge_keys = ("thread_id", "agent", "id", "role", "content_type", "phase") - if any(prev_data.get(k) != incoming.get(k) for k in merge_keys): - return False - - prev_content = prev_data.get("content") or "" - incoming_content = incoming.get("content") or "" - incoming_finish = incoming.get("finish_reason") - - if incoming_content: - if len(prev_content.encode("utf-8")) + len(incoming_content.encode("utf-8")) > self._max_merged_bytes: - return False - prev_data["content"] = f"{prev_content}{incoming_content}" - - if incoming_finish is not None: - prev_data["finish_reason"] = incoming_finish - - return bool(incoming_content) or (incoming_finish is not None) - - def _try_merge_tool_call_chunks(self, prev_event: Dict[str, Any], incoming: Dict[str, Any]) -> bool: - prev_data = prev_event.get("data") - if not isinstance(prev_data, dict): - return False - - merge_keys = ("thread_id", "agent", "id") - if any(prev_data.get(k) != incoming.get(k) for k in merge_keys): - return False - - prev_chunks = prev_data.get("tool_call_chunks") - incoming_chunks = incoming.get("tool_call_chunks") - if not (isinstance(prev_chunks, list) and isinstance(incoming_chunks, list)): - return False - if len(prev_chunks) != 1 or len(incoming_chunks) != 1: - return False - - prev_chunk = prev_chunks[0] - incoming_chunk = incoming_chunks[0] - if not (isinstance(prev_chunk, dict) and isinstance(incoming_chunk, dict)): - return False - - prev_call_id = prev_chunk.get("id") - incoming_call_id = incoming_chunk.get("id") - if prev_call_id is not None or incoming_call_id is not None: - if prev_call_id != incoming_call_id: - return False - else: - if prev_chunk.get("index") != incoming_chunk.get("index"): - return False - - prev_args = prev_chunk.get("args") or "" - incoming_args = incoming_chunk.get("args") or "" - if not isinstance(prev_args, str) or not isinstance(incoming_args, str): - return False - - if incoming_args: - if len(prev_args.encode("utf-8")) + len(incoming_args.encode("utf-8")) > self._max_merged_bytes: - return False - prev_chunk["args"] = f"{prev_args}{incoming_args}" - - return bool(incoming_args) - - class RunSSEProducer: """LangGraph workflow SSE producer with reasoning lifecycle tracking, tool-call dedup, and content normalisation.""" @@ -448,9 +151,13 @@ def __init__( # Track reasoning status per agent for lifecycle management self.reasoning_active: Set[str] = set() # Monotonic clock at each agent's open reasoning block, so its closing - # signal can carry how long the model thought. Stored frames replay - # verbatim, so the duration survives a reload without a schema change. + # signal can carry how long the model thought. self._reasoning_started_at: Dict[str, float] = {} + # How long the main agent thought, for replay. + self._reasoning = ReasoningDurations() + # The main agent's provenance records as emitted, for the finalize to + # write without reading them back out of the event archive. + self.provenance_events: List[Dict[str, Any]] = [] # Track reasoning block index per agent to detect block transitions # When index changes (e.g., 0→1), a separator (\n\n) is needed between blocks @@ -523,6 +230,8 @@ def __init__( # resume from is the I8 lie). self.interrupt_verified: Optional[bool] = None self.interrupt_reason: Optional[str] = None + # A credit pause's denial, which the finalize stamps on the row. + self.pause_message: Optional[str] = None # ---- Stop reconciliation state (decision T3-A) ------------------- # Open artifacts whose last-emitted status was not terminal. Keyed by @@ -748,6 +457,9 @@ async def stream_workflow( prov_data = self._resolve_provenance_event( event_data, agent_from_stream ) + self.provenance_events.append( + {"event": "provenance", "data": dict(prov_data)} + ) yield self._format_sse_event("provenance", prov_data) continue @@ -890,8 +602,11 @@ async def stream_workflow( # State updates (stream_mode="updates") carry no SSE payloads of # their own — the mode stays subscribed because interrupts - # arrive through it (handled above via "__interrupt__"). + # arrive through it (handled above via "__interrupt__"). A + # main-graph node's update names the message it committed. if stream_mode == "updates": + if not agent_from_stream and isinstance(event_data, dict): + self._reasoning.node_updated(event_data) continue # Process message chunks (stream_mode="messages") @@ -933,6 +648,16 @@ async def stream_workflow( f"[RAW_REASONING] agent={agent_name} reasoning_content={reasoning_raw}" ) + node = (message_metadata or {}).get("langgraph_node") + if ( + not agent_from_stream + and not is_compaction_chunk + and node != "tools" + and isinstance(message_chunk, (AIMessage, AIMessageChunk)) + and message_chunk.id + ): + self._reasoning.saw_root_chunk(str(node), message_chunk.id) + # Process the message chunk async for event in self._process_message_chunk( message_chunk, @@ -949,7 +674,11 @@ async def stream_workflow( if self._pending_interrupts: self.interrupt_verified = await self._verify_interrupt_durable(graph) if self.interrupt_verified: - self.interrupt_reason = self._derive_interrupt_reason() + self.interrupt_reason, self.pause_message = classify_interrupts( + intr + for pending in self._pending_interrupts + for intr in pending.get("__interrupt__", ()) + ) for pending in self._pending_interrupts: interrupt_event = self._handle_interrupt(pending) if interrupt_event: @@ -1097,16 +826,6 @@ async def _verify_interrupt_durable(self, graph) -> bool: ) return False - def _derive_interrupt_reason(self) -> Optional[str]: - """Classify the buffered interrupt into the ledger's reason column.""" - from src.server.contracts.status import classify_interrupt_reason - - return classify_interrupt_reason( - intr - for event_data in self._pending_interrupts - for intr in event_data.get("__interrupt__", ()) - ) - def _handle_interrupt(self, event_data: dict) -> Optional[str]: """Format an ``__interrupt__`` event as an SSE string.""" interrupt_obj = event_data["__interrupt__"][0] @@ -1161,7 +880,7 @@ def _resolve_provenance_event( if namespace_tuple: agent = self._extract_agent_name(namespace_tuple, event_data) else: - agent = "main" + agent = MAIN_PROVENANCE_LANE prov_data = { key: value for key, value in event_data.items() if key != "type" } @@ -1788,8 +1507,24 @@ def _format_reasoning_signal( started = self._reasoning_started_at.pop(agent_name, None) if started is not None: data["elapsed_ms"] = int((time.monotonic() - started) * 1000) + if not is_compaction: + self._reasoning.closed(message_id, data["elapsed_ms"]) return self._format_sse_event(event_type, data) + def record(self, stop_events: Optional[List[Dict[str, Any]]] = None) -> ProducerRecord: + """What the finalize stores beside the status that the checkpoint + does not keep. ``stop_events`` are the lanes a stop drained, whose + provenance reaches the row only through this finalize.""" + return ProducerRecord( + provenance=self.provenance_events + + [ + e + for e in stop_events or () + if isinstance(e, dict) and e.get("event") == "provenance" + ], + reasoning_ms=self._reasoning.as_facts(), + ) + def _format_sse_event(self, event_type: str, data: dict[str, Any], *, accumulate: bool = True) -> str: """ Format data as SSE (Server-Sent Events) string with sequence numbering. @@ -1841,110 +1576,17 @@ def format_error_event( *, exc: Optional[BaseException] = None, ) -> str: - """Format an error event as SSE string. - - When ``exc`` is passed the event carries ``error_kind`` (``upstream`` - or ``internal``), ``status_code`` (when available), and ``hints`` for - the frontend to render user-actionable guidance. The legacy ``error`` - and ``message`` fields stay so older clients keep working. - - Hints are chosen by status and then filtered by who holds the credential - (see ``user_owns_credential``), so a platform-billed turn never asks the - user to check a key they do not have. - - Args: - error_message: Raw error text (usually ``str(exc)``). - exc: The exception itself — enables classification. Optional to - keep the legacy signature working for paths that only have - a prebuilt message. - - Returns: - SSE-formatted error event. - """ - data: Dict[str, Any] = { - "thread_id": self.thread_id, - "error": _sanitize_error_text(error_message), - "message": "An error occurred during processing", - } - if exc is not None: - info = classify_stream_exception(exc) - trace = find_resilience_trace(exc) - if trace is not None and info["kind"] == "internal": - # The trace proves the failure happened inside a model call. - # A generic wrapper exception (no recognizable SDK module in - # the chain) must not demote it to "internal" — the frontend - # routes internal errors to a generic banner that drops the - # model / attempted-models context. - primary_status: Any = None - attempted_raw = trace.get("attempted_models") - if ( - isinstance(attempted_raw, list) - and attempted_raw - and isinstance(attempted_raw[0], dict) - ): - primary_status = attempted_raw[0].get("status_code") - info = { - "kind": "upstream", - "status_code": primary_status - if isinstance(primary_status, int) - else None, - "provider_module": None, - } - data["error_kind"] = info["kind"] - if info["status_code"] is not None: - data["status_code"] = info["status_code"] - if info["provider_module"]: - data["provider_module"] = info["provider_module"] - if info["kind"] == "upstream": - # Order matters — frontend renders the hints as a list, so the - # most relevant hint for this status goes first. 5xx/429 are - # provider outages, not the user's credentials; showing - # "check your API key" first on a 503 is misleading. - status = info.get("status_code") - if status in (401, 403): - hints = ["api_key", "model_access", "try_another_model"] - elif status == 404: - hints = ["model_access", "try_another_model"] - elif status in _REQUEST_REFUSED_STATUSES: - hints = ["try_another_model"] - elif status == 429 or (isinstance(status, int) and status >= 500): - hints = ["provider_status", "try_another_model"] - else: - # No status (network error) — could be anything; show all. - hints = [ - "api_key", - "model_access", - "provider_status", - "try_another_model", - ] - # The status says what failed; the credential says who can act - # on it. Drop the hints that would send a user to fix a key - # they do not hold. Every branch above keeps - # "try_another_model", so the list never empties. - if not user_owns_credential( - getattr(self.agent_config, "credential_source", None) - ): - hints = [h for h in hints if h not in _CREDENTIAL_HINTS] - data["hints"] = hints - if trace is not None: - primary_model = trace.get("model") - if isinstance(primary_model, str) and primary_model: - data["model"] = primary_model - attempted = trace.get("attempted_models") - if isinstance(attempted, list): - data["attempted_models"] = [ - { - "model": entry.get("model"), - "error": _sanitize_error_text( - str(entry.get("error") or "") - ), - "status_code": entry.get("status_code"), - "attempts": entry.get("attempts"), - } - for entry in attempted - if isinstance(entry, dict) - ] - return self._format_sse_event("error", data) + """Format an error event as SSE string; ``exc`` enables the + classification ``error_event_data`` describes.""" + return self._format_sse_event( + "error", + error_event_data( + self.thread_id, + error_message, + exc=exc, + credential_source=getattr(self.agent_config, "credential_source", None), + ), + ) def _format_credit_usage_event( self, diff --git a/src/server/services/runs/stream_errors.py b/src/server/services/runs/stream_errors.py new file mode 100644 index 000000000..7a357d8d2 --- /dev/null +++ b/src/server/services/runs/stream_errors.py @@ -0,0 +1,269 @@ +"""Chat-stream failure classification, shared by the producer's error frame +and the ledger metadata a failed model call leaves on its run.""" + +import re +from typing import Any, Dict, Optional + +from src.config import settings as app_settings +from src.server.utils.error_sanitization import sanitize_error_text + +# Chat-stream failures fall into two buckets and the user-facing remedy is +# different for each: +# +# - ``upstream`` — the LLM provider we called returned an error (their +# server 500'd, the key is rejected, rate-limited, etc). +# User should check their key / plan / provider status. +# - ``internal`` — our own pipeline failed (middleware bug, our DB, a +# schema mismatch in the payload we built). User can't +# do anything; we should log loudly and show a generic +# retry message. +# +# Classification is a module-prefix check with exception-chain walking — any +# exception in the chain sourced from a known provider SDK flips the whole +# failure to ``upstream``. + +# Provider SDK and LangChain-wrapper module prefixes. Any exception in the +# cause chain whose ``__module__`` matches one of these flips the whole +# failure to ``upstream``. Keep in sync with the SDKs wired up in +# ``src/llms/llm.py`` — missing a prefix means the user sees "our service +# failed" for what's really a provider error. +# +# ``httpx`` is in this list as a last-resort catch: a bare httpx exception +# that reaches the stream error handler has almost always come from the +# LangChain call path (SDKs raise via httpx). If our own service calls +# (credit checks, workspace manager) ever start raising bare httpx errors to +# the stream path we should wrap them in a distinct exception type before +# they bubble; classification is a UI hint, not a diagnostic source of truth. +_UPSTREAM_MODULE_PREFIXES: tuple[str, ...] = ( + # Raw provider SDKs + "anthropic", + "openai", + "google.api_core", + "google.genai", + "google.generativeai", + "cohere", + "httpx", + # Our own wrappers around a provider stream. They quote the provider and + # are raised fresh rather than chained, so nothing below would match them. + "src.llms.extension", + # LangChain wrappers — their exceptions may not chain through the raw SDK + # when the wrapper normalizes errors, so match them directly. + "langchain_openai", + "langchain_anthropic", + "langchain_deepseek", + "langchain_qwq", + "langchain_google_genai", + "langchain_google_vertexai", + "langchain_mistralai", + "langchain_together", + "langchain_groq", + "groq", +) + +_STATUS_CODE_RE = re.compile(r"\b([45]\d{2})\b") + + +def _parse_status_from_message(text: str) -> Optional[int]: + match = _STATUS_CODE_RE.search(text) + return int(match.group(1)) if match else None + + +# Statuses meaning the provider refused the request we sent: bad shape, +# unsupported content, oversized payload. Never the caller's credential and +# never the provider's health, so the only honest hint is to move to a model +# that accepts the request. 401/403/404 are absent because they have their own +# branches; they are about access, not about the request body. +_REQUEST_REFUSED_STATUSES: frozenset[int] = frozenset({400, 405, 413, 422}) + +# Hints that ask the user to go fix a credential. Only truthful when the user +# is the one holding it. +_CREDENTIAL_HINTS: frozenset[str] = frozenset({"api_key", "model_access"}) + + +def user_owns_credential(credential_source: Any) -> bool: + """Whether the reader of the error holds the credential that ran the call. + + Both halves are load-bearing, because ``platform`` means opposite things in + the two host modes: in OSS the key is the operator's own ``.env`` entry and + the operator is the user, while on the hosted service the user has no key at + all and "check your API key" sends them to a page that is not the cause. + An unknown source fails closed, since a wrong credential hint is worse than + a missing one. + """ + if app_settings.HOST_MODE == "oss": + return True + return str(credential_source) in ("oauth", "byok") + + +def find_resilience_trace(exc: BaseException) -> Optional[Dict[str, Any]]: + """Find the attempt trace attached by ``ModelResilienceMiddleware``. + + The middleware sets ``__model_resilience__`` (see ``RESILIENCE_TRACE_ATTR`` + in ``src/ptc_agent/agent/middleware/model_resilience.py``) on the primary + model's exception before re-raising it. Walks the cause chain defensively + in case a wrapper exception ends up on top. + """ + seen: set[int] = set() + current: Optional[BaseException] = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + trace = getattr(current, "__model_resilience__", None) + if isinstance(trace, dict): + return trace + current = current.__cause__ or current.__context__ + return None + + +def model_call_failure( + exc: BaseException, credential_source: Any +) -> Dict[str, Any]: + """Ledger metadata for a run that failed on its model call; ``{}`` when + the failure was anything else. + + The resilience trace is the signal, not the provider SDK prefix + ``classify_stream_exception`` matches: the middleware attaches it only to + a model call's exception, while an httpx 401 from a tool or data API + would pass for a rejected model key. The primary's status is the one its + credential earned, and that credential is the run's own. + """ + trace = find_resilience_trace(exc) + attempted = trace.get("attempted_models") if trace else None + if not (isinstance(attempted, list) and attempted and isinstance(attempted[0], dict)): + return {} + status = attempted[0].get("status_code") + if not isinstance(status, int): + return {} + return { + "error_status_code": status, + "error_credential_owned": user_owns_credential(credential_source), + } + + +def classify_stream_exception(exc: BaseException) -> Dict[str, Any]: + """Classify a chat-stream exception as ``upstream`` or ``internal``. + + Walks ``__cause__`` / ``__context__`` so a wrapped provider error (e.g. + a LangChain exception caused by ``anthropic.InternalServerError``) is + still recognized as upstream. Returns a dict with ``kind``, + ``status_code`` (when carried on the exception or parseable from its + message), and ``provider_module`` (the matched SDK prefix, or None). + """ + seen: set[int] = set() + current: Optional[BaseException] = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + module = getattr(type(current), "__module__", "") or "" + for prefix in _UPSTREAM_MODULE_PREFIXES: + if module == prefix or module.startswith(prefix + "."): + status = getattr(current, "status_code", None) + if not isinstance(status, int): + status = _parse_status_from_message(str(current)) + return { + "kind": "upstream", + "status_code": status if isinstance(status, int) else None, + "provider_module": prefix, + } + current = current.__cause__ or current.__context__ + + return {"kind": "internal", "status_code": None, "provider_module": None} + + +def error_event_data( + thread_id: str, + error_message: str, + *, + exc: Optional[BaseException] = None, + credential_source: Any = None, +) -> Dict[str, Any]: + """The ``error`` frame's data. + + With ``exc`` the frame carries ``error_kind`` (``upstream`` or + ``internal``), ``status_code`` when available, and ``hints`` for the + frontend to render user-actionable guidance; the legacy ``error`` and + ``message`` fields stay so older clients keep working. Hints are chosen by + status and then filtered by who holds the credential (see + ``user_owns_credential``), so a platform-billed turn never asks the user + to check a key they do not have. + """ + data: Dict[str, Any] = { + "thread_id": thread_id, + "error": sanitize_error_text(error_message), + "message": "An error occurred during processing", + } + if exc is not None: + info = classify_stream_exception(exc) + trace = find_resilience_trace(exc) + if trace is not None and info["kind"] == "internal": + # The trace proves the failure happened inside a model call. + # A generic wrapper exception (no recognizable SDK module in + # the chain) must not demote it to "internal" — the frontend + # routes internal errors to a generic banner that drops the + # model / attempted-models context. + primary_status: Any = None + attempted_raw = trace.get("attempted_models") + if ( + isinstance(attempted_raw, list) + and attempted_raw + and isinstance(attempted_raw[0], dict) + ): + primary_status = attempted_raw[0].get("status_code") + info = { + "kind": "upstream", + "status_code": primary_status + if isinstance(primary_status, int) + else None, + "provider_module": None, + } + data["error_kind"] = info["kind"] + if info["status_code"] is not None: + data["status_code"] = info["status_code"] + if info["provider_module"]: + data["provider_module"] = info["provider_module"] + if info["kind"] == "upstream": + # Order matters — frontend renders the hints as a list, so the + # most relevant hint for this status goes first. 5xx/429 are + # provider outages, not the user's credentials; showing + # "check your API key" first on a 503 is misleading. + status = info.get("status_code") + if status in (401, 403): + hints = ["api_key", "model_access", "try_another_model"] + elif status == 404: + hints = ["model_access", "try_another_model"] + elif status in _REQUEST_REFUSED_STATUSES: + hints = ["try_another_model"] + elif status == 429 or (isinstance(status, int) and status >= 500): + hints = ["provider_status", "try_another_model"] + else: + # No status (network error) — could be anything; show all. + hints = [ + "api_key", + "model_access", + "provider_status", + "try_another_model", + ] + # The status says what failed; the credential says who can act + # on it. Drop the hints that would send a user to fix a key + # they do not hold. Every branch above keeps + # "try_another_model", so the list never empties. + if not user_owns_credential(credential_source): + hints = [h for h in hints if h not in _CREDENTIAL_HINTS] + data["hints"] = hints + if trace is not None: + primary_model = trace.get("model") + if isinstance(primary_model, str) and primary_model: + data["model"] = primary_model + attempted = trace.get("attempted_models") + if isinstance(attempted, list): + data["attempted_models"] = [ + { + "model": entry.get("model"), + "error": sanitize_error_text( + str(entry.get("error") or "") + ), + "status_code": entry.get("status_code"), + "attempts": entry.get("attempts"), + } + for entry in attempted + if isinstance(entry, dict) + ] + return data diff --git a/src/server/services/runs/subagent_archive.py b/src/server/services/runs/subagent_archive.py index 449e63913..2a3f29594 100644 --- a/src/server/services/runs/subagent_archive.py +++ b/src/server/services/runs/subagent_archive.py @@ -161,12 +161,8 @@ async def persist_collected_events( """ if sandbox: try: - from src.server.services.persistence.image_capture import ( - capture_and_rewrite_images, - ) - - await capture_and_rewrite_images( - events, sandbox, thread_id=thread_id, workspace_id=workspace_id, + await _capture_lane_images( + events, sandbox, response_id, thread_id, workspace_id ) except Exception: logger.warning( @@ -201,6 +197,44 @@ async def persist_collected_events( return False +async def _capture_lane_images( + events: list[dict], + sandbox, + response_id: str, + thread_id: str, + workspace_id: str, +) -> None: + """Hook B: capture the images the lanes reference and rewrite them in the + events, recording each lane's path-to-URL map on the runs the turn + launched for it. + + Not on the thread's checkpoint, as the main turn's capture is: these + lanes settle after their turn did, and a record appended now would land + on whatever turn is newest, after the stored projection of this one. + """ + from src.server.services.persistence.image_capture import ( + capture_images, + sandbox_image_paths, + ) + from src.server.services.subagent_run_coordinator import record_lane_images + + by_task: dict[str, list[dict]] = {} + for event in events: + data = event.get("data") + agent = data.get("agent") if isinstance(data, dict) else None + if isinstance(agent, str) and agent.startswith("task:"): + by_task.setdefault(agent.removeprefix("task:"), []).append(event) + # Read before the capture: it rewrites the paths it resolves. + paths = {task_id: sandbox_image_paths(lane) for task_id, lane in by_task.items()} + path_to_url = await capture_images( + events, sandbox, thread_id=thread_id, workspace_id=workspace_id + ) + for task_id, lane_paths in paths.items(): + images = {p: path_to_url[p] for p in lane_paths if p in path_to_url} + if images: + await record_lane_images(thread_id, response_id, task_id, images) + + @dataclass class ArchiveBuffer: """One collection's archive writes: the events not yet landed, and diff --git a/src/server/services/runs/subagent_collection.py b/src/server/services/runs/subagent_collection.py index 14dc2b605..78d351fb8 100644 --- a/src/server/services/runs/subagent_collection.py +++ b/src/server/services/runs/subagent_collection.py @@ -197,9 +197,10 @@ def _owned_pending(tasks: list, response_id: str) -> dict[asyncio.Task, Any]: async def _replay_into(buf: ArchiveBuffer, task) -> None: - buf.withheld |= not await replay_owned_task_events( + replayed = await replay_owned_task_events( buf.thread_id, task, buf.response_id, buf.unwritten ) + buf.withheld |= not replayed async def _replay_settled(tasks: list, pending: dict, buf: ArchiveBuffer) -> None: @@ -555,8 +556,7 @@ async def collect_orphaned_subagent_results( try: # ``unwritten``: the turn collector's events that have not landed yet. buf = ArchiveBuffer( - response_id, thread_id, workspace_id, user_id, sandbox, - unwritten=unwritten, + response_id, thread_id, workspace_id, user_id, sandbox, unwritten ) _settle_finished(tasks, response_id) diff --git a/src/server/services/share_redaction.py b/src/server/services/share_redaction.py index c999ba935..3e8a7d479 100644 --- a/src/server/services/share_redaction.py +++ b/src/server/services/share_redaction.py @@ -25,7 +25,14 @@ # ``tool_call_chunks`` go whole: they are the arguments again, streamed in # pieces that often carry no call id to match, and no replay reads them. An # interrupt asks the owner, and no share renders or answers one. -_DROPPED_EVENTS = frozenset({"provenance", "tool_call_chunks", "interrupt"}) +# +# Exception text can name a provider's organization, an internal address or a +# sandbox, and scrubbing credentials from it removes none of those. A turn's +# closing ``error`` event is nothing but that text, so it goes whole; a run's +# lifecycle frame and a model fallback keep their status and lose the text. +# No share renders either. +_DROPPED_EVENTS = frozenset({"provenance", "tool_call_chunks", "interrupt", "error"}) +_ERROR_TEXT_EVENTS = frozenset({"workflow_lifecycle", "model_fallback"}) _PRIVATE_ARTIFACT_KEYS = (RECEIPT_KEY, "provenance") # The owner's own turn, and the resume that answers an approval names every # order it decided by the attempt's ledger id. A viewer cannot answer one and @@ -204,6 +211,8 @@ def event(self, event_type: str, data: dict[str, Any]) -> dict[str, Any] | None: # depth. sandbox_state is server-side runtime state. data = without_keys(data, _OWNER_ID_KEYS) data.pop("sandbox_state", None) + if event_type in _ERROR_TEXT_EVENTS: + data.pop("error", None) if event_type == "tool_calls": self._tool_calls(data) elif event_type == "tool_call_result": diff --git a/src/server/services/subagent_run_coordinator.py b/src/server/services/subagent_run_coordinator.py index bbd6a0c95..1460f6866 100644 --- a/src/server/services/subagent_run_coordinator.py +++ b/src/server/services/subagent_run_coordinator.py @@ -18,6 +18,7 @@ from ptc_agent.agent.middleware.background_subagent.registry import TaskRunRejected from src.server.contracts.status import REPORT_BACK_STATUSES +from src.server.database import replay_facts as facts_db from src.server.database.runs import subagent_runs as sr_db logger = logging.getLogger(__name__) @@ -33,6 +34,44 @@ def v2_stream_key(thread_id: str, task_run_id: str) -> str: return f"subagent:stream:{thread_id}:{task_run_id}" +def _refresh_launching_turn(thread_id: Any, parent_run_id: Any) -> None: + """Re-project the turn that launched a run, whose stored replay was keyed + on the run as it was. A run with no parent launched no turn.""" + if not parent_run_id: + return + try: + from src.server.services.history.replay.refresh import ( + schedule_replay_refresh, + ) + + schedule_replay_refresh(str(thread_id), str(parent_run_id)) + except Exception: + pass + + +async def record_lane_images( + thread_id: str, response_id: str, task_id: str, images: Dict[str, str] +) -> None: + """Record a lane's captured images on the runs the turn launched for it, + then re-project that turn, whose stored lines were keyed on the facts + before these. Best effort: images that fail to land only leave replay + showing the sandbox paths. + """ + try: + recorded = await facts_db.append_lane_images( + thread_id, response_id, task_id, images + ) + except Exception: + logger.warning( + f"[subagent_coordinator] lane images not recorded for task " + f"{task_id} of run {response_id}", + exc_info=True, + ) + return + if recorded: + _refresh_launching_turn(thread_id, response_id) + + class SubagentRunCoordinator: """Admission-authoritative task-run coordination for one thread.""" @@ -234,6 +273,9 @@ async def finalize_task_run( prune_if_archived_soon(self.thread_id) except Exception: pass + # The launching turn's replay could not be stored while this run + # still wrote; now it can. + _refresh_launching_turn(self.thread_id, run.get("parent_run_id")) if not defer_run_end: await self._append_v2_frame( task_run_id, @@ -254,6 +296,29 @@ async def finalize_task_run( pass return result + async def record_replay_facts( + self, task_run_id: str, *, steering_returned: list + ) -> bool: + """Keep the steering a settled run handed back for its replay, then + re-project the turn that launched it. Best effort: False when the + returns did not land, which only leaves replay without them. + """ + try: + row = await facts_db.append_returned_steering( + task_run_id, steering_returned + ) + except Exception: + logger.warning( + f"[subagent_coordinator] returned steering not recorded for " + f"run {task_run_id}", + exc_info=True, + ) + return False + if row is None: + return False + _refresh_launching_turn(row["thread_id"], row.get("parent_run_id")) + return True + async def append_run_end( self, task_run_id: str, *, task_id: str, outcome: str ) -> None: diff --git a/src/server/services/transcripts.py b/src/server/services/transcripts.py index def3c29b2..f97b8adab 100644 --- a/src/server/services/transcripts.py +++ b/src/server/services/transcripts.py @@ -20,14 +20,16 @@ import logging import uuid from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime from typing import TYPE_CHECKING, Any from langchain_core.messages import AnyMessage from ptc_agent.agent.transcript import ( + EarlierTurnsMissing, TranscriptTarget, + Window, build_directory, load_manifest, ) @@ -143,13 +145,17 @@ class _Job: checkpoint_id: str | None #: The manifest of the stored copy this one replaces, if there is one. previous: str | None + #: What the window trimmed from the head of ``messages``. + window: Window #: A task's render from the checkpoint: the run it read. run: TaskRun | None = None def render(self, *, inline: bool, full: bool = False) -> _Rendered: """The copy, each file named by digest. Only the segments that changed since ``previous`` render and carry bytes; the rest are carried over. - CPU only, so it runs in a worker thread.""" + CPU only, so it runs in a worker thread. Raises + ``EarlierTurnsMissing`` when ``previous`` cannot stand in for the + trimmed runs (see ``_whole``).""" from src.server.database.thread_transcripts import StoredTranscript directory = build_directory( @@ -157,6 +163,7 @@ def render(self, *, inline: bool, full: bool = False) -> _Rendered: unit=self.agent.unit, header=self.header, previous=None if full else load_manifest(self.previous), + window=self.window, ) copy = StoredTranscript( self.fingerprint, self.checkpoint_id, directory.manifest, run=self.run @@ -180,6 +187,30 @@ def render(self, *, inline: bool, full: bool = False) -> _Rendered: return _Rendered(copy, blobs) +async def _whole(job: _Job) -> _Job: + """``job`` over the thread's whole message list: the runs the window + trimmed are read back from the turn slices, for a render with no stored + copy to carry them from (a first export, a schema change, a fork that + dropped the copy, a copy moved under its render).""" + from src.server.services.history.window import earlier_runs + + earlier = await earlier_runs(job.agent.thread_id, job.checkpoint_id, job.window) + if earlier is None: + raise EarlierTurnsMissing( + f"{job.agent.directory}: {job.window.runs} trimmed runs unreadable" + ) + return replace(job, messages=[*earlier, *job.messages], window=Window()) + + +async def _render_job(job: _Job, *, inline: bool, full: bool = False) -> _Rendered: + try: + return await asyncio.to_thread(job.render, inline=inline, full=full) + except EarlierTurnsMissing: + logger.info(f"Transcript {job.agent.directory} renders its trimmed runs from slices") + whole = await _whole(job) + return await asyncio.to_thread(whole.render, inline=inline, full=full) + + async def _store_blobs(user_id: str, contents: dict[str, bytes]) -> None: from src.server.database.workspace_file_blobs import registered_blobs, store_blobs @@ -206,7 +237,7 @@ async def _save(target: _Target, job: _Job, rendered: _Rendered) -> bool: if job.previous is None: raise logger.info(f"Transcript {agent.directory} moved under its render; rendering it in full") - full = await asyncio.to_thread(job.render, inline=target.inline, full=True) + full = await _render_job(job, inline=target.inline, full=True) await _store_blobs(target.user_id, full.blobs) return await save_stored(agent.thread_id, agent.prefix, target.user_id, full.copy) @@ -256,7 +287,7 @@ async def behind_in_store(threads: list[tuple[str, str | None]]) -> list[Behind] async def _render_and_save(target: _Target, job: _Job) -> bool: - rendered = await asyncio.to_thread(job.render, inline=target.inline) + rendered = await _render_job(job, inline=target.inline) return await _save(target, job, rendered) @@ -270,7 +301,8 @@ async def _render_own( rendered_at = (state.config or {}).get("configurable", {}).get("checkpoint_id") if own is not None and _newer(own.checkpoint_id, rendered_at): return None - messages = list((state.values or {}).get("messages") or []) + values = state.values or {} + messages = list(values.get("messages") or []) if not messages: # Nothing to show; a copy already stored is left for the next # render that has something, rather than emptied. @@ -282,6 +314,7 @@ async def _render_own( _fingerprint(thread.checkpoint_id), rendered_at, own.manifest if own is not None else None, + Window.of(values), ) return await _render_and_save(target, job) @@ -310,6 +343,8 @@ async def _render_task( _fingerprint(_task_print(run)), history.checkpoint_id, previous.manifest if previous is not None else None, + # Only the main agent's checkpoint is trimmed (``with_window``). + Window(), run, ) return await _render_and_save(target, job) @@ -480,7 +515,9 @@ async def one(thread_id: str) -> None: return counts -async def save_live(transcript: TranscriptTarget, messages: Sequence[AnyMessage]) -> bool: +async def save_live( + transcript: TranscriptTarget, messages: Sequence[AnyMessage], *, window: Window +) -> bool: """Store one agent's transcript from messages in hand, ahead of the render from the checkpoint: compaction points the model at it mid-turn, before the turn end renders it. Only the segments that changed since the stored @@ -490,8 +527,9 @@ async def save_live(transcript: TranscriptTarget, messages: Sequence[AnyMessage] to that checkpoint, so a render from there (a turn end whose export has not landed, an export that read a running task before this compaction) leaves the copy in place; the next checkpoint the agent writes still - renders over it, and only what came after this save. Returns whether it - landed or had nothing to change.""" + renders over it, and only what came after this save. ``window`` was + trimmed from the head of ``messages``. Returns whether it landed or had + nothing to change.""" from src.server.database.conversation import ( get_thread_by_id, get_thread_checkpoint_id, @@ -532,8 +570,9 @@ async def save_live(transcript: TranscriptTarget, messages: Sequence[AnyMessage] LIVE_FINGERPRINT, checkpoint_id, stored.manifest if stored is not None else None, + window, ) - rendered = await asyncio.to_thread(job.render, inline=target.inline) + rendered = await _render_job(job, inline=target.inline) if stored is not None and rendered.copy.manifest == stored.manifest: return True return await _save(target, job, rendered) diff --git a/src/server/utils/checkpoint_helpers.py b/src/server/utils/checkpoint_helpers.py index 141f89bd4..d2e917238 100644 --- a/src/server/utils/checkpoint_helpers.py +++ b/src/server/utils/checkpoint_helpers.py @@ -4,9 +4,10 @@ patterns used across workflow endpoints. """ -import logging +import asyncio from collections import defaultdict from contextlib import nullcontext +from dataclasses import dataclass, field from functools import wraps from typing import Any, Callable, TypeVar @@ -16,8 +17,6 @@ from psycopg.rows import dict_row from psycopg_pool import AsyncConnectionPool -logger = logging.getLogger(__name__) - def _setup(): """The app's setup module, read at call time. @@ -110,6 +109,44 @@ def get_checkpointer(): return checkpointer +# LangGraph records a pending write on this channel when a node pauses via +# interrupt(). We match the stored channel name directly rather than importing +# the constant, which langgraph made private in v1.0 (deprecated import); the +# on-disk channel name is a stable storage-format detail. +INTERRUPT_CHANNEL = "__interrupt__" + + +@dataclass +class Boundary: + """A turn boundary on the current branch.""" + + checkpoint_id: str + parent_checkpoint_id: str | None + metadata: dict[str, Any] + # The checkpoint's interrupt records. On a resume boundary they are the + # answered interrupts of the turn it resumes: interrupt and resume writes + # attach to the same checkpoint. + interrupts: list[dict[str, Any]] = field(default_factory=list) + + @property + def is_input(self) -> bool: + """A new message; any other boundary resumes the turn before it.""" + return self.metadata.get("source") == "input" + + +def interrupt_records(pending_writes: Any) -> list[dict[str, Any]]: + """``{"id", "value"}`` records from a checkpoint's ``__interrupt__`` writes.""" + records: list[dict[str, Any]] = [] + for _task_id, channel, value in pending_writes or (): + if channel != INTERRUPT_CHANNEL: + continue + for intr in value if isinstance(value, (list, tuple)) else [value]: + records.append( + {"id": getattr(intr, "id", None), "value": getattr(intr, "value", None)} + ) + return records + + def is_turn_boundary(cp_tuple: Any) -> bool: """A checkpoint that starts a conversational turn. @@ -137,7 +174,7 @@ def _resumed_by_later_run(branch: list[Any], i: int) -> bool: """ cp = branch[i] if not any( - channel == "__interrupt__" for _, channel, _ in (cp.pending_writes or []) + channel == INTERRUPT_CHANNEL for _, channel, _ in (cp.pending_writes or []) ): return False run_id = (cp.metadata or {}).get("run_id") @@ -150,24 +187,123 @@ def _resumed_by_later_run(branch: list[Any], i: int) -> bool: return False -# Pending-write channels the walk consumers read: ``__resume__`` marks a HITL -# resume boundary (is_turn_boundary), ``__interrupt__`` carries the answered -# interrupt payloads (history reader). -_BOUNDARY_WRITE_CHANNELS = ("__resume__", "__interrupt__") +_BRANCH_TIP_SQL = """ + SELECT + (SELECT checkpoint_id FROM checkpoints + WHERE thread_id = %(thread_id)s AND checkpoint_ns = %(ns)s + AND checkpoint_id = %(requested)s) AS requested, + (SELECT checkpoint_id FROM checkpoints + WHERE thread_id = %(thread_id)s AND checkpoint_ns = %(ns)s + ORDER BY checkpoint_id DESC LIMIT 1) AS newest +""" -_warned_skeleton_fallback = False +# Parent links from the tip, filtered by is_turn_boundary and +# _resumed_by_later_run, oldest first. Each step is a LATERAL primary-key +# probe: as a plain join, a generic plan (the checkpointer prepares every +# statement) may hash the thread's whole history once per step, which goes +# quadratic on a long thread. The LIMIT keeps the planner from flattening the +# probe back into that join. For the later-run test, ``moves_after`` counts the +# non-update checkpoints after each one and ``move_no`` numbers them from the +# tip, so the first non-update checkpoint after a row is the one whose +# ``move_no`` equals that row's ``moves_after``. +_BRANCH_BOUNDARIES_SQL = """ + WITH RECURSIVE branch AS ( + SELECT checkpoint_id, parent_checkpoint_id, + metadata->>'source' AS source, metadata->>'run_id' AS run_id + FROM checkpoints + WHERE thread_id = %(thread_id)s AND checkpoint_ns = %(ns)s + AND checkpoint_id = %(tip)s + UNION ALL + SELECT p.checkpoint_id, p.parent_checkpoint_id, p.source, p.run_id + FROM branch b CROSS JOIN LATERAL ( + SELECT c.checkpoint_id, c.parent_checkpoint_id, + c.metadata->>'source' AS source, + c.metadata->>'run_id' AS run_id + FROM checkpoints c + WHERE c.thread_id = %(thread_id)s AND c.checkpoint_ns = %(ns)s + AND c.checkpoint_id = b.parent_checkpoint_id + LIMIT 1 + ) p + ), marks AS ( + SELECT checkpoint_id, + bool_or(channel = '__resume__') AS resumed, + bool_or(channel = '__interrupt__') AS interrupted + FROM checkpoint_writes + WHERE thread_id = %(thread_id)s AND checkpoint_ns = %(ns)s + AND channel IN ('__resume__', '__interrupt__') + GROUP BY checkpoint_id + ), steps AS ( + SELECT b.checkpoint_id, b.parent_checkpoint_id, b.source, b.run_id, + coalesce(m.resumed, false) AS resumed, + coalesce(m.interrupted, false) AS interrupted, + count(*) FILTER (WHERE b.source IS DISTINCT FROM 'update') OVER ( + ORDER BY b.checkpoint_id DESC + ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING + ) AS moves_after, + count(*) FILTER (WHERE b.source IS DISTINCT FROM 'update') OVER ( + ORDER BY b.checkpoint_id DESC ROWS UNBOUNDED PRECEDING + ) AS move_no + FROM branch b LEFT JOIN marks m USING (checkpoint_id) + ), bounds AS ( + SELECT s.checkpoint_id, s.parent_checkpoint_id + FROM steps s + LEFT JOIN steps nxt + ON nxt.move_no = s.moves_after + AND nxt.source IS DISTINCT FROM 'update' + WHERE s.source = 'input' + OR s.resumed + OR (s.interrupted + AND nxt.source = 'loop' + AND nxt.run_id IS NOT NULL + AND nxt.run_id IS DISTINCT FROM s.run_id) + ) + SELECT b.checkpoint_id, b.parent_checkpoint_id, m.metadata + FROM bounds b CROSS JOIN LATERAL ( + SELECT c.metadata + FROM checkpoints c + WHERE c.thread_id = %(thread_id)s AND c.checkpoint_ns = %(ns)s + AND c.checkpoint_id = b.checkpoint_id + LIMIT 1 + ) m + ORDER BY b.checkpoint_id +""" +# The boundary rule needs only which channels were written, which the walk +# reads in SQL; the one payload a consumer reads is the interrupts. +_BOUNDARY_INTERRUPTS_SQL = """ + SELECT checkpoint_id, task_id, channel, type, blob + FROM checkpoint_writes + WHERE thread_id = %(thread_id)s AND checkpoint_ns = %(ns)s + AND checkpoint_id = ANY(%(ids)s) AND channel = '__interrupt__' + ORDER BY checkpoint_id, task_id, idx +""" -async def _list_skeletons_via_tables( - checkpointer: AsyncPostgresSaver, thread_id: str, checkpoint_ns: str = "" -) -> list[CheckpointTuple]: - """Single-namespace checkpoint skeletons straight from the saver's tables. - ``alist`` eagerly joins every checkpoint's channel blobs and writes — - hundreds of ms on long threads — while the walk only needs ids, parents, - metadata, and the two boundary write channels. Couples to the - checkpoint-postgres *schema* (stable, versioned) instead of its API; any - failure here falls back to the ``alist`` path. +def _pick_tip( + thread_id: str, + requested: str | None, + requested_exists: bool, + newest: str | None, + strict: bool, +) -> str | None: + if requested is not None and not requested_exists and strict: + raise CheckpointBranchTipNotFound(thread_id, requested) + return requested if requested_exists else newest + + +async def _walk_via_tables( + checkpointer: AsyncPostgresSaver, + thread_id: str, + requested: str | None, + strict: bool, + checkpoint_ns: str, +) -> tuple[list[Boundary], str | None]: + """The branch walk run in Postgres. + + A long thread has tens of thousands of checkpoints and a few hundred + turns, and decoding every row's metadata to keep the boundaries held the + event loop for seconds. Couples to the checkpoint-postgres *schema* + (stable, versioned) instead of its API. """ conn_or_pool = checkpointer.conn if isinstance(conn_or_pool, AsyncConnectionPool): @@ -175,91 +311,96 @@ async def _list_skeletons_via_tables( else: conn_ctx = nullcontext(conn_or_pool) # caller-owned single connection + params = {"thread_id": thread_id, "ns": checkpoint_ns, "requested": requested} + write_rows: list[dict[str, Any]] = [] async with conn_ctx as conn, conn.cursor(row_factory=dict_row) as cur: - await cur.execute( - """SELECT checkpoint_id, parent_checkpoint_id, metadata - FROM checkpoints - WHERE thread_id = %s AND checkpoint_ns = %s - ORDER BY checkpoint_id DESC""", - (thread_id, checkpoint_ns), + await cur.execute(_BRANCH_TIP_SQL, params) + tip_row = await cur.fetchone() + tip_id = _pick_tip( + thread_id, + requested, + tip_row["requested"] is not None, + tip_row["newest"], + strict, ) + if tip_id is None: + return [], None + await cur.execute(_BRANCH_BOUNDARIES_SQL, {**params, "tip": tip_id}) rows = await cur.fetchall() - await cur.execute( - """SELECT checkpoint_id, task_id, channel, type, blob - FROM checkpoint_writes - WHERE thread_id = %s AND checkpoint_ns = %s - AND channel = ANY(%s) - ORDER BY checkpoint_id, task_id, idx""", - (thread_id, checkpoint_ns, list(_BOUNDARY_WRITE_CHANNELS)), - ) - write_rows = await cur.fetchall() + if rows: + await cur.execute( + _BOUNDARY_INTERRUPTS_SQL, + {**params, "ids": [r["checkpoint_id"] for r in rows]}, + ) + write_rows = await cur.fetchall() writes_by_cp: dict[str, list[tuple[str, str, Any]]] = defaultdict(list) for w in write_rows: value = checkpointer.serde.loads_typed((w["type"], bytes(w["blob"]))) writes_by_cp[w["checkpoint_id"]].append((w["task_id"], w["channel"], value)) - - def _config(checkpoint_id: str) -> dict[str, Any]: - return { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint_id, - } - } - - return [ - CheckpointTuple( - config=_config(r["checkpoint_id"]), - checkpoint={}, # skeleton: channel values are never read by the walk + boundaries = [ + Boundary( + checkpoint_id=r["checkpoint_id"], + parent_checkpoint_id=r["parent_checkpoint_id"], metadata=r["metadata"] or {}, - parent_config=( - _config(r["parent_checkpoint_id"]) - if r["parent_checkpoint_id"] - else None - ), - pending_writes=writes_by_cp.get(r["checkpoint_id"], []), + interrupts=interrupt_records(writes_by_cp.get(r["checkpoint_id"])), ) for r in rows ] + return boundaries, tip_id + + +def _boundary(cp: CheckpointTuple) -> Boundary: + parent = (cp.parent_config or {}).get("configurable") or {} + return Boundary( + checkpoint_id=cp.config["configurable"]["checkpoint_id"], + parent_checkpoint_id=parent.get("checkpoint_id"), + metadata=dict(cp.metadata or {}), + interrupts=interrupt_records(cp.pending_writes), + ) -async def _list_skeletons_via_alist( - checkpointer: Any, thread_id: str, checkpoint_ns: str = "" -) -> list[CheckpointTuple]: - """Public-API fallback: full tuples via ``alist``, one namespace only.""" +async def _walk_via_alist( + checkpointer: Any, + thread_id: str, + requested: str | None, + strict: bool, + checkpoint_ns: str, +) -> tuple[list[Boundary], str | None]: + """The walk for any other saver: every checkpoint via ``alist``, walked + in process.""" config = build_checkpoint_config(thread_id) config["configurable"]["checkpoint_ns"] = checkpoint_ns - return [cp_tuple async for cp_tuple in checkpointer.alist(config)] - + checkpoints = [cp_tuple async for cp_tuple in checkpointer.alist(config)] + cp_by_id = {cp.config["configurable"]["checkpoint_id"]: cp for cp in checkpoints} + tip_id = _pick_tip( + thread_id, + requested, + requested in cp_by_id, + checkpoints[0].config["configurable"]["checkpoint_id"] if checkpoints else None, + strict, + ) + if tip_id is None: + return [], None -async def _list_checkpoint_skeletons( - checkpointer: Any, thread_id: str, checkpoint_ns: str = "" -) -> list[CheckpointTuple]: - """Newest-first single-namespace checkpoints, as walk-sufficient skeletons. + current_branch: set[str] = set() + cursor: str | None = tip_id + while cursor and cursor in cp_by_id: + current_branch.add(cursor) + parent = cp_by_id[cursor].parent_config + cursor = parent["configurable"].get("checkpoint_id") if parent else None - Skeletons carry config/metadata/parent_config plus pending writes for the - boundary channels only — no channel values. Fast path reads the postgres - tables directly; anything else (in-memory savers, schema drift) uses the - public ``alist`` API, so the walk survives either the schema or the API - changing — just not both at once. - """ - global _warned_skeleton_fallback - if isinstance(checkpointer, AsyncPostgresSaver): - try: - return await _list_skeletons_via_tables( - checkpointer, thread_id, checkpoint_ns - ) - except Exception: - if not _warned_skeleton_fallback: - _warned_skeleton_fallback = True - logger.warning( - "[CHECKPOINT] Table-level checkpoint listing failed; " - "falling back to alist (slow). Check checkpoint-postgres " - "schema compatibility.", - exc_info=True, - ) - return await _list_skeletons_via_alist(checkpointer, thread_id, checkpoint_ns) + branch = [ + cp + for cp in reversed(checkpoints) # alist is newest-first + if cp.config["configurable"]["checkpoint_id"] in current_branch + ] + boundaries = [ + _boundary(cp) + for i, cp in enumerate(branch) + if is_turn_boundary(cp) or _resumed_by_later_run(branch, i) + ] + return boundaries, tip_id async def walk_current_branch_boundaries( @@ -269,59 +410,68 @@ async def walk_current_branch_boundaries( *, strict_branch_tip: bool = False, checkpoint_ns: str = "", -) -> tuple[list[Any], str | None]: - """Chronological turn-boundary checkpoints on the thread's current branch. +) -> tuple[list[Boundary], str | None]: + """Chronological turn boundaries on the thread's current branch. Edit/regenerate fork the checkpoint graph, so only ancestors of the branch tip count as turns: the tip is ``branch_tip_checkpoint_id`` when present and on the graph, else the newest checkpoint. With ``strict_branch_tip=True``, a supplied-but-missing tip raises ``CheckpointBranchTipNotFound`` instead of - silently reading the newest (possibly uncommitted) state. Canonical branch - walk shared by ``checkpoint_handler.get_thread_turns`` (turn CRUD) and the - history reader. + silently reading the newest (possibly uncommitted) state. Returns + ``(boundaries, tip_id)``, boundaries oldest-first; ``tip_id`` is None only + when the thread has none. - Returns ``(boundaries, tip_id)`` — boundaries oldest-first as skeleton - ``CheckpointTuple``s (no channel values; pending writes limited to the - boundary channels); ``tip_id`` is None only when the thread has none. + A Postgres saver walks in SQL only, and an error there raises: the + in-process walk decodes every checkpoint of the thread on the event loop, + the stall the SQL walk exists to remove, so it is no fallback. """ - checkpoints = await _list_checkpoint_skeletons( - checkpointer, thread_id, checkpoint_ns + walk = ( + _walk_via_tables + if isinstance(checkpointer, AsyncPostgresSaver) + else _walk_via_alist + ) + return await walk( + checkpointer, + thread_id, + branch_tip_checkpoint_id, + strict_branch_tip, + checkpoint_ns, ) - if not checkpoints: - if strict_branch_tip and branch_tip_checkpoint_id is not None: - raise CheckpointBranchTipNotFound(thread_id, branch_tip_checkpoint_id) - return [], None - cp_by_id = { - cp.config["configurable"]["checkpoint_id"]: cp for cp in checkpoints - } - if branch_tip_checkpoint_id is not None: - tip = cp_by_id.get(branch_tip_checkpoint_id) - if tip is None: - if strict_branch_tip: - raise CheckpointBranchTipNotFound( - thread_id, branch_tip_checkpoint_id - ) - tip = checkpoints[0] - else: - tip = checkpoints[0] # alist is newest-first - tip_id: str = tip.config["configurable"]["checkpoint_id"] - current_branch: set[str] = set() - cursor: str | None = tip_id - while cursor and cursor in cp_by_id: - current_branch.add(cursor) - parent = cp_by_id[cursor].parent_config - cursor = parent["configurable"].get("checkpoint_id") if parent else None +async def update_at_commit( + graph: Any, + thread_id: str, + built_on: str | None, + values: dict[str, Any], + *, + timeout: float | None = None, +) -> bool: + """Write ``values`` onto the thread's checkpoint ``built_on``, and move the + thread's commit pointer onto the write when it still points there. True + when the pointer moved; ``timeout`` bounds the write alone. + + ``built_on`` is both the write's anchor and the CAS guard. Unanchored, a + turn that starts on another worker between the caller's read and the + write would re-parent the write onto that turn's uncommitted checkpoint + while the CAS still passed, publishing partial state as the commit + pointer; so with no checkpoint to build on, nothing is written. + """ + if not built_on: + return False + from src.server.database.conversation import threads_write - branch = [ - cp - for cp in reversed(checkpoints) - if cp.config["configurable"]["checkpoint_id"] in current_branch - ] - boundaries = [ - cp - for i, cp in enumerate(branch) - if is_turn_boundary(cp) or _resumed_by_later_run(branch, i) - ] - return boundaries, tip_id + config = { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": "", + "checkpoint_id": built_on, + } + } + new_config = await asyncio.wait_for(graph.aupdate_state(config, values), timeout) + new_id = ((new_config or {}).get("configurable") or {}).get("checkpoint_id") + if not new_id: + return False + return await threads_write.advance_thread_checkpoint_id( + thread_id, from_checkpoint_id=built_on, to_checkpoint_id=new_id + ) diff --git a/src/server/utils/pg_sanitize.py b/src/server/utils/pg_sanitize.py index 381f31f27..408bb4f21 100644 --- a/src/server/utils/pg_sanitize.py +++ b/src/server/utils/pg_sanitize.py @@ -92,6 +92,12 @@ def _safe_dumps(value: Any) -> str: return _NUL_ESCAPE.sub(r"\1", s) +def jsonb_round_trip(value: Any) -> Any: + """``value`` as a ``SafeJson`` bind stores it, read back: what a reader of + the column will get.""" + return json.loads(_safe_dumps(value)) + + # Items per slice: a few milliseconds of encoding between event-loop yields. _JSON_ARRAY_SLICE = 2000 diff --git a/src/tools/secretary/utils.py b/src/tools/secretary/utils.py index 6da1018cb..0b8dde8ba 100644 --- a/src/tools/secretary/utils.py +++ b/src/tools/secretary/utils.py @@ -298,10 +298,11 @@ def _join_recent_turns(turn_texts: list[str]) -> str: async def extract_text_from_thread( thread_id: str, turns: int = 1, timezone: str = "UTC" ) -> dict[str, Any]: - """Extract text content from a thread's SSE events, one labeled block per run. + """Extract the text each run of a thread committed, one labeled block per run. Reads from Redis if the thread is actively running, otherwise reads - from the database. Filters for message_chunk events with text content. + from the database: each settled run's text comes from its turn slice, + falling back to the main lane of its stored events. Args: thread_id: The conversation thread ID @@ -360,7 +361,7 @@ async def extract_text_from_thread( text = _join_recent_turns([ _run_block( run, - _qualify_file_paths(_text_from_response(run), workspace_id), + _qualify_file_paths(run["main_text"], workspace_id), zone, ) for run in runs @@ -464,7 +465,7 @@ async def _extract_from_redis(thread_id: str, run_id: str) -> str: def _text_from_response(response: dict[str, Any]) -> str: - """Concatenate the text of one turn's ``message_chunk`` SSE events.""" + """The fallback text of one turn: its stored main-lane ``message_chunk`` events.""" chunks: list[str] = [] for event in response.get("sse_events") or []: if not isinstance(event, dict): @@ -474,6 +475,9 @@ def _text_from_response(response: dict[str, Any]) -> str: data = event.get("data", {}) if not isinstance(data, dict): continue + agent = data.get("agent") + if isinstance(agent, str) and agent.startswith("task:"): + continue # a subagent lane, not what the main agent said if data.get("content_type") == "text": content = data.get("content", "") if content: @@ -481,6 +485,47 @@ def _text_from_response(response: dict[str, Any]) -> str: return "".join(chunks) +async def _with_main_text( + thread_id: str, runs: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Copies of ``runs`` each carrying ``main_text``, what the turn committed + (``committed_texts``). A run that has no slice to read falls back to its + stored events' main lane, read for those runs alone.""" + if not runs: + return [] + # Local import: this module loads at agent import time. + from src.server.services.history.committed import committed_texts + + texts = await committed_texts( + thread_id, + {str(run["conversation_response_id"]): run.get("turn_index") for run in runs}, + ) + events = await _stored_events_of( + [ + response_id + for run in runs + if (response_id := str(run["conversation_response_id"])) not in texts + ] + ) + out: list[dict[str, Any]] = [] + for run in runs: + response_id = str(run["conversation_response_id"]) + text = texts.get(response_id) + if text is None: + text = _text_from_response({"sse_events": events.get(response_id)}) + out.append({**run, "main_text": text}) + return out + + +async def _stored_events_of(response_ids: list[str]) -> dict[str, Any]: + if not response_ids: + return {} + from src.server.database.conversation.replay_rows import get_replay_responses + + rows = await get_replay_responses(response_ids) + return {response_id: row.get("sse_events") for response_id, row in rows.items()} + + async def _runs_to_read(thread_id: str, turns: int = 1) -> list[dict[str, Any]]: """The settled runs to show, oldest -> newest, for the most-recent ``turns`` turns. @@ -496,21 +541,28 @@ async def _runs_to_read(thread_id: str, turns: int = 1) -> list[dict[str, Any]]: if turns != 1: limit = _MAX_HISTORY_TURNS if turns <= 0 else min(turns, _MAX_HISTORY_TURNS) - return await get_recent_responses_for_thread(thread_id, limit=limit) + return await _with_main_text( + thread_id, await get_recent_responses_for_thread(thread_id, limit=limit) + ) # Hot path is limit=1; a text-less newest run pays one wider read. - runs = await get_recent_responses_for_thread(thread_id, limit=1) - if not runs or _text_from_response(runs[0]): + runs = await _with_main_text( + thread_id, await get_recent_responses_for_thread(thread_id, limit=1) + ) + if not runs or runs[0]["main_text"]: return runs - window = await get_recent_responses_for_thread( - thread_id, limit=_EMPTY_LATEST_FALLBACK_TURNS + window = await _with_main_text( + thread_id, + await get_recent_responses_for_thread( + thread_id, limit=_EMPTY_LATEST_FALLBACK_TURNS + ), ) if not window: return runs newest = window[-1] - if _text_from_response(newest): # a newer run settled between the reads + if newest["main_text"]: # a newer run settled between the reads return [newest] for run in reversed(window[:-1]): - if _text_from_response(run): + if run["main_text"]: return [run, newest] return [newest] diff --git a/tests/integration/test_backfill_turn_slices_db.py b/tests/integration/test_backfill_turn_slices_db.py new file mode 100644 index 000000000..e4ccde6ed --- /dev/null +++ b/tests/integration/test_backfill_turn_slices_db.py @@ -0,0 +1,231 @@ +"""scripts/utils/backfill_turn_slices.py against real Postgres: the first pass, +and the re-key after a deploy changes the projection or the slice cut. + +A deploy leaves every stored row in place under a key no read accepts, so the +selection has to find rows that are present but stale, not only absent ones. +One thread has a turn whose run died before its first checkpoint: it never +gets a row, so that thread stays a candidate each pass confirms current +without projecting anything. +""" + +from __future__ import annotations + +import argparse +import uuid +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.graph import START, StateGraph + +import src.server.utils.checkpointer # noqa: F401 (installs the walk guard) +from ptc_agent.agent.state import DeltaAgentState +from scripts.utils import backfill_turn_slices as backfill +from src.server.services.history import slices +from src.server.services.history.replay import keys + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + +DEAD_TURN = 1 + + +def _graph(saver): + def reply(state): + n = len(state["messages"]) + return {"messages": [AIMessage(content=f"reply {n}", id=f"ai-{uuid.uuid4()}")]} + + return ( + StateGraph(DeltaAgentState) + .add_node("reply", reply) + .add_edge(START, "reply") + .compile(checkpointer=saver) + ) + + +async def _thread(saver, workspace_id: str, *, dead: int | None) -> str: + """Three turns run and finalized the way a turn is; ``dead``'s run fails + before the graph writes anything.""" + from src.server.database.conversation import create_thread + from src.server.database.runs.lifecycle import ( + QuerySpec, + RunOutcome, + finalize_run, + start_run, + ) + + graph = _graph(saver) + thread_id = str(uuid.uuid4()) + await create_thread( + conversation_thread_id=thread_id, + workspace_id=workspace_id, + current_status="completed", + msg_type="ptc", + ) + for text in ("q0", "q1", "q2"): + run_id = str(uuid.uuid4()) + row = await start_run( + run_id=run_id, + thread_id=thread_id, + request_key=str(uuid.uuid4()), + query=QuerySpec(query_id=str(uuid.uuid4()), content=text, query_type="initial"), + metadata={"msg_type": "ptc"}, + ) + turn_index = row["turn_index"] + if turn_index == dead: + await finalize_run( + run_id=run_id, thread_id=thread_id, outcome=RunOutcome(status="error") + ) + continue + config = { + "configurable": {"thread_id": thread_id}, + "metadata": {"run_id": run_id, "turn_index": turn_index}, + } + await graph.ainvoke({"messages": [HumanMessage(content=text)]}, config) + state = await graph.aget_state({"configurable": {"thread_id": thread_id}}) + await finalize_run( + run_id=run_id, + thread_id=thread_id, + outcome=RunOutcome(status="completed"), + checkpoint_id=state.config["configurable"]["checkpoint_id"], + ) + return thread_id + + +@pytest_asyncio.fixture(loop_scope="session") +async def threads(seed_workspace, test_db_pool, patched_get_db_connection, monkeypatch): + from src.server.app import setup + from src.server.services.history.reader import CheckpointHistoryReader + + saver = AsyncPostgresSaver(test_db_pool) + monkeypatch.setattr(setup, "checkpointer", saver) + CheckpointHistoryReader.reset_instance() + workspace_id = str(seed_workspace["workspace_id"]) + out = SimpleNamespace( + whole=await _thread(saver, workspace_id, dead=None), + dead=await _thread(saver, workspace_id, dead=DEAD_TURN), + ) + yield out + CheckpointHistoryReader.reset_instance() + for thread_id in (out.whole, out.dead): + await saver.adelete_thread(thread_id) + + +def _args(*threads: str, apply: bool = True, recheck: bool = False) -> argparse.Namespace: + return argparse.Namespace( + apply=apply, + days=None, + user=None, + workspace=None, + thread=list(threads), + recheck=recheck, + concurrency=2, + pause=0, + ) + + +async def _candidates(*threads: str) -> dict[str, tuple[int, int]]: + return {c.thread_id: (c.missing, c.stale) for c in await backfill._threads(_args(*threads))} + + +async def _rows(*threads: str) -> list[tuple]: + from psycopg.rows import tuple_row + + from src.server.database.pool import get_db_connection + + async with get_db_connection() as conn, conn.cursor(row_factory=tuple_row) as cur: + await cur.execute( + "SELECT s.conversation_thread_id::text, r.turn_index, s.slice_key," + " s.lines_key, s.lines IS NOT NULL, s.built_at FROM turn_slices s" + " JOIN conversation_responses r USING (conversation_response_id)" + " WHERE s.conversation_thread_id = ANY(%s::uuid[]) ORDER BY 1, 2", + (list(threads),), + ) + return await cur.fetchall() + + +def _all_current(rows: list[tuple]) -> bool: + epoch = f"{keys.lines_epoch()}." + return all( + slice_key == slices.SLICE_KEY and lines_key.startswith(epoch) and has_lines + for _, _, slice_key, lines_key, has_lines, _ in rows + ) + + +async def test_a_first_pass_stores_every_turn_and_a_rerun_projects_nothing(threads, capsys): + both = (threads.whole, threads.dead) + assert await _candidates(*both) == {threads.whole: (3, 0), threads.dead: (3, 0)} + + assert await backfill._run(_args(*both, apply=False)) == 0 + assert await _rows(*both) == [] + + assert await backfill._run(_args(*both)) == 0 + rows = await _rows(*both) + assert [(t, i) for t, i, *_ in rows] == sorted( + [(threads.whole, i) for i in (0, 1, 2)] + [(threads.dead, i) for i in (0, 2)] + ) + assert _all_current(rows) + assert "2 projected (5 turns)" in capsys.readouterr().out + + # The dead turn never gets a row, so its thread stays a candidate; the + # read path's own check finds nothing to project there. + assert await _candidates(*both) == {threads.dead: (1, 0)} + assert await backfill._run(_args(*both)) == 0 + assert "0 projected (0 turns), 1 already current" in capsys.readouterr().out + assert await _rows(*both) == rows + + +@pytest.mark.parametrize("change", ["projection", "slice cut"]) +async def test_a_deploy_that_changes_the_projection_rekeys_every_stored_turn( + threads, monkeypatch, capsys, change +): + both = (threads.whole, threads.dead) + await backfill._run(_args(*both)) + before = await _rows(*both) + + if change == "slice cut": + monkeypatch.setattr(slices, "SLICE_KEY", f"{slices.SLICE_KEY}-next") + monkeypatch.setattr( + keys, "_PROJECTION_VERSION", f"{keys.LINES_VERSION}.{slices.SLICE_KEY}.next" + ) + assert not _all_current(before) + assert await _candidates(*both) == {threads.whole: (0, 3), threads.dead: (1, 2)} + + capsys.readouterr() + assert await backfill._run(_args(*both)) == 0 + after = await _rows(*both) + assert [(t, i) for t, i, *_ in after] == [(t, i) for t, i, *_ in before] + assert _all_current(after) + assert "2 projected (5 turns)" in capsys.readouterr().out + + assert await _candidates(*both) == {threads.dead: (1, 0)} + assert await backfill._run(_args(*both)) == 0 + assert await _rows(*both) == after + + +async def test_a_bulk_row_update_is_found_by_a_recheck(threads, capsys): + """An UPDATE leaves every key's epoch current but moves the row versions + the keys cover, so only the read path's own check sees it.""" + from src.server.database.pool import get_db_connection + + both = (threads.whole, threads.dead) + await backfill._run(_args(*both)) + before = {(t, i): key for t, i, _, key, *_ in await _rows(*both)} + async with get_db_connection() as conn: + await conn.execute( + "UPDATE conversation_responses SET sse_events = NULL" + " WHERE conversation_thread_id = %s::uuid", + (threads.whole,), + ) + assert await _candidates(*both) == {threads.dead: (1, 0)} + + capsys.readouterr() + assert await backfill._run(_args(*both, recheck=True)) == 0 + assert "1 projected (3 turns), 1 already current" in capsys.readouterr().out + after = {(t, i): key for t, i, _, key, *_ in await _rows(*both)} + assert [k for k in before if before[k] != after[k]] == [(threads.whole, i) for i in (0, 1, 2)] + assert _all_current(await _rows(*both)) + + assert await backfill._run(_args(*both, recheck=True)) == 0 + assert "0 projected (0 turns), 2 already current" in capsys.readouterr().out diff --git a/tests/integration/test_checkpoint_walk_db.py b/tests/integration/test_checkpoint_walk_db.py new file mode 100644 index 000000000..96554f692 --- /dev/null +++ b/tests/integration/test_checkpoint_walk_db.py @@ -0,0 +1,198 @@ +"""The branch walk on real Postgres agrees with the in-process walk. + +The table path restates ``is_turn_boundary`` and ``_resumed_by_later_run`` in +SQL, and the ``alist`` walk other savers take runs them in Python. A boundary +rule added to only one side shows up here as the two paths disagreeing. +""" + +from __future__ import annotations + +import uuid + +import psycopg +import pytest +import pytest_asyncio +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.graph import START, StateGraph +from langgraph.types import Command, interrupt + +from ptc_agent.agent.state import DeltaAgentState +from src.server.utils import checkpoint_helpers +from src.server.utils.checkpoint_helpers import ( + CheckpointBranchTipNotFound, + _walk_via_alist, + _walk_via_tables, + walk_current_branch_boundaries, +) + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +@pytest_asyncio.fixture(loop_scope="session") +async def saver(test_db_pool): + return AsyncPostgresSaver(test_db_pool) # binds the running loop + + +def _graph(saver, *, asks=True): + """``ask`` waits on the user and ``reply`` follows. ``asks=False`` is + another agent holding the same node, one that runs without asking.""" + + def ask(state): + if not asks: + return {"messages": [AIMessage(content="refused", id="ai-refused")]} + answer = interrupt({"action_requests": [{"description": "go?"}]}) + return {"messages": [AIMessage(content=f"asked: {answer}", id="ai-asked")]} + + def reply(state): + return {"messages": [AIMessage(content="reply", id=f"ai-{len(state['messages'])}")]} + + return ( + StateGraph(DeltaAgentState) + .add_node("ask", ask) + .add_node("reply", reply) + .add_edge(START, "ask") + .add_edge("ask", "reply") + .compile(checkpointer=saver) + ) + + +def _cfg(thread_id, *, run_id=None, checkpoint_id=None): + cfg = {"configurable": {"thread_id": thread_id}} + if checkpoint_id: + cfg["configurable"]["checkpoint_id"] = checkpoint_id + if run_id: + cfg["metadata"] = {"run_id": run_id} + return cfg + + +async def _interrupted_turn(graph, thread_id): + await graph.ainvoke( + {"messages": [HumanMessage(content="q0", id="h-0")]}, + _cfg(thread_id, run_id="run-0"), + ) + state = await graph.aget_state(_cfg(thread_id)) + return state.config["configurable"]["checkpoint_id"], state.interrupts[0].id + + +async def _both_walks(saver, thread_id, tip=None, strict=False): + """Walk both paths, assert they agree, and return the boundary ids.""" + tables = await _walk_via_tables(saver, thread_id, tip, strict, "") + assert tables == await _walk_via_alist(saver, thread_id, tip, strict, "") + return [b.checkpoint_id for b in tables[0]] + + +async def test_resume_of_cancelled_calls_is_its_own_boundary(saver): + thread = str(uuid.uuid4()) + graph = _graph(saver) + interrupt_tip, interrupt_id = await _interrupted_turn(graph, thread) + await graph.aupdate_state( + _cfg(thread), + {"messages": [AIMessage(content="Not run", id="ai-cancelled")]}, + as_node="ask", + ) + await graph.ainvoke(Command(resume={interrupt_id: "yes"}), _cfg(thread, run_id="run-1")) + + boundaries = await _both_walks(saver, thread) + assert len(boundaries) == 2 + assert boundaries[1] == interrupt_tip + + +async def test_resume_by_an_agent_that_does_not_ask_is_its_own_boundary(saver): + thread = str(uuid.uuid4()) + interrupt_tip, interrupt_id = await _interrupted_turn(_graph(saver), thread) + await _graph(saver, asks=False).ainvoke( + Command(resume={interrupt_id: "yes"}), _cfg(thread, run_id="run-1") + ) + + assert (await _both_walks(saver, thread))[1:] == [interrupt_tip] + + +async def test_interrupt_continued_by_its_own_run_is_no_boundary(saver): + thread = str(uuid.uuid4()) + graph = _graph(saver) + await _interrupted_turn(graph, thread) + await graph.aupdate_state( + _cfg(thread, run_id="run-0"), + {"messages": [AIMessage(content="note", id="ai-note")]}, + as_node="ask", + ) + await graph.ainvoke(None, _cfg(thread, run_id="run-0")) + + assert len(await _both_walks(saver, thread)) == 1 + + +async def test_new_message_over_an_unanswered_interrupt_is_one_boundary(saver): + thread = str(uuid.uuid4()) + graph = _graph(saver) + interrupt_tip, _ = await _interrupted_turn(graph, thread) + await graph.ainvoke( + {"messages": [HumanMessage(content="q1", id="h-1")]}, + _cfg(thread, run_id="run-1"), + ) + + boundaries = await _both_walks(saver, thread) + assert len(boundaries) == 2 + assert interrupt_tip not in boundaries + + +async def test_answered_interrupt_and_an_interrupt_at_the_tip(saver): + thread = str(uuid.uuid4()) + graph = _graph(saver) + _, interrupt_id = await _interrupted_turn(graph, thread) + assert len(await _both_walks(saver, thread)) == 1 # nothing after it yet + + await graph.ainvoke(Command(resume={interrupt_id: "yes"}), _cfg(thread, run_id="run-1")) + assert len(await _both_walks(saver, thread)) == 2 # the resume writes __resume__ + boundaries, _ = await _walk_via_tables(saver, thread, None, False, "") + assert not boundaries[1].is_input + assert boundaries[1].interrupts == [ + {"id": interrupt_id, "value": {"action_requests": [{"description": "go?"}]}} + ] + + +async def test_a_fork_walks_only_its_own_branch(saver): + thread = str(uuid.uuid4()) + graph = _graph(saver, asks=False) + await graph.ainvoke({"messages": [HumanMessage(content="q0", id="h-0")]}, _cfg(thread)) + first_turn_tip = (await graph.aget_state(_cfg(thread))).config["configurable"]["checkpoint_id"] + await graph.ainvoke({"messages": [HumanMessage(content="q1", id="h-1")]}, _cfg(thread)) + old_tip = (await graph.aget_state(_cfg(thread))).config["configurable"]["checkpoint_id"] + await graph.ainvoke( + {"messages": [HumanMessage(content="q1 edited", id="h-1b")]}, + _cfg(thread, checkpoint_id=first_turn_tip), + ) + + newest = await _both_walks(saver, thread) + old = await _both_walks(saver, thread, tip=old_tip, strict=True) + assert len(newest) == len(old) == 2 + assert newest[0] == old[0] and newest[1] != old[1] + + with pytest.raises(CheckpointBranchTipNotFound): + await _walk_via_tables(saver, thread, "missing", True, "") + with pytest.raises(CheckpointBranchTipNotFound): + await _walk_via_alist(saver, thread, "missing", True, "") + + +async def test_a_failed_sql_walk_raises_rather_than_walking_in_process( + saver, monkeypatch +): + """Schema drift (a renamed column, say) fails the read where it happened + instead of decoding the whole thread on the event loop.""" + thread = str(uuid.uuid4()) + graph = _graph(saver, asks=False) + await graph.ainvoke({"messages": [HumanMessage(content="q0", id="h-0")]}, _cfg(thread)) + monkeypatch.setattr( + checkpoint_helpers, + "_BRANCH_BOUNDARIES_SQL", + checkpoint_helpers._BRANCH_BOUNDARIES_SQL.replace( + "parent_checkpoint_id", "parent_id" + ), + ) + + async def no_alist(*args, **kwargs): + raise AssertionError("walked in process") + + monkeypatch.setattr(checkpoint_helpers, "_walk_via_alist", no_alist) + with pytest.raises(psycopg.errors.UndefinedColumn): + await walk_current_branch_boundaries(saver, thread) diff --git a/tests/integration/test_finalize_run_db.py b/tests/integration/test_finalize_run_db.py new file mode 100644 index 000000000..c6a6598af --- /dev/null +++ b/tests/integration/test_finalize_run_db.py @@ -0,0 +1,295 @@ +"""finalize_run against real Postgres. + +The finalize merges its events onto whatever mid-run appenders wrote, writes +the run's replay facts with its status, and rewrites the main lane's +provenance from the lane's own record, so only a real database can show that +an append queued on the row lock reaches the archive, that a subagent lane's +records survive, and that a lost CAS writes nothing. +""" + +from __future__ import annotations + +import asyncio +import uuid +from unittest.mock import patch + +import pytest +import pytest_asyncio + +# Session loop: the pool lives there, and the lock test needs a second connection. +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + +# What a finalize hands back, won or lost. +ROW_KEYS = { + "conversation_response_id", + "conversation_thread_id", + "turn_index", + "run_seq", + "status", + "interrupt_reason", + "metadata", +} + +CHUNK = {"event": "message_chunk", "data": {"agent": "ptc", "content": "m"}} +COMPACT = {"event": "context_window", "data": {"agent": "ptc", "action": "compact"}} + + +def _provenance(url: str) -> dict: + return { + "event": "provenance", + "source_type": "web", + "identifier": url, + "agent": "ptc", + } + + +def _finalize(run_id: str, thread_id: str, status: str, **fields): + from src.server.database.runs.lifecycle import RunOutcome, finalize_run + + return finalize_run( + run_id=run_id, thread_id=thread_id, outcome=RunOutcome(status=status, **fields) + ) + + +def _record(*events: dict, reasoning_ms=None): + """A producer's record of ``events``' provenance.""" + from src.server.database.runs.lifecycle import ProducerRecord + + return ProducerRecord( + provenance=[e for e in events if e["event"] == "provenance"], + reasoning_ms=reasoning_ms, + ) + + +@pytest_asyncio.fixture(loop_scope="session") +async def open_run(seed_workspace, patched_get_db_connection): + from src.server.database.conversation import create_thread + from src.server.database.runs.lifecycle import start_run + + thread_id = str(uuid.uuid4()) + await create_thread( + conversation_thread_id=thread_id, + workspace_id=str(seed_workspace["workspace_id"]), + current_status="in_progress", + msg_type="ptc", + ) + run_id = str(uuid.uuid4()) + await start_run( + run_id=run_id, + thread_id=thread_id, + request_key=str(uuid.uuid4()), + metadata={"msg_type": "ptc", "user_id": seed_workspace["user_id"]}, + ) + return thread_id, run_id + + +async def _sse_events(conn, run_id: str): + cur = await conn.execute( + "SELECT sse_events FROM conversation_responses WHERE conversation_response_id = %s", + (run_id,), + ) + return (await cur.fetchone())["sse_events"] + + +async def _provenance_ids(conn, run_id: str) -> list[str]: + cur = await conn.execute( + "SELECT identifier FROM provenance_records " + "WHERE conversation_response_id = %s ORDER BY identifier", + (run_id,), + ) + return [r["identifier"] for r in await cur.fetchall()] + + +async def test_events_merge_onto_mid_run_appends_and_encode_once(open_run, db_conn): + from src.server.database.conversation import append_sse_event + from src.server.utils import pg_sanitize + + thread_id, run_id = open_run + early = _provenance("https://early.example") + assert await append_sse_event(thread_id, early) + events = [CHUNK, _provenance("https://final.example")] + + encodes = [] + real_dumps = pg_sanitize._safe_dumps + + def counting_dumps(value): + if value is events: + encodes.append(value) + return real_dumps(value) + + with patch.object(pg_sanitize, "_safe_dumps", counting_dumps): + result = await _finalize( + run_id, thread_id, "completed", sse_events=events, record=_record(*events) + ) + + assert result.applied + assert set(result.run) == ROW_KEYS + assert result.run["status"] == "completed" + assert len(encodes) == 1 + assert await _sse_events(db_conn, run_id) == [early, *events] + # The lane's own record, not the merged archive: no mid-run writer + # appends provenance, so the archive holds nothing the record lacks. + assert await _provenance_ids(db_conn, run_id) == ["https://final.example"] + # The outbox jobs are built from the returned row alone. + cur = await db_conn.execute( + "SELECT payload FROM hook_outbox WHERE run_id = %s AND hook_type = 'user_feed'", + (run_id,), + ) + assert (await cur.fetchone())["payload"]["run_seq"] == result.run["run_seq"] + + +async def test_append_queued_on_the_row_lock_reaches_the_archive( + open_run, db_conn, test_db_pool +): + from src.server.database.conversation import append_sse_event + + thread_id, run_id = open_run + late = _provenance("https://late.example") + + async with test_db_pool.connection() as holder: + async with holder.transaction(): + assert await append_sse_event(thread_id, late, conn=holder) + finalize = asyncio.create_task( + _finalize( + run_id, + thread_id, + "completed", + sse_events=[CHUNK], + record=_record(CHUNK), + ) + ) + await asyncio.sleep(0.2) + assert not finalize.done() # queued behind the appender's row lock + assert (await finalize).applied + + assert await _sse_events(db_conn, run_id) == [late, CHUNK] + # Provenance is the lane's own record, which the finalize carried. + assert await _provenance_ids(db_conn, run_id) == [] + + +async def test_lost_cas_writes_nothing_and_returns_the_survivor(open_run, db_conn): + thread_id, run_id = open_run + won_events = [CHUNK, _provenance("https://won.example")] + first = await _finalize( + run_id, thread_id, "completed", sse_events=won_events, record=_record(*won_events) + ) + assert first.applied + + lost_events = [COMPACT, _provenance("https://lost.example")] + second = await _finalize( + run_id, + thread_id, + "error", + errors=["late"], + sse_events=lost_events, + record=_record(*lost_events), + ) + + assert not second.applied + assert set(second.run) == ROW_KEYS + assert second.run["status"] == "completed" + assert await _sse_events(db_conn, run_id) == won_events + assert await _provenance_ids(db_conn, run_id) == ["https://won.example"] + + +@pytest.mark.parametrize("appended", [None, COMPACT], ids=["empty", "mid-run-append"]) +async def test_finalize_without_events_keeps_the_column(open_run, db_conn, appended): + from src.server.database.conversation import append_sse_event + + thread_id, run_id = open_run + if appended is not None: + assert await append_sse_event(thread_id, appended) + + result = await _finalize(run_id, thread_id, "cancelled") + + assert result.applied + assert await _sse_events(db_conn, run_id) == ( + [appended] if appended is not None else None + ) + assert await _provenance_ids(db_conn, run_id) == [] + + +async def _replay_facts(conn, run_id: str): + cur = await conn.execute( + "SELECT replay_facts FROM conversation_responses " + "WHERE conversation_response_id = %s", + (run_id,), + ) + return (await cur.fetchone())["replay_facts"] + + +async def test_replay_facts_land_with_the_status_and_survive_a_lost_cas( + open_run, db_conn +): + thread_id, run_id = open_run + durations = {"msg-1": [1840, 920]} + first = await _finalize( + run_id, + thread_id, + "completed", + sse_events=[CHUNK], + record=_record(reasoning_ms=durations), + ) + assert first.applied + assert await _replay_facts(db_conn, run_id) == {"reasoning_ms": durations} + + second = await _finalize( + run_id, + thread_id, + "error", + errors=["late"], + record=_record(reasoning_ms={"msg-2": [5]}), + ) + assert not second.applied + assert await _replay_facts(db_conn, run_id) == {"reasoning_ms": durations} + + +@pytest.mark.parametrize("recovered", [False, True], ids=["fail-open", "recovery"]) +async def test_finalize_without_facts_leaves_the_column_null( + open_run, db_conn, recovered +): + """Recovery and fail-open settles hold no producer: their turns replay + reasoning durations from the stored events, as a legacy turn does.""" + thread_id, run_id = open_run + record = _record(CHUNK) if recovered else None + assert (await _finalize(run_id, thread_id, "cancelled", record=record)).applied + assert await _replay_facts(db_conn, run_id) is None + + +async def test_main_lane_provenance_leaves_subagent_lanes(open_run, db_conn): + """A subagent's archive writes its lane before the parent settles; the + finalize rewrites only the main lane.""" + from src.server.database.conversation.responses import replace_agent_events + + thread_id, run_id = open_run + task_record = { + "event": "provenance", + "data": { + "agent": "task:abc", + "source_type": "web", + "identifier": "https://task.example", + }, + } + assert await replace_agent_events(run_id, [task_record]) + assert await _provenance_ids(db_conn, run_id) == ["https://task.example"] + + main_record = { + "event": "provenance", + "data": { + "agent": "main", + "source_type": "web", + "identifier": "https://main.example", + }, + } + result = await _finalize( + run_id, + thread_id, + "completed", + sse_events=[CHUNK, main_record], + record=_record(main_record), + ) + assert result.applied + assert await _provenance_ids(db_conn, run_id) == [ + "https://main.example", + "https://task.example", + ] diff --git a/tests/integration/test_fork_turns_db.py b/tests/integration/test_fork_turns_db.py index f287914df..0ba616a4c 100644 --- a/tests/integration/test_fork_turns_db.py +++ b/tests/integration/test_fork_turns_db.py @@ -5,7 +5,7 @@ every turn after it sits one place lower among the boundaries than its number. A fork deletes rows from the turn it replaces onward, and the turn it replaces is fixed by the checkpoint it forks at: deleting by -position would drop the rows of a turn the branch keeps. +position would drop the rows (and the slice) of a turn the branch keeps. """ from __future__ import annotations @@ -64,7 +64,7 @@ async def _start(thread_id, text, *, fork=None, turn_index=None, query=True): async def _run(graph, thread_id, run_id, turn_index, text, *, checkpoint_id=None): """The graph half of a turn: stamped the way ``build_graph_config`` stamps a run, then committed the way the finalize commits it.""" - from src.server.database.runs.lifecycle import finalize_run + from src.server.database.runs.lifecycle import RunOutcome, finalize_run configurable = {"thread_id": thread_id} if checkpoint_id: @@ -79,7 +79,7 @@ async def _run(graph, thread_id, run_id, turn_index, text, *, checkpoint_id=None await finalize_run( run_id=run_id, thread_id=thread_id, - status="completed", + outcome=RunOutcome(status="completed"), checkpoint_id=state.config["configurable"]["checkpoint_id"], ) @@ -87,9 +87,10 @@ async def _run(graph, thread_id, run_id, turn_index, text, *, checkpoint_id=None @pytest_asyncio.fixture(loop_scope="session") async def thread(seed_workspace, test_db_pool, patched_get_db_connection): """Turns 0, 2 and 3 ran; turn 1's run failed before the graph wrote - anything.""" + anything. Turn 2's slice is stored, as the turn finalize stores it.""" from src.server.database.conversation import create_thread - from src.server.database.runs.lifecycle import finalize_run + from src.server.database.conversation import turn_slices as slices_db + from src.server.database.runs.lifecycle import RunOutcome, finalize_run saver = AsyncPostgresSaver(test_db_pool) graph = _graph(saver) @@ -105,9 +106,20 @@ async def thread(seed_workspace, test_db_pool, patched_get_db_connection): run_id, turn_index = await _start(thread_id, text) runs[turn_index] = run_id if turn_index == DEAD_TURN: - await finalize_run(run_id=run_id, thread_id=thread_id, status="error") + await finalize_run( + run_id=run_id, thread_id=thread_id, outcome=RunOutcome(status="error") + ) else: await _run(graph, thread_id, run_id, turn_index, text) + await slices_db.upsert_turn_slices( + [ + slices_db.TurnRow( + runs[2], + thread_id, + slices_db.StoredSlice("in-2", "tail-2", "k", "msgpack", b"\x80"), + ) + ] + ) with patch( "src.server.handlers.checkpoint_handler.get_checkpointer", return_value=saver, @@ -147,6 +159,16 @@ async def _rows(conn, table, thread_id): return [r["turn_index"] for r in await cur.fetchall()] +async def _slices(conn, thread_id): + cur = await conn.execute( + "SELECT r.turn_index FROM turn_slices s " + "JOIN conversation_responses r USING (conversation_response_id) " + "WHERE s.conversation_thread_id = %s", + (thread_id,), + ) + return [r["turn_index"] for r in await cur.fetchall()] + + async def test_turns_carry_the_turn_number_not_the_position(thread): turns = await _turns(thread.id) @@ -172,6 +194,7 @@ async def test_editing_the_latest_turn_keeps_every_turn_before_it(thread, db_con assert turn_index == 3 assert await _rows(db_conn, "conversation_responses", thread.id) == [0, 1, 2, 3] assert await _rows(db_conn, "conversation_queries", thread.id) == [0, 1, 2, 3] + assert await _slices(db_conn, thread.id) == [2] # Every boundary on the new branch still names a turn with rows. assert [t.turn_index for t in await _turns(thread.id)] == [0, 2, 3] @@ -189,6 +212,7 @@ async def test_regenerating_the_latest_turn_keeps_its_question(thread, db_conn): assert turn_index == 3 assert await _rows(db_conn, "conversation_responses", thread.id) == [0, 1, 2, 3] assert await _rows(db_conn, "conversation_queries", thread.id) == [0, 1, 2, 3] + assert await _slices(db_conn, thread.id) == [2] async def test_a_client_count_cannot_move_the_fork(thread, db_conn): @@ -204,6 +228,7 @@ async def test_a_client_count_cannot_move_the_fork(thread, db_conn): assert (edit.from_turn, regenerate.from_turn) == (3, 3) await _start(thread.id, "q3 edited", fork=edit) assert await _rows(db_conn, "conversation_responses", thread.id) == [0, 1, 2, 3] + assert await _slices(db_conn, thread.id) == [2] async def test_rewriting_the_failed_message_replaces_it(thread, db_conn): @@ -224,6 +249,7 @@ async def test_rewriting_the_failed_message_replaces_it(thread, db_conn): assert turn_index == DEAD_TURN assert await _rows(db_conn, "conversation_responses", thread.id) == [0, 1] + assert await _slices(db_conn, thread.id) == [] assert [t.turn_index for t in await _turns(thread.id)] == [0, 1] diff --git a/tests/integration/test_legacy_facts_db.py b/tests/integration/test_legacy_facts_db.py new file mode 100644 index 000000000..4aecbc7ee --- /dev/null +++ b/tests/integration/test_legacy_facts_db.py @@ -0,0 +1,301 @@ +"""The legacy facts backfill against real Postgres. + +A backfill write races the run's own writers, so only a real database can show +that it keeps the keys they set, lands only on the row version it derived +from, is not repeated, and that a later write to the stored events takes the +facts away again; and that a backfilled row's events are left out of the +replay read. The checkpoint side is the unit suites' mocked reader: the rows +are what is under test here. +""" + +from __future__ import annotations + +import uuid +from collections import Counter +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from langchain_core.messages import AIMessage + +from scripts.utils import backfill_replay_facts as backfill +from src.server.services.history.replay import legacy +from tests.unit.server.services.history.replay_builders import ( + ThreadHistory, + _cache_probe, + _mock_reader, + _turn, +) + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + +TIP = "cp-tip" +RUN_FACTS = {"v": 1, "reasoning_ms": {"msg_1": [250]}} +STORED = [ + {"event": "message_chunk", "data": {"id": "lc-1", "content_type": "text", "content": "answer"}}, + {"event": "context_window", "data": {"action": "offload", "kind": "tool_args"}}, +] +LATER = {"event": "provenance", "data": {"source_type": "web", "identifier": "https://a.example"}} + + +@pytest_asyncio.fixture +async def turn(seed_workspace, seed_response, patched_get_db_connection, monkeypatch): + """A one-turn thread whose completed row has stored events and the run + facts its finalize wrote, over a checkpoint that projects the turn.""" + from src.server.database.conversation import create_query, create_thread + + thread_id = str(uuid.uuid4()) + await create_thread( + conversation_thread_id=thread_id, + workspace_id=str(seed_workspace["workspace_id"]), + current_status="completed", + msg_type="ptc", + ) + await create_query( + conversation_query_id=str(uuid.uuid4()), + conversation_thread_id=thread_id, + turn_index=0, + content="question", + query_type="initial", + ) + response_id = str(uuid.uuid4()) + await seed_response(response_id, thread_id, 0, sse_events=STORED) + + async def sql(query, params=()): + async with patched_get_db_connection() as conn: + cur = await conn.execute(query, params) + return await cur.fetchone() if cur.description else None + + await sql( + "UPDATE conversation_threads SET latest_checkpoint_id = %s " + "WHERE conversation_thread_id = %s", + (TIP, thread_id), + ) + await sql( + "UPDATE conversation_responses SET replay_facts = %s::jsonb " + "WHERE conversation_response_id = %s", + (_json(RUN_FACTS), response_id), + ) + _cache_probe(monkeypatch) + _mock_reader( + monkeypatch, + ThreadHistory(thread_id=thread_id, turns=[_turn(0, [AIMessage("answer", id="ai-1")])]), + ) + + async def row(): + return await sql( + "SELECT replay_facts, sse_events, metadata, xmin::text AS xmin " + "FROM conversation_responses WHERE conversation_response_id = %s", + (response_id,), + ) + + return SimpleNamespace(thread_id=thread_id, response_id=response_id, sql=sql, row=row) + + +def _json(value) -> str: + import json + + return json.dumps(value) + + +async def _backfill(thread_id: str) -> Counter: + stats: Counter = Counter() + mismatched: list[str] = [] + await backfill._backfill_thread(thread_id, stats, mismatched) + assert mismatched == [] + return stats + + +async def _lines(thread_id: str): + """The thread's turn 0 as a replay projects it now, from its rows.""" + from src.server.database.conversation import get_replay_thread_data + from src.server.services.history.replay import project_turns + + rows = await get_replay_thread_data(thread_id) + lines = await project_turns(rows, TIP, turn_indexes=[0], reproject=True) + return lines[0] + + +async def test_a_backfill_keeps_live_keys_runs_once_and_ends_event_reads(turn): + from src.server.database.conversation import get_replay_responses + + before = await _lines(turn.thread_id) + + stats = await _backfill(turn.thread_id) + + assert (stats["turns/written"], stats["turns/refreshed"]) == (1, 1) + facts = (await turn.row())["replay_facts"] + assert {k: facts[k] for k in RUN_FACTS} == RUN_FACTS + assert facts["legacy"]["v"] == legacy.FACTS_VERSION + assert [s["event"] for s in facts["legacy"]["merge"]["signals"]] == ["context_window"] + + # A rerun finds nothing pending and leaves the row's version alone. + version = (await turn.row())["xmin"] + assert await _backfill(turn.thread_id) == Counter() + assert (await turn.row())["xmin"] == version + + backfilled = await get_replay_responses( + [turn.response_id], legacy_facts=legacy.FACTS_VERSION + ) + assert backfilled[turn.response_id]["sse_events"] is None + verbatim = await get_replay_responses([turn.response_id]) + assert verbatim[turn.response_id]["sse_events"] == STORED + + # With the column gone the turn replays the same. + await turn.sql( + "UPDATE conversation_responses SET sse_events = NULL " + "WHERE conversation_response_id = %s", + (turn.response_id,), + ) + assert await _lines(turn.thread_id) == before + + +async def test_a_row_without_stored_events_is_marked_known_empty(turn): + await turn.sql( + "UPDATE conversation_responses SET sse_events = NULL " + "WHERE conversation_response_id = %s", + (turn.response_id,), + ) + + await _backfill(turn.thread_id) + + facts = (await turn.row())["replay_facts"] + assert facts["legacy"] == {"v": legacy.FACTS_VERSION} + assert await backfill._pending_rows(turn.thread_id) == {} + + +@pytest.mark.parametrize("events", [None, STORED], ids=["no_events", "with_events"]) +async def test_a_turn_without_a_boundary_is_marked_only_when_it_stored_nothing( + turn, seed_response, events +): + """A turn the checkpoints do not cover yet replays as its stub; facts + known empty would hide its events from the projection that covers it.""" + stub_id = str(uuid.uuid4()) + await seed_response(stub_id, turn.thread_id, 1, status="interrupted", sse_events=events) + + stats = await _backfill(turn.thread_id) + + facts = await turn.sql( + "SELECT replay_facts FROM conversation_responses WHERE conversation_response_id = %s", + (stub_id,), + ) + if events is None: + assert stats["turns/stub"] == 1 + assert facts["replay_facts"]["legacy"] == {"v": legacy.FACTS_VERSION} + assert set(await backfill._pending_rows(turn.thread_id)) == set() + else: + assert stats["turns/stub_with_events"] == 1 + assert "legacy" not in (facts["replay_facts"] or {}) + assert set(await backfill._pending_rows(turn.thread_id)) == {1} + + +async def test_a_row_written_meanwhile_is_left_for_the_next_run(turn, monkeypatch): + from src.server.database.conversation.responses import append_sse_event + + real_write = backfill._write_facts + + async def after_a_live_append(rows): + await append_sse_event(turn.thread_id, LATER) + return await real_write(rows) + + monkeypatch.setattr(backfill, "_write_facts", after_a_live_append) + stats = await _backfill(turn.thread_id) + + assert (stats["turns/written"], stats["turns/moved"]) == (0, 1) + assert stats["turns/refreshed"] == 0 + assert "legacy" not in (await turn.row())["replay_facts"] + + monkeypatch.setattr(backfill, "_write_facts", real_write) + assert (await _backfill(turn.thread_id))["turns/written"] == 1 + signals = (await turn.row())["replay_facts"]["legacy"]["merge"]["signals"] + assert [s["event"] for s in signals] == ["context_window", "provenance"] + + +async def test_facts_that_would_replay_differently_are_never_written(turn, monkeypatch): + """A row cannot hold a NUL, so facts carrying one would lose it.""" + from src.server.services.history.replay import stored_merge + + real_derive = stored_merge.derive_turn + + def derive_with_a_nul(*args, **kwargs): + facts = real_derive(*args, **kwargs) + facts["merge"]["signals"][0]["data"]["kind"] = "tool\x00args" + return facts + + monkeypatch.setattr(stored_merge, "derive_turn", derive_with_a_nul) + stats: Counter = Counter() + mismatched: list[str] = [] + await backfill._backfill_thread(turn.thread_id, stats, mismatched) + + assert mismatched == [f"{turn.thread_id} turn 0"] + assert (stats["turns/mismatch"], stats["turns/written"]) == (1, 0) + assert (await turn.row())["replay_facts"] == RUN_FACTS + assert set(await backfill._pending_rows(turn.thread_id)) == {0} + + +@pytest.mark.parametrize("writer", ["append", "replace"]) +async def test_a_write_to_the_stored_events_takes_the_facts_away(turn, writer): + from src.server.database.conversation.responses import ( + append_sse_event, + replace_agent_events, + ) + + await _backfill(turn.thread_id) + assert "legacy" in (await turn.row())["replay_facts"] + + if writer == "append": + assert await append_sse_event(turn.thread_id, LATER) + else: + assert await replace_agent_events( + turn.response_id, + [{"event": "message_chunk", "data": {"agent": "task:k1", "content": "t"}}], + ) + + facts = (await turn.row())["replay_facts"] + assert facts == RUN_FACTS + assert set(await backfill._pending_rows(turn.thread_id)) == {0} + + +async def test_credit_pauses_are_stamped_once(turn, seed_response): + from src.server.contracts.status import INTERRUPT_REASON_CREDIT_PAUSE + from src.server.database.automation_executions import get_settling_run + + def pause(message): + request = {"type": INTERRUPT_REASON_CREDIT_PAUSE} + if message is not None: + request["message"] = message + return [{"event": "interrupt", "data": {"action_requests": [request]}}] + + ids = {} + for turn_index, (name, events, metadata) in enumerate( + [ + ("worded", pause("Out of credits."), None), + ("wordless", pause(None), None), + ("finalized", pause("Older words."), {"credit_pause_message": "Stamped."}), + ], + start=1, + ): + ids[name] = str(uuid.uuid4()) + await seed_response( + ids[name], turn.thread_id, turn_index, status="interrupted", sse_events=events + ) + await turn.sql( + "UPDATE conversation_responses SET interrupt_reason = %s, " + "metadata = metadata || %s::jsonb WHERE conversation_response_id = %s", + (INTERRUPT_REASON_CREDIT_PAUSE, _json(metadata or {}), ids[name]), + ) + + args = SimpleNamespace(thread=[turn.thread_id], workspace=None, user=None, days=None) + pending = await backfill._pause_rows(args) + assert sorted(pending) == sorted([ids["worded"], ids["wordless"]]) + + assert await backfill._stamp_pauses(pending) == (2, 1) + assert await backfill._pause_rows(args) == [] + assert await backfill._stamp_pauses(pending) == (0, 0) + + stamps = {} + for name, response_id in ids.items(): + run = await get_settling_run(response_id) + assert run["sse_events"] is None # settling no longer reads them + stamps[name] = run["metadata"]["credit_pause_message"] + assert stamps == {"worded": "Out of credits.", "wordless": None, "finalized": "Stamped."} diff --git a/tests/integration/test_main_state_db.py b/tests/integration/test_main_state_db.py new file mode 100644 index 000000000..1d9b50e27 --- /dev/null +++ b/tests/integration/test_main_state_db.py @@ -0,0 +1,115 @@ +"""The main agent's private fields survive a write through another graph, on +the Postgres saver, which keeps a primitive channel value inside the +checkpoint row: a graph that leaves such a field undeclared writes the next +checkpoint without it. The writers are the history reader's ui-record append +and a manual compaction, whose build leaves the subagent switch out. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from langchain.agents import create_agent +from langchain.agents.middleware import AgentMiddleware +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.graph import START, StateGraph +from langsmith import tracing_context + +from ptc_agent.agent.main_state import MainAgentState +from ptc_agent.agent.state import DeltaAgentState +from src.server.database.conversation import threads_write +from src.server.services.history.reader import CheckpointHistoryReader + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + +THREAD = "5e1f0000-0000-4000-8000-00000000f1e1" +KEPT = {"_subagents_trimmed": True, "_notes_calls_trimmed": 7} + + +def _config() -> dict[str, Any]: + return {"configurable": {"thread_id": THREAD}} + + +def _graph(schema: type, saver: Any) -> Any: + return ( + StateGraph(schema) + .add_node("noop", lambda state: {}) + .add_edge(START, "noop") + .compile(checkpointer=saver) + ) + + +class _Sets(AgentMiddleware): + """Sets both fields, as a trim's carries do.""" + + def before_agent(self, state: Any, runtime: Any) -> dict[str, Any]: + return dict(KEPT) + + +def _agent(saver: Any, *middleware: AgentMiddleware) -> Any: + """A main-agent build with no subagent switch, as thread maintenance's is.""" + return create_agent( + GenericFakeChatModel(messages=iter([AIMessage("a", id="a-1")])), + tools=[], + middleware=list(middleware), + checkpointer=saver, + state_schema=MainAgentState, + ) + + +async def _seeded(saver: Any) -> Any: + """A thread whose last turn left both fields set.""" + await saver.adelete_thread(THREAD) + with tracing_context(enabled=False): + await _agent(saver, _Sets()).ainvoke( + {"messages": [HumanMessage("q", id="h-1")]}, _config() + ) + + +async def _kept(saver: Any) -> dict[str, Any]: + values = (await _graph(MainAgentState, saver).aget_state(_config())).values + return {key: values.get(key) for key in KEPT} + + +async def test_a_graph_without_them_erases_them(test_db_pool) -> None: + """The hazard the schema closes, so the tests below can see it.""" + saver = AsyncPostgresSaver(test_db_pool) + await _seeded(saver) + await _graph(DeltaAgentState, saver).aupdate_state(_config(), {"ui": []}) + assert await _kept(saver) == dict.fromkeys(KEPT) + + +async def test_a_ui_record_append_keeps_them(test_db_pool, monkeypatch) -> None: + saver = AsyncPostgresSaver(test_db_pool) + await _seeded(saver) + + async def advance(thread_id, *, from_checkpoint_id, to_checkpoint_id): + return True + + monkeypatch.setattr(threads_write, "advance_thread_checkpoint_id", advance) + await CheckpointHistoryReader(saver).append_ui_record(THREAD, "image_map", {"a": "b"}) + state = await _graph(MainAgentState, saver).aget_state(_config()) + assert [r["name"] for r in state.values["ui"]] == ["image_map"] + assert await _kept(saver) == KEPT + + +async def test_a_manual_compaction_keeps_them(test_db_pool, monkeypatch) -> None: + from src.server.handlers.thread_maintenance import _update_graph_state + + saver = AsyncPostgresSaver(test_db_pool) + await _seeded(saver) + graph = _agent(saver) + + async def advance(thread_id, *, from_checkpoint_id, to_checkpoint_id): + return True + + monkeypatch.setattr(threads_write, "advance_thread_checkpoint_id", advance) + state = await graph.aget_state(_config()) + await _update_graph_state( + graph, state, {"_offloaded_tool_call_ids": {"call-1"}}, THREAD, "offload" + ) + assert (await graph.aget_state(_config())).values["_offloaded_tool_call_ids"] == {"call-1"} + assert await _kept(saver) == KEPT diff --git a/tests/integration/test_repair_crossed_threads_db.py b/tests/integration/test_repair_crossed_threads_db.py new file mode 100644 index 000000000..c3ab67ee8 --- /dev/null +++ b/tests/integration/test_repair_crossed_threads_db.py @@ -0,0 +1,510 @@ +"""scripts/ops/repair_crossed_threads.py on threads that crossed into +DeltaChannel the way production's did. + +A crossed thread starts with turns under a plain ``add_messages`` state, each +fanning out to parallel tool nodes, which that reducer applied in node order. +Its first turn under DeltaChannel loaded by replaying the stored writes from +the root, which orders a checkpoint's writes by task id, so the tool results +of each legacy turn came back in another order and stayed in it. Turn slices +are cut from the legacy lists and keep the original order, so the window's +coverage check never holds the head. Six parallel tools make a turn whose +task ids happen to sort in node order one in 720, and there are two. + +Twin threads are built from the same script, one repaired and one left +crossed, and run the same turns after it through compaction with the window +on: what the model and the summary model are sent must not tell them apart. +""" + +from __future__ import annotations + +import uuid +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Annotated, Any, TypedDict + +import psycopg +import pytest +import pytest_asyncio +from langchain.agents import create_agent +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import add_messages +from psycopg.rows import dict_row, tuple_row + +import src.server.utils.checkpointer # noqa: F401 (installs the walk guard) +from ptc_agent.agent.main_state import MainAgentState +from ptc_agent.agent.middleware.compaction.compact import Summarizer +from ptc_agent.agent.middleware.compaction.middleware import CompactionMiddleware +from ptc_agent.agent.middleware.compaction.notes import notes_window_carry +from ptc_agent.agent.middleware.compaction.types import OffloadSettings +from ptc_agent.agent.middleware.compaction.window import window_cut +from ptc_agent.agent.middleware.subagent_switch import subagents_window_carry +from ptc_agent.agent.transcript import Window +from scripts.ops import repair_crossed_threads as repair +from scripts.ops import thin_message_snapshots as thinning +from scripts.ops._delta_loads import Mode, Status +from tests.unit.middleware.compaction import window_harness as h + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + +TOOLS = "abcdef" +LEGACY_TURNS = 2 +DELTA_TURNS = 5 +#: Messages in view, the system prompt counted, at which a summary runs. +THRESHOLD = 28 +KEEP = 4 +OFFLOAD = OffloadSettings(keep_messages=KEEP, max_length=200, idle_seconds=0.0) +SETTINGS = repair.Settings(keep_messages=KEEP, offload=OFFLOAD) +VERIFY_ALL = repair.Options(verify_all=True) + + +class _Legacy(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +def _number(state) -> int: + return sum(isinstance(m, HumanMessage) for m in state["messages"]) + + +def _legacy_graph(saver): + """A turn that calls every tool at once and answers once they are back.""" + + def ask(state): + n = _number(state) + calls = [{"name": f"tool_{x}", "id": f"call-{n}-{x}", "args": {"n": n}} for x in TOOLS] + return {"messages": [AIMessage("", id=f"ai-{n}-ask", tool_calls=calls)]} + + def tool(x: str): + def run(state): + n = _number(state) + return { + "messages": [ + ToolMessage(f"result {n}{x}", tool_call_id=f"call-{n}-{x}", id=f"tool-{n}-{x}") + ] + } + + return run + + def answer(state): + n = _number(state) + return {"messages": [AIMessage(f"Answer {n}.", id=f"ai-{n}-end")]} + + graph = StateGraph(_Legacy).add_node("ask", ask).add_node("answer", answer) + graph.add_edge(START, "ask") + for x in TOOLS: + graph.add_node(f"tool_{x}", tool(x)).add_edge("ask", f"tool_{x}") + graph.add_edge([f"tool_{x}" for x in TOOLS], "answer").add_edge("answer", END) + return graph.compile(checkpointer=saver) + + +def _agent(saver, thread_id: str, model, summary, *, threshold: int, asked: list | None = None): + """Compaction as the main agent's stack holds it, with the window on when + ``asked`` collects the answers of the server's own coverage check.""" + compaction = CompactionMiddleware( + Summarizer(summary, limit=10**9, counter=len), + token_threshold=threshold, + keep_messages=KEEP, + offload=OFFLOAD, + backend=_backend(), + workspace_id="ws-1", + ) + if asked is not None: + from src.server.services.history.window import runs_held + + async def coverage(head: Window) -> bool: + held = await runs_held(thread_id, head) + asked.append((head.runs, held)) + return held + + compaction = compaction.with_window( + coverage, [notes_window_carry(thread_id), subagents_window_carry] + ) + return create_agent( + model, + system_prompt="You are the test analyst.", + tools=[], + middleware=[compaction], + checkpointer=saver, + state_schema=MainAgentState, + ) + + +class _Mount: + """The computer's file mount as far as compaction uses it: a transcript + saved the way the server saves one mid-turn.""" + + async def save_transcript(self, target, messages, *, window: Window) -> bool: + from src.server.services import transcripts + + return await transcripts.save_live(target, messages, window=window) + + +def _backend() -> Any: + mount = _Mount() + + async def settled_livefs(workspace_id=None): + return mount + + return SimpleNamespace( + livefs=mount, settled_livefs=settled_livefs, normalize_path=lambda path: path + ) + + +def _sent(requests: list, thread_id: str) -> list[str]: + """Requests as text, the thread's own ids and run-time ids normalized.""" + normalize = h.Normalizer() + return [ + normalize(h.dump(r).replace(thread_id, "").replace(thread_id[:8], "")) + for r in requests + ] + + +async def _turn(saver, thread_id: str, graph, text: str) -> str: + """One turn the way the server runs it: its rows started, the graph run + with the turn stamped on its config, the run finalized at its tip. The + input carries no id, as production's legacy inputs did not.""" + from src.server.database.runs.lifecycle import QuerySpec, RunOutcome, finalize_run, start_run + + run_id = str(uuid.uuid4()) + row = await start_run( + run_id=run_id, + thread_id=thread_id, + request_key=str(uuid.uuid4()), + query=QuerySpec(query_id=str(uuid.uuid4()), content=text, query_type="initial"), + metadata={"msg_type": "ptc"}, + ) + config = { + "configurable": {"thread_id": thread_id}, + "metadata": {"run_id": run_id, "turn_index": row["turn_index"]}, + } + await graph.ainvoke({"messages": [HumanMessage(text)]}, config) + tip = await _newest(saver, thread_id) + await finalize_run( + run_id=run_id, + thread_id=thread_id, + outcome=RunOutcome(status="completed"), + checkpoint_id=tip, + ) + return tip + + +async def _newest(saver, thread_id: str) -> str: + newest = await saver.aget_tuple({"configurable": {"thread_id": thread_id}}) + return newest.config["configurable"]["checkpoint_id"] + + +async def _project(thread_id: str) -> None: + """Every turn's slice, as ``backfill_turn_slices`` stores them.""" + from src.server.services.history.replay import load_thread_inputs, project_turns + + rows, tip = await load_thread_inputs(thread_id) + await project_turns(rows, tip, last_n_turns=10**6) + + +async def _crossed(saver, test_db_uri: str, workspace_id: str) -> str: + """Legacy turns, then turns under DeltaChannel compacting as they go, the + first of which crossed; then the rollout before the repair: the stored + lists kept as seeds, and every turn's slice stored.""" + from src.server.database.conversation import create_thread + + thread_id = str(uuid.uuid4()) + await create_thread( + conversation_thread_id=thread_id, + workspace_id=workspace_id, + current_status="completed", + msg_type="ptc", + ) + legacy = _legacy_graph(saver) + for n in range(1, LEGACY_TURNS + 1): + await _turn(saver, thread_id, legacy, f"Question {n}") + model, summary = h.RecordingModel(), h.SummaryModel() + for n in range(LEGACY_TURNS + 1, LEGACY_TURNS + DELTA_TURNS + 1): + model.script = [AIMessage(f"Answer {n}.", id=f"ai-{n}-end")] + graph = _agent(saver, thread_id, model, summary, threshold=THRESHOLD) + tip = await _turn(saver, thread_id, graph, f"Question {n}") + assert summary.requests, "the delta turns never compacted" + thinned = await thinning.thin_thread( + test_db_uri, thread_id, tip, Mode.APPLY, thinning.Options(keep_all=True, verify_all=True) + ) + assert not thinned.failed + await _project(thread_id) + return thread_id + + +async def _head(thread_id: str) -> tuple[Window, list[AnyMessage], list[AnyMessage]]: + """What the window's next trim would drop at the tip, the tip's messages + it covers, and the same messages as the stored slices hold them.""" + from src.server.services.history import window as window_reads + from src.server.services.history.reader import CheckpointHistoryReader + + reader = CheckpointHistoryReader.get_instance() + values = (await reader.aget_state(thread_id)).values + messages = values["messages"] + cut = window_cut(messages, values.get("_summarization_event")) + head = Window.of(values).extend(messages[:cut]) + anchors, _ = await reader.aget_turn_anchors(thread_id) + stored = await window_reads._head(thread_id, anchors, head.messages, extract=False) + return head, messages[:cut], stored + + +def _keys(messages) -> list[str]: + return [repair._key(m) for m in messages] + + +async def _repair(test_db_uri: str, thread_id: str, mode: Mode = Mode.APPLY): + return await repair.repair_thread(test_db_uri, thread_id, mode, VERIFY_ALL, SETTINGS) + + +async def _rows(test_db_uri: str, thread_id: str) -> list[tuple]: + """Every row of the thread, with its row version.""" + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + out = [] + for table, key in ( + ("checkpoints", "thread_id"), + ("checkpoint_blobs", "thread_id"), + ("checkpoint_writes", "thread_id"), + ("conversation_threads", "conversation_thread_id::text"), + ("turn_slices", "conversation_thread_id::text"), + ): + cur = await conn.execute( + f"SELECT md5(t::text), xmin::text FROM {table} t WHERE {key} = %s ORDER BY 1", + (thread_id,), + ) + out.append((table, await cur.fetchall())) + return out + + +async def _pointer(test_db_uri: str, thread_id: str) -> str: + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + cur = await conn.execute( + "SELECT latest_checkpoint_id FROM conversation_threads" + " WHERE conversation_thread_id = %s", + (thread_id,), + ) + return (await cur.fetchone())[0] + + +@pytest_asyncio.fixture(loop_scope="session") +async def saver(test_db_pool, patched_get_db_connection, monkeypatch): + """The test pool as the server's checkpointer. The transcript store reads + rows as tuples, as the app pool returns them, not as the test pool's dicts.""" + from src.server.app import setup + from src.server.services.history.reader import CheckpointHistoryReader + + @asynccontextmanager + async def tuple_rows(): + async with test_db_pool.connection() as conn: + conn.row_factory = tuple_row + try: + yield conn + finally: + conn.row_factory = dict_row + + monkeypatch.setattr("src.server.database.thread_transcripts.get_db_connection", tuple_rows) + saver = AsyncPostgresSaver(test_db_pool) + monkeypatch.setattr(setup, "checkpointer", saver) + CheckpointHistoryReader.reset_instance() + yield saver + CheckpointHistoryReader.reset_instance() + + +@pytest_asyncio.fixture(loop_scope="session") +async def twins(saver, test_db_uri, seed_workspace): + """Two threads from one script: ``repaired`` is repaired, ``crossed`` is not.""" + workspace_id = str(seed_workspace["workspace_id"]) + repaired = await _crossed(saver, test_db_uri, workspace_id) + crossed = await _crossed(saver, test_db_uri, workspace_id) + yield repaired, crossed + for thread_id in (repaired, crossed): + await saver.adelete_thread(thread_id) + + +async def test_a_crossed_thread_is_refused_until_repaired(saver, test_db_uri, twins): + thread_id, _ = twins + head, tip_head, stored = await _head(thread_id) + from src.server.services.history.window import runs_held + + # The crossing as production's: the slices hold the tip's head messages in + # another order, so the window's coverage check refuses every trim. + legacy = LEGACY_TURNS * (len(TOOLS) + 3) + assert head.runs > LEGACY_TURNS and head.messages > legacy + assert sorted(_keys(tip_head)) == sorted(_keys(stored)) + assert _keys(tip_head[:legacy]) != _keys(stored[:legacy]) + assert _keys(tip_head[legacy:]) == _keys(stored[legacy:]) + assert not await runs_held(thread_id, head) + + before = await _rows(test_db_uri, thread_id) + planned = await _repair(test_db_uri, thread_id, Mode.DRY_RUN) + assert planned.status is Status.PLANNED, planned.reason + report = planned.report + assert (report.crossed, report.refused, report.cut) == ("yes", "", head.messages) + assert report.requests_match and report.requests == 4 + assert report.span[1] < report.replayed + rehearsed = await _repair(test_db_uri, thread_id, Mode.REHEARSE) + assert rehearsed.status is Status.REHEARSED and not rehearsed.failed + assert await _rows(test_db_uri, thread_id) == before + + tip = await _newest(saver, thread_id) + applied = await _repair(test_db_uri, thread_id) + assert applied.status is Status.WRITTEN and not applied.failed, applied.reason + assert applied.committed_match and all(check.match for check in applied.checks) + assert applied.report.requests_match + assert applied.report.reprojected_same + # The trim's own writes are not left on the old tip, where a load of any + # other child of it would replay them. + assert dict(await _rows(test_db_uri, thread_id))["checkpoint_writes"] == dict(before)[ + "checkpoint_writes" + ] + + # The window a trim of an uncrossed thread would leave: the slices' head. + from src.server.services.history.reader import CheckpointHistoryReader + + reader = CheckpointHistoryReader.get_instance() + values = (await reader.aget_state(thread_id)).values + window = Window.of(values) + canonical = Window().extend(stored) + assert (window.runs, window.messages, window.digest, dict(window.earlier)) == ( + canonical.runs, + canonical.messages, + canonical.digest, + dict(canonical.earlier), + ) + assert await runs_held(thread_id, window) + kept = (await reader.aget_state(thread_id, tip)).values["messages"][head.messages :] + assert [m.id for m in values["messages"]] == [m.id for m in kept] + assert _keys(values["messages"]) == _keys(kept) + + # A rerun finds nothing to do and writes nothing. + after = await _rows(test_db_uri, thread_id) + for mode in (Mode.DRY_RUN, Mode.APPLY): + again = await _repair(test_db_uri, thread_id, mode) + assert again.status is Status.NOTHING and again.report.crossed == "no" + assert await _rows(test_db_uri, thread_id) == after + + +async def test_a_last_turn_that_fails_to_store_is_reported( + saver, test_db_uri, seed_workspace, monkeypatch +): + """The repair commits, but the turn it moved is not stored again, so the + old row, with the messages it held, is no proof the turn projects.""" + + async def refused(rows): + raise RuntimeError("store refused") + + thread_id = await _crossed(saver, test_db_uri, str(seed_workspace["workspace_id"])) + monkeypatch.setattr( + "src.server.database.conversation.turn_slices.upsert_turn_slices", refused + ) + applied = await _repair(test_db_uri, thread_id) + assert applied.status is Status.WRITTEN and applied.committed_match, applied.reason + assert applied.report.reprojected + assert applied.report.reprojected_same is False + await saver.adelete_thread(thread_id) + + +async def test_the_repair_changes_nothing_the_models_are_sent(saver, test_db_uri, twins): + """The same turns after the repair on both twins, a summary among them: + the repaired one trims where the crossed one is refused, and the model + and the summary model are sent the same on both.""" + from src.server.services.history.reader import CheckpointHistoryReader + + repaired, crossed = twins + result = await _repair(test_db_uri, repaired) + assert result.status is Status.WRITTEN and not result.failed, result.reason + reader = CheckpointHistoryReader.get_instance() + runs = Window.of((await reader.aget_state(repaired)).values).runs + + sent: dict[str, tuple[list, list, list]] = {} + windows: dict[str, list[int]] = {} + start = LEGACY_TURNS + DELTA_TURNS + 1 + for thread_id in (repaired, crossed): + model, summary, asked = h.RecordingModel(), h.SummaryModel(), [] + windows[thread_id] = [] + # A plain turn, one that summarizes, and one that trims to that summary. + for offset, threshold in enumerate((10**6, 0, 10**6)): + n = start + offset + model.script = [AIMessage(f"Answer {n}.", id=f"ai-{n}-end")] + graph = _agent(saver, thread_id, model, summary, threshold=threshold, asked=asked) + await _turn(saver, thread_id, graph, f"Question {n}") + await _project(thread_id) + values = (await reader.aget_state(thread_id)).values + windows[thread_id].append(Window.of(values).runs) + sent[thread_id] = ( + _sent(model.requests, thread_id), _sent(summary.requests, thread_id), asked + ) + + assert sent[repaired][:2] == sent[crossed][:2] + assert len(sent[repaired][0]) == 3 and len(sent[repaired][1]) == 1 + # The summary cites the transcript, which the repaired thread renders + # with its trimmed turns read back from their slices. + assert "" in sent[repaired][1][0] + # Only the trim after the summary asks of the repaired thread, and is held; + # every trim the crossed one would make is refused. + assert [held for _, held in sent[repaired][2]] == [True] + assert sent[crossed][2] and not any(held for _, held in sent[crossed][2]) + assert windows[repaired][:2] == [runs, runs] and windows[repaired][2] > runs + assert windows[crossed] == [0, 0, 0] + + +async def test_a_thread_is_refused_while_a_run_is_live_or_a_slice_is_missing( + saver, test_db_uri, seed_workspace +): + from src.server.database.runs.lifecycle import QuerySpec, RunOutcome, finalize_run, start_run + + thread_id = await _crossed(saver, test_db_uri, str(seed_workspace["workspace_id"])) + before = await _rows(test_db_uri, thread_id) + pointer = await _pointer(test_db_uri, thread_id) + + run_id = str(uuid.uuid4()) + await start_run( + run_id=run_id, + thread_id=thread_id, + request_key=str(uuid.uuid4()), + query=QuerySpec(query_id=str(uuid.uuid4()), content="live", query_type="initial"), + ) + live = await _repair(test_db_uri, thread_id) + assert live.report.refused == "a run is in progress" + await finalize_run(run_id=run_id, thread_id=thread_id, outcome=RunOutcome(status="cancelled")) + + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + await conn.execute( + "DELETE FROM turn_slices s USING conversation_responses r" + " WHERE s.conversation_response_id = r.conversation_response_id" + " AND s.conversation_thread_id = %s AND r.turn_index = 0", + (thread_id,), + ) + missing = await _repair(test_db_uri, thread_id) + assert missing.report.refused.startswith("a turn before the cut has no stored slice") + + # Neither wrote a checkpoint or moved the pointer; the run and the slice + # are the test's own changes to the other rows. + after = dict(await _rows(test_db_uri, thread_id)) + for table, rows in before: + if table.startswith("checkpoint"): + assert after[table] == rows, table + assert await _pointer(test_db_uri, thread_id) == pointer + await saver.adelete_thread(thread_id) + + +async def test_a_kept_message_stored_without_an_id_is_written_with_the_planned_one( + saver, test_db_uri, seed_workspace +): + """A message injected through ``aupdate_state`` is stored without an id, + and the trim stamps one on every kept message: the repair writes the id + its plan checked, not a second fresh one.""" + from src.server.services.history.reader import CheckpointHistoryReader + from src.server.utils.checkpoint_helpers import update_at_commit + + thread_id = await _crossed(saver, test_db_uri, str(seed_workspace["workspace_id"])) + graph = _agent(saver, thread_id, h.RecordingModel(), h.SummaryModel(), threshold=10**6) + tip = await _newest(saver, thread_id) + assert await update_at_commit(graph, thread_id, tip, {"messages": [HumanMessage("A note.")]}) + reader = CheckpointHistoryReader.get_instance() + assert (await reader.aget_state(thread_id)).values["messages"][-1].id is None + + applied = await _repair(test_db_uri, thread_id) + assert applied.status is Status.WRITTEN and not applied.failed, applied.reason + note = (await reader.aget_state(thread_id)).values["messages"][-1] + assert note.content == "A note." and note.id is not None + await saver.adelete_thread(thread_id) diff --git a/tests/integration/test_replace_agent_events_db.py b/tests/integration/test_replace_agent_events_db.py index ae6cc7a87..ac04a1445 100644 --- a/tests/integration/test_replace_agent_events_db.py +++ b/tests/integration/test_replace_agent_events_db.py @@ -2,7 +2,7 @@ The strip and append run in SQL under a row lock, so only a real database can show that a concurrent append survives, that other agents' rows keep their -order, and that provenance is re-derived from the rows the write leaves. +order, and that the written lanes' provenance is rewritten from the batch. """ from __future__ import annotations @@ -155,8 +155,25 @@ async def test_append_waiting_on_the_lock_survives( assert events.count(STALE_TASK_ROW) == 1 -async def test_provenance_follows_the_rows_the_write_leaves(response_row, db_conn): +async def test_a_lane_write_leaves_other_lanes_records(response_row, db_conn): + """The main lane's records are the finalize's to write; a subagent's + archive replaces its own lane and nothing else.""" from src.server.database.conversation.responses import replace_agent_events + from src.server.database.provenance import ( + MAIN_PROVENANCE_LANE, + sync_provenance_for_response, + ) + + async with db_conn.transaction(): + await sync_provenance_for_response( + db_conn, + conversation_response_id=response_row, + conversation_thread_id=await _thread_of(db_conn, response_row), + turn_index=0, + sse_events=[MAIN_PROVENANCE], + strict=True, + lanes=(MAIN_PROVENANCE_LANE,), + ) assert await replace_agent_events(response_row, [TASK_PROVENANCE]) @@ -166,6 +183,15 @@ async def test_provenance_follows_the_rows_the_write_leaves(response_row, db_con ] +async def _thread_of(conn, response_id: str) -> str: + cur = await conn.execute( + "SELECT conversation_thread_id FROM conversation_responses " + "WHERE conversation_response_id = %s", + (response_id,), + ) + return str((await cur.fetchone())["conversation_thread_id"]) + + async def test_a_stripped_agents_provenance_leaves_with_it( make_response_row, db_conn ): diff --git a/tests/integration/test_replay_facts_db.py b/tests/integration/test_replay_facts_db.py new file mode 100644 index 000000000..c29089d72 --- /dev/null +++ b/tests/integration/test_replay_facts_db.py @@ -0,0 +1,161 @@ +"""Subagent run facts against real Postgres. + +A run's facts are learned after its settle CAS, by two writers that can land +together (the archive's lane images, the steering sweep's returns), and replay +keys a turn's stored lines on the facts of the runs it launched, so only a +real database can show that both writes keep theirs and that the facts reach +the launching turn's index. +""" + +from __future__ import annotations + +import asyncio +import uuid + +import pytest +import pytest_asyncio + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + + +@pytest_asyncio.fixture(loop_scope="session") +async def launched_run(seed_workspace, patched_get_db_connection): + """A thread whose second turn launched one subagent run.""" + from src.server.database.conversation import create_thread + from src.server.database.runs import subagent_runs + from src.server.database.runs.lifecycle import start_run + + thread_id = str(uuid.uuid4()) + await create_thread( + conversation_thread_id=thread_id, + workspace_id=str(seed_workspace["workspace_id"]), + current_status="in_progress", + msg_type="ptc", + ) + response_id = str(uuid.uuid4()) + await start_run( + run_id=response_id, + thread_id=thread_id, + request_key=str(uuid.uuid4()), + turn_index=1, + metadata={"msg_type": "ptc", "user_id": seed_workspace["user_id"]}, + ) + task_run_id = str(uuid.uuid4()) + await subagent_runs.start_task_run( + task_run_id=task_run_id, + thread_id=thread_id, + task_id="k1", + cause="init", + parent_run_id=response_id, + ) + await subagent_runs.finalize_task_run(task_run_id=task_run_id, status="completed") + return thread_id, response_id, task_run_id + + +async def _run_facts(thread_id): + """The run facts a replay of the thread reads, with its other rows.""" + from src.server.database.conversation import get_replay_thread_data + + return (await get_replay_thread_data(thread_id)).run_facts + + +async def _run(thread_id, task_id, parent_run_id, *, predecessor=None): + from src.server.database.runs import subagent_runs + + task_run_id = str(uuid.uuid4()) + await subagent_runs.start_task_run( + task_run_id=task_run_id, + thread_id=thread_id, + task_id=task_id, + cause="resume" if predecessor else "init", + parent_run_id=parent_run_id, + predecessor_run_id=predecessor, + ) + await subagent_runs.finalize_task_run(task_run_id=task_run_id, status="completed") + return task_run_id + + +async def test_facts_merge_per_key_and_name_the_launching_turn(launched_run): + from src.server.database import replay_facts + + thread_id, response_id, task_run_id = launched_run + returned = {"content": "also check AMD", "input_id": "in-1", "reason": "run_ended"} + + imaged, steered = await asyncio.gather( + replay_facts.append_lane_images( + thread_id, response_id, "k1", {"/home/workspace/a.png": "https://cdn/a.png"} + ), + replay_facts.append_returned_steering(task_run_id, [returned]), + ) + assert await replay_facts.append_lane_images( + thread_id, response_id, "k1", {"/home/workspace/b.png": "https://cdn/b.png"} + ) + + assert imaged is True + assert str(steered["thread_id"]) == thread_id + assert str(steered["parent_run_id"]) == response_id + + rows = await _run_facts(thread_id) + assert len(rows) == 1 + assert rows[0]["turn_index"] == 1 + assert str(rows[0]["task_run_id"]) == task_run_id + assert rows[0]["replay_facts"] == { + "images": { + "/home/workspace/a.png": "https://cdn/a.png", + "/home/workspace/b.png": "https://cdn/b.png", + }, + "steering_returned": [returned], + } + + +async def test_lane_images_reach_the_turns_runs_of_that_task_only(launched_run): + """A task resumed within the turn has two runs there, both of which the + lane's events may come from; another task's run and a later turn's run of + the same task keep their own facts.""" + from src.server.database import replay_facts + from src.server.database.runs.lifecycle import RunOutcome, finalize_run, start_run + + thread_id, response_id, first = launched_run + resumed = await _run(thread_id, "k1", response_id, predecessor=first) + await _run(thread_id, "k2", response_id) + await finalize_run( + run_id=response_id, thread_id=thread_id, outcome=RunOutcome(status="completed") + ) + later_turn = str(uuid.uuid4()) + await start_run( + run_id=later_turn, + thread_id=thread_id, + request_key=str(uuid.uuid4()), + turn_index=2, + metadata={"msg_type": "ptc"}, + ) + await _run(thread_id, "k1", later_turn, predecessor=resumed) + + images = {"/home/workspace/a.png": "https://cdn/a.png"} + assert await replay_facts.append_lane_images(thread_id, response_id, "k1", images) + + rows = await _run_facts(thread_id) + assert {str(r["task_run_id"]) for r in rows} == {first, resumed} + assert all(r["replay_facts"] == {"images": images} for r in rows) + + +async def test_nothing_to_record_writes_nothing(launched_run): + from src.server.database import replay_facts + + thread_id, response_id, task_run_id = launched_run + assert await replay_facts.append_returned_steering(task_run_id, []) is None + assert await replay_facts.append_lane_images(thread_id, response_id, "k1", {}) is False + assert await _run_facts(thread_id) == [] + + +async def test_an_unknown_run_reports_none(launched_run): + from src.server.database import replay_facts + + thread_id, _, _ = launched_run + returned = {"content": "c", "input_id": "in-1", "reason": "run_ended"} + images = {"/home/workspace/a.png": "https://cdn/a.png"} + assert await replay_facts.append_returned_steering(str(uuid.uuid4()), [returned]) is None + assert ( + await replay_facts.append_lane_images(thread_id, str(uuid.uuid4()), "k1", images) + is False + ) diff --git a/tests/integration/test_replay_walk_failure_db.py b/tests/integration/test_replay_walk_failure_db.py new file mode 100644 index 000000000..20a1a2e8a --- /dev/null +++ b/tests/integration/test_replay_walk_failure_db.py @@ -0,0 +1,129 @@ +"""A branch walk that fails in Postgres, as replay and its endpoint see it. + +The walk raises rather than falling back to the in-process walk, so replay's +own fallback is what keeps a thread readable: ``auto`` serves the stored +events, and ``checkpoint`` (no fallback by request) fails the request. +""" + +from __future__ import annotations + +import uuid +from unittest.mock import AsyncMock, patch + +import psycopg +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.graph import START, StateGraph + +from ptc_agent.agent.state import DeltaAgentState +from src.server.database.conversation.replay_rows import ThreadRows +from src.server.services.history.reader import CheckpointHistoryReader +from src.server.services.history.replay import read_replay_page +from src.server.utils import checkpoint_helpers +from tests.conftest import create_test_app + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio(loop_scope="session")] + +STORED = [{"event": "message_chunk", "data": {"content": "stored answer"}}] + + +@pytest_asyncio.fixture(loop_scope="session") +async def broken_walk(test_db_pool, monkeypatch): + """A one-turn thread whose branch walk fails on a renamed column, read + through a real reader on the real saver. Returns the thread's rows.""" + saver = AsyncPostgresSaver(test_db_pool) + graph = ( + StateGraph(DeltaAgentState) + .add_node("agent", lambda s: {"messages": [AIMessage("answer", id="ai-0")]}) + .add_edge(START, "agent") + .compile(checkpointer=saver) + ) + thread_id = str(uuid.uuid4()) + await graph.ainvoke( + {"messages": [HumanMessage("question", id="h-0")]}, + {"configurable": {"thread_id": thread_id}, "metadata": {"turn_index": 0}}, + ) + tip = (await graph.aget_state({"configurable": {"thread_id": thread_id}})).config[ + "configurable" + ]["checkpoint_id"] + monkeypatch.setattr(CheckpointHistoryReader, "_instance", CheckpointHistoryReader(saver)) + monkeypatch.setattr( + checkpoint_helpers, + "_BRANCH_BOUNDARIES_SQL", + checkpoint_helpers._BRANCH_BOUNDARIES_SQL.replace( + "parent_checkpoint_id", "parent_id" + ), + ) + thread = {"conversation_thread_id": thread_id, "latest_checkpoint_id": tip} + queries = [ + {"turn_index": 0, "content": "question", "type": "initial", "created_at": "t0"} + ] + responses = { + 0: { + "conversation_response_id": str(uuid.uuid4()), + "turn_index": 0, + "status": "completed", + "sse_events": STORED, + } + } + return ThreadRows( + thread_id=thread_id, + owner_id="test-user-123", + thread=thread, + queries=queries, + responses_by_turn=responses, + usages=[], + provenance=[], + run_facts=[], + ) + + +async def test_auto_replays_the_stored_events(broken_walk): + page, source = await read_replay_page(broken_walk) + assert source == "sse" + assert any("stored answer" in line.data_json for line in page.lines) + + +async def test_checkpoint_raises_the_walk_error(broken_walk): + with pytest.raises(psycopg.errors.UndefinedColumn): + await read_replay_page(broken_walk, source="checkpoint") + + +@pytest.mark.parametrize( + "source, status, replayed_from", + [("auto", 200, "sse"), ("checkpoint", 500, None)], +) +async def test_the_endpoint_answer(broken_walk, source, status, replayed_from): + from src.server.app.threads import router + + thread_id = broken_walk.thread_id + with ( + patch( + "src.server.app.threads.messaging.get_replay_thread_data", + new=AsyncMock(return_value=broken_walk), + ), + patch( + "src.server.services.history.snapshot.build_thread_snapshot", + new=AsyncMock(return_value=None), + ), + ): + async with AsyncClient( + transport=ASGITransport(app=create_test_app(router)), + base_url="http://test", + ) as client: + resp = await client.get( + f"/api/v1/threads/{thread_id}/messages/replay", + params={"source": source}, + ) + + assert resp.status_code == status + if replayed_from: + assert resp.headers["X-Replay-Source"] == replayed_from + assert "stored answer" in resp.text + else: + assert resp.json()["detail"].startswith( + 'Failed to replay thread: column "parent_id" does not exist' + ) diff --git a/tests/integration/test_thin_message_snapshots_db.py b/tests/integration/test_thin_message_snapshots_db.py new file mode 100644 index 000000000..6e0b6ed31 --- /dev/null +++ b/tests/integration/test_thin_message_snapshots_db.py @@ -0,0 +1,549 @@ +"""scripts/ops/thin_message_snapshots.py on threads built through the Postgres saver. + +Threads snapshot ``messages`` every second update so a few turns leave several +snapshots, and a migrated thread starts with turns under a plain +``add_messages`` state, the way production's did. Each test writes its own +thread and the conversation row that names its tip. +""" + +from __future__ import annotations + +import uuid +from typing import Annotated, TypedDict + +import psycopg +import pytest +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langgraph.channels.delta import DeltaChannel +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.graph import START, StateGraph +from langgraph.graph.message import add_messages + +import src.server.utils.checkpointer # noqa: F401 (installs the walk guard) +from ptc_agent.agent.state import DeltaAgentState, messages_delta_reducer +from scripts.ops import thin_message_snapshots as thinning +from scripts.ops._delta_loads import ( + Loader, + Mode, + Rewrite, + Status, + Timeouts, + Walk, + messages_digest, + read_tree, + rewrite_thread, + unsupported_saver, +) + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + +# Every snapshot is a candidate: no spacing keep, no recent-turn keep. +_EVERYTHING = thinning.Options(spacing=10**6, keep_recent_turns=0, verify_all=True) + + +class _PreDeltaState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + +class _SnapshotEverySecondUpdate(TypedDict): + messages: Annotated[ + list[AnyMessage], DeltaChannel(messages_delta_reducer, snapshot_frequency=2) + ] + + +def _reply(state) -> dict: + question = state["messages"][-1] + return {"messages": [AIMessage(f"re {question.content}", id=f"a-{uuid.uuid4()}")]} + + +def _graph(schema, saver): + return StateGraph(schema).add_node("reply", _reply).add_edge(START, "reply").compile(checkpointer=saver) + + +async def _turns(saver, thread_id: str, count: int, *, start: int = 0, at: str | None = None, + pre_delta: int = 0, delta=_SnapshotEverySecondUpdate) -> str: + """Run ``count`` turns from ``at`` (the newest checkpoint without it) and + return the checkpoint the last one ended on.""" + tip = at + for i in range(start, start + count): + config = {"configurable": {"thread_id": thread_id}} + if tip is not None: + config["configurable"]["checkpoint_id"] = tip + schema = _PreDeltaState if i < pre_delta else delta + await _graph(schema, saver).ainvoke({"messages": [HumanMessage(f"q{i}")]}, config) + # Runs are sequential, so the newest checkpoint ends the turn just run. + newest = [c async for c in saver.alist({"configurable": {"thread_id": thread_id}}, limit=1)] + tip = newest[0].config["configurable"]["checkpoint_id"] + assert tip is not None + return tip + + +async def _register(test_db_uri: str, workspace_id: str, thread_id: str, tip_id: str) -> None: + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + await conn.execute( + "INSERT INTO conversation_threads (conversation_thread_id, workspace_id, current_status," + " thread_index, latest_checkpoint_id) VALUES (%s, %s, 'completed', 0, %s)", + (thread_id, workspace_id, tip_id), + ) + + +async def _loads(test_db_uri: str, thread_id: str) -> dict[str, tuple[str, int]]: + """Every root checkpoint's messages as the reader serves them: digest and walk length.""" + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + cur = await conn.execute( + "SELECT checkpoint_id FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = ''" + " ORDER BY checkpoint_id", + (thread_id,), + ) + ids = [r[0] for r in await cur.fetchall()] + loader = Loader(conn) + out = {} + for cid in ids: + messages, walk = await loader.state_messages(thread_id, cid) + out[cid] = (messages_digest(loader.serde, messages), walk.chain) + return out + + +async def _rows(test_db_uri: str, thread_id: str) -> tuple[list, list]: + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + cur = await conn.execute( + "SELECT checkpoint_id, checkpoint::text FROM checkpoints WHERE thread_id = %s ORDER BY 1", + (thread_id,), + ) + checkpoints = await cur.fetchall() + cur = await conn.execute( + "SELECT version, type, md5(blob) FROM checkpoint_blobs" + " WHERE thread_id = %s AND channel = 'messages' ORDER BY 1", + (thread_id,), + ) + return checkpoints, await cur.fetchall() + + +async def _new_thread(test_db_pool, test_db_uri, workspace, *, turns: int, pre_delta: int = 0, + delta=_SnapshotEverySecondUpdate) -> tuple[str, str]: + thread_id = str(uuid.uuid4()) + tip = await _turns(AsyncPostgresSaver(test_db_pool), thread_id, turns, pre_delta=pre_delta, delta=delta) + await _register(test_db_uri, str(workspace["workspace_id"]), thread_id, tip) + return thread_id, tip + + +@pytest.mark.parametrize("pre_delta", [0, 3], ids=["delta", "migrated"]) +async def test_thinning_keeps_every_load_the_same(test_db_pool, test_db_uri, seed_workspace, pre_delta): + thread_id, tip = await _new_thread(test_db_pool, test_db_uri, seed_workspace, turns=8, pre_delta=pre_delta) + before = await _loads(test_db_uri, thread_id) + _, blobs_before = await _rows(test_db_uri, thread_id) + + out = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.APPLY, _EVERYTHING) + + assert out.status is Status.WRITTEN, (out.status, out.reason) + assert out.report.dropped > 0 and out.report.reclaimed > 0 + assert len(out.checks) == len(before) and all(c.match for c in out.checks) + assert out.report.anchors_match and out.committed_match + _, blobs_after = await _rows(test_db_uri, thread_id) + assert len(blobs_after) == len(blobs_before) - out.report.dropped + out.report.seeded + after = await _loads(test_db_uri, thread_id) + assert {c: d for c, (d, _) in after.items()} == {c: d for c, (d, _) in before.items()} + assert after[tip][1] <= before[tip][1] + # Idempotent: a second pass finds nothing it may drop. + again = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.APPLY, _EVERYTHING) + assert again.status is Status.NOTHING, (again.status, again.reason) + + +async def test_a_rehearsal_and_a_dry_run_leave_the_thread_as_it_was(test_db_pool, test_db_uri, seed_workspace): + thread_id, tip = await _new_thread(test_db_pool, test_db_uri, seed_workspace, turns=6, pre_delta=2) + rows = await _rows(test_db_uri, thread_id) + + dry = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.DRY_RUN, _EVERYTHING) + rehearsed = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.REHEARSE, _EVERYTHING) + + assert dry.status is Status.PLANNED and dry.report.dropped > 0 and not dry.checks + assert rehearsed.status is Status.REHEARSED and rehearsed.report.dropped_bytes > 0 + assert rehearsed.checks and all(c.match for c in rehearsed.checks) + assert await _rows(test_db_uri, thread_id) == rows + + +async def test_a_fork_tip_keeps_the_snapshot_it_loads(test_db_pool, test_db_uri, seed_workspace): + saver = AsyncPostgresSaver(test_db_pool) + thread_id = str(uuid.uuid4()) + turn_ends = [] + for i in range(6): + turn_ends.append(await _turns(saver, thread_id, 1, start=i, at=turn_ends[-1] if turn_ends else None)) + fork_tip = await _turns(saver, thread_id, 1, start=100, at=turn_ends[2]) + tip = await _turns(saver, thread_id, 2, start=6, at=turn_ends[-1]) + await _register(test_db_uri, str(seed_workspace["workspace_id"]), thread_id, tip) + before = await _loads(test_db_uri, thread_id) + + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + tree, sizes = await read_tree(conn, thread_id) + flags = {c for c, n in tree.nodes.items() if n.flagged} + seed, _ = tree.walk(fork_tip, flags) + fork_version = ( + tree.nodes[fork_tip].version + if tree.nodes[fork_tip].version in sizes + else tree.nodes[seed].version + ) + assert fork_version in sizes + + out = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.APPLY, _EVERYTHING) + + assert out.status is Status.WRITTEN, (out.status, out.reason) + assert out.report.dropped > 0 and out.report.kept["branch tip"] >= 1 + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + _, sizes_after = await read_tree(conn, thread_id) + assert fork_version in sizes_after + after = await _loads(test_db_uri, thread_id) + assert after[fork_tip] == before[fork_tip] + assert {c: d for c, (d, _) in after.items()} == {c: d for c, (d, _) in before.items()} + + +async def test_a_thread_whose_turns_have_no_stored_lines_is_not_a_candidate( + test_db_pool, test_db_uri, seed_workspace, seed_response, patched_get_db_connection +): + from src.server.database.conversation.queries import create_query + + thread_id, tip = await _new_thread(test_db_pool, test_db_uri, seed_workspace, turns=2) + for turn in range(2): + await create_query(str(uuid.uuid4()), thread_id, turn, f"q{turn}", "initial" if turn == 0 else "follow_up") + await seed_response(str(uuid.uuid4()), thread_id, turn) + + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + candidacy = await thinning.preconditions(conn, thread_id, Loader(conn)) + unchecked = await thinning.preconditions(conn, thread_id, Loader(conn), coverage=False) + + assert candidacy == thinning.NotCandidate("turns without current stored lines", 2) + assert unchecked == thinning.Candidate(tip) # --keep all makes nothing slower + + +async def test_a_thread_with_a_run_in_progress_is_not_a_candidate( + test_db_pool, test_db_uri, seed_workspace, seed_response, patched_get_db_connection +): + thread_id, tip = await _new_thread(test_db_pool, test_db_uri, seed_workspace, turns=2) + await seed_response(str(uuid.uuid4()), thread_id, 2, status="in_progress") + + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + candidacy = await thinning.preconditions(conn, thread_id, Loader(conn), coverage=False) + out = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.APPLY, _EVERYTHING) + + assert candidacy == thinning.NotCandidate("a run is in progress") + assert (out.status, out.reason) == (Status.SKIPPED, "a run is in progress") + + +async def test_an_unverified_saver_release_is_refused( + test_db_pool, test_db_uri, seed_workspace, monkeypatch +): + """On the per-thread path too, which a caller reaches without ``run()``.""" + import importlib.metadata + + thread_id, tip = await _new_thread(test_db_pool, test_db_uri, seed_workspace, turns=8) + rows = await _rows(test_db_uri, thread_id) + monkeypatch.setattr(importlib.metadata, "version", lambda dist: "0.0.1") + + out = await thinning.thin_thread(test_db_uri, thread_id, tip, Mode.APPLY, _EVERYTHING) + + assert "not a verified release" in (unsupported_saver() or "") + assert out.status is Status.ERROR and "not a verified release" in out.reason + assert await _rows(test_db_uri, thread_id) == rows + + +# -------------------------------------------------------------------------- +# --keep all: lists stored before DeltaChannel, at production's cadence +# -------------------------------------------------------------------------- + +_KEEP_ALL = thinning.Options(keep_all=True, verify_all=True) + + +async def _load(test_db_uri: str, thread_id: str, checkpoint_id: str) -> tuple[list, Walk]: + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + return await Loader(conn).state_messages(thread_id, checkpoint_id) + + +async def _flagged(test_db_uri: str, thread_id: str) -> set[str]: + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + cur = await conn.execute( + "SELECT checkpoint_id FROM checkpoints WHERE thread_id = %s" + " AND checkpoint -> 'channel_values' -> 'messages' IS NOT NULL", + (thread_id,), + ) + return {r[0] for r in await cur.fetchall()} + + +async def _listed_for_keep_all(test_db_uri: str, thread_id: str) -> bool: + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + cur = await conn.execute(thinning._UNFLAGGED_LISTS_SQL, (0,)) + return thread_id in {r[0] for r in await cur.fetchall()} + + +def _continues(before: list, after: list) -> bool: + """``after`` is ``before`` unchanged, ids and order included, then one new turn.""" + serde = JsonPlusSerializer() + prefix = after[: len(before)] + return ( + [m.id for m in prefix] == [m.id for m in before] + and messages_digest(serde, prefix) == messages_digest(serde, before) + and [m.type for m in after[len(before) :]] == ["human", "ai"] + ) + + +async def _migrated(test_db_pool, test_db_uri, workspace, *, pre_delta: int, delta: int): + """``pre_delta`` turns before DeltaChannel, then ``delta`` at production's + cadence; returns the thread and the checkpoint each turn ended on.""" + saver = AsyncPostgresSaver(test_db_pool) + thread_id = str(uuid.uuid4()) + ends: list[str] = [] + for i in range(pre_delta + delta): + at = ends[-1] if ends else None + ends.append(await _turns(saver, thread_id, 1, start=i, at=at, pre_delta=pre_delta, + delta=DeltaAgentState)) + await _register(test_db_uri, str(workspace["workspace_id"]), thread_id, ends[-1]) + return thread_id, ends + + +async def _keep_all(test_db_uri: str, thread_id: str, tip: str, mode: Mode = Mode.APPLY): + return await thinning.thin_thread(test_db_uri, thread_id, tip, mode, _KEEP_ALL) + + +async def test_keep_all_flags_the_lists_of_a_thread_that_has_not_crossed( + test_db_pool, test_db_uri, seed_workspace +): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=3, delta=0) + before = await _loads(test_db_uri, thread_id) + _, blobs_before = await _rows(test_db_uri, thread_id) + + out = await _keep_all(test_db_uri, thread_id, ends[-1]) + + assert out.status is Status.WRITTEN, (out.status, out.reason) + report = out.report + assert (report.dropped, report.seeded, report.promoted) == (0, 0, report.blobs) + assert all(c.match for c in out.checks) and out.committed_match + assert set(ends) <= await _flagged(test_db_uri, thread_id) + assert await _loads(test_db_uri, thread_id) == before + assert (await _rows(test_db_uri, thread_id))[1] == blobs_before # no blob written + assert not await _listed_for_keep_all(test_db_uri, thread_id) + + +async def test_after_keep_all_the_first_delta_turn_continues_the_stored_list( + test_db_pool, test_db_uri, seed_workspace +): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=3, delta=0) + await _keep_all(test_db_uri, thread_id, ends[-1]) + stored, _ = await _load(test_db_uri, thread_id, ends[-1]) + + tip = await _turns(AsyncPostgresSaver(test_db_pool), thread_id, 1, start=3, delta=DeltaAgentState) + + after, walk = await _load(test_db_uri, thread_id, tip) + assert _continues(stored, after) + assert walk.seed and walk.writes == 2 # seeded at the list: the input and the reply + + +async def test_after_keep_all_a_branch_forked_at_an_older_list_continues_it( + test_db_pool, test_db_uri, seed_workspace +): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=3, delta=0) + await _keep_all(test_db_uri, thread_id, ends[-1]) + stored, _ = await _load(test_db_uri, thread_id, ends[0]) + + tip = await _turns( + AsyncPostgresSaver(test_db_pool), thread_id, 1, start=100, at=ends[0], delta=DeltaAgentState + ) + + after, walk = await _load(test_db_uri, thread_id, tip) + assert _continues(stored, after) + assert walk.seed and walk.writes == 2 + + +async def test_keep_all_seeds_a_crossed_thread_with_the_value_it_loads_today( + test_db_pool, test_db_uri, seed_workspace +): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=2, delta=1) + before = await _loads(test_db_uri, thread_id) + + out = await _keep_all(test_db_uri, thread_id, ends[-1]) + + assert out.status is Status.WRITTEN, (out.status, out.reason) + assert (out.report.dropped, out.report.seeded) == (0, 1) + assert all(c.match for c in out.checks) and out.committed_match + tip = out.committed + assert tip is not None and tip.before.pages and not tip.before.seed # walked to the root + assert tip.after is not None and tip.after.chain < tip.before.chain + assert {c: d for c, (d, _) in (await _loads(test_db_uri, thread_id)).items()} == { + c: d for c, (d, _) in before.items() + } + + +async def test_keep_all_flags_the_lists_a_crossed_walk_passed_over( + test_db_pool, test_db_uri, seed_workspace +): + """A list a delta walk passes over is flagged too, or a branch off it + would replay the old writes from the root.""" + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=2, delta=1) + + out = await _keep_all(test_db_uri, thread_id, ends[-1]) + + assert out.status is Status.WRITTEN, (out.status, out.reason) + assert ends[0] in await _flagged(test_db_uri, thread_id) + stored, _ = await _load(test_db_uri, thread_id, ends[0]) + tip = await _turns( + AsyncPostgresSaver(test_db_pool), thread_id, 1, start=100, at=ends[0], delta=DeltaAgentState + ) + after, walk = await _load(test_db_uri, thread_id, tip) + assert _continues(stored, after) + assert walk.seed and walk.writes == 2 + + +@pytest.mark.parametrize(("pre_delta", "delta"), [(3, 0), (2, 1)], ids=["not crossed", "crossed"]) +async def test_keep_all_dry_run_and_rehearsal_leave_the_thread_as_it_was( + test_db_pool, test_db_uri, seed_workspace, pre_delta, delta +): + thread_id, ends = await _migrated( + test_db_pool, test_db_uri, seed_workspace, pre_delta=pre_delta, delta=delta + ) + rows = await _rows(test_db_uri, thread_id) + + dry = await _keep_all(test_db_uri, thread_id, ends[-1], Mode.DRY_RUN) + rehearsed = await _keep_all(test_db_uri, thread_id, ends[-1], Mode.REHEARSE) + + assert dry.status is Status.PLANNED and dry.report.promoted and not dry.checks + assert rehearsed.status is Status.REHEARSED and all(c.match for c in rehearsed.checks) + assert await _rows(test_db_uri, thread_id) == rows + assert await _listed_for_keep_all(test_db_uri, thread_id) + + +@pytest.mark.parametrize(("pre_delta", "delta"), [(3, 0), (2, 1)], ids=["not crossed", "crossed"]) +async def test_keep_all_rerun_is_a_no_op(test_db_pool, test_db_uri, seed_workspace, pre_delta, delta): + thread_id, ends = await _migrated( + test_db_pool, test_db_uri, seed_workspace, pre_delta=pre_delta, delta=delta + ) + assert (await _keep_all(test_db_uri, thread_id, ends[-1])).status is Status.WRITTEN + rows = await _rows(test_db_uri, thread_id) + + again = await _keep_all(test_db_uri, thread_id, ends[-1]) + + assert again.status is Status.NOTHING, (again.status, again.reason) + assert await _rows(test_db_uri, thread_id) == rows + assert not await _listed_for_keep_all(test_db_uri, thread_id) + + +async def test_keep_all_skips_a_checkpoint_that_changes_before_the_write( + test_db_pool, test_db_uri, seed_workspace, monkeypatch +): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=1, delta=1) + rows = await _rows(test_db_uri, thread_id) + write = thinning._write + + async def a_turn_rewrites_a_list_first(conn, *args): + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as other: + await other.execute( + "UPDATE checkpoints SET metadata = metadata WHERE thread_id = %s AND checkpoint_id = %s", + (thread_id, ends[0]), + ) + return await write(conn, *args) + + monkeypatch.setattr(thinning, "_write", a_turn_rewrites_a_list_first) + out = await _keep_all(test_db_uri, thread_id, ends[-1]) + + assert (out.status, out.reason) == (Status.SKIPPED, "a checkpoint changed concurrently") + assert await _rows(test_db_uri, thread_id) == rows + + +async def test_keep_all_never_seeds_over_a_stored_list(test_db_pool, test_db_uri, seed_workspace): + """A list stored where a crossing walk would be seeded is flagged as one.""" + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=1, delta=1) + async with await psycopg.AsyncConnection.connect(test_db_uri, autocommit=True) as conn: + tree, sizes = await read_tree(conn, thread_id) + first_delta = next( + c for c in tree.order if tree.nodes[c].version and tree.nodes[c].version not in sizes + ) + await conn.execute( + "INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)" + " VALUES (%s, '', 'messages', %s, 'msgpack', '\\x90'::bytea)", + (thread_id, tree.nodes[first_delta].version), + ) + before = await _loads(test_db_uri, thread_id) + _, blobs = await _rows(test_db_uri, thread_id) + + out = await _keep_all(test_db_uri, thread_id, ends[-1]) + + assert out.status is Status.WRITTEN, (out.status, out.reason) + assert first_delta in await _flagged(test_db_uri, thread_id) + _, blobs_after = await _rows(test_db_uri, thread_id) + assert set(blobs) <= set(blobs_after) and len(blobs_after) == len(blobs) + out.report.seeded + assert {c: d for c, (d, _) in (await _loads(test_db_uri, thread_id)).items()} == { + c: d for c, (d, _) in before.items() + } + + +async def test_keep_all_leaves_a_thread_written_after_deltachannel( + test_db_pool, test_db_uri, seed_workspace +): + thread_id, tip = await _new_thread(test_db_pool, test_db_uri, seed_workspace, turns=4) + + out = await _keep_all(test_db_uri, thread_id, tip) + + assert out.status is Status.NOTHING, (out.status, out.reason) + assert not await _listed_for_keep_all(test_db_uri, thread_id) + + +# -------------------------------------------------------------------------- +# The harness: nothing commits unverified +# -------------------------------------------------------------------------- + + +async def test_a_thread_transaction_may_sit_idle_no_longer_than_a_statement(test_db_uri): + seen = {} + + async def plan(conn, loader): + cur = await conn.execute( + "SELECT current_setting('idle_in_transaction_session_timeout')," + " current_setting('statement_timeout')" + ) + seen["idle"], seen["statement"] = await cur.fetchone() + return Rewrite(None) + + out = await rewrite_thread(test_db_uri, "t", Mode.DRY_RUN, Timeouts(statement_s=7), plan) + + assert out.status is Status.NOTHING + assert seen == {"idle": "7s", "statement": "7s"} + + +async def _rewrite_tip(test_db_uri, thread_id, tip, statement, verify=None): + async def plan(conn, loader): + async def write(): + await conn.execute(statement, (thread_id,)) + + return Rewrite(None, [tip], tip, write, verify) + + return await rewrite_thread(test_db_uri, thread_id, Mode.APPLY, Timeouts(), plan) + + +async def test_a_write_that_changes_a_load_is_rolled_back(test_db_pool, test_db_uri, seed_workspace): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=2, delta=0) + rows = await _rows(test_db_uri, thread_id) + + out = await _rewrite_tip( + test_db_uri, thread_id, ends[-1], + "UPDATE checkpoint_blobs SET type = 'msgpack', blob = '\\x90'::bytea" + " WHERE thread_id = %s AND channel = 'messages'", # every list now loads empty + ) + + assert out.status is Status.MISMATCH and out.failed + assert [c.match for c in out.checks] == [False] + assert await _rows(test_db_uri, thread_id) == rows + + +async def test_a_write_its_planner_rejects_is_rolled_back(test_db_pool, test_db_uri, seed_workspace): + thread_id, ends = await _migrated(test_db_pool, test_db_uri, seed_workspace, pre_delta=2, delta=0) + rows = await _rows(test_db_uri, thread_id) + + async def reject(checks): + assert all(c.match for c in checks.values()) + return "rejected" + + out = await _rewrite_tip( + test_db_uri, thread_id, ends[-1], + "UPDATE checkpoints SET metadata = metadata || '{\"x\": 1}' WHERE thread_id = %s", + reject, + ) + + assert (out.status, out.reason) == (Status.MISMATCH, "rolled back: rejected") + assert await _rows(test_db_uri, thread_id) == rows diff --git a/tests/integration/test_turn_slices_db.py b/tests/integration/test_turn_slices_db.py new file mode 100644 index 000000000..6b125f983 --- /dev/null +++ b/tests/integration/test_turn_slices_db.py @@ -0,0 +1,137 @@ +"""Batched slice writes against real Postgres: one statement per batch, a +row whose parent is gone dropped alone, and lines kept off a successor +slice.""" + +from __future__ import annotations + +import uuid + +import pytest + +from src.server.database.conversation.turn_slices import ( + RunRow, + StoredSlice, + TurnRow, + TurnUpdate, +) + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + + +def _turn_row(response_id, thread_id, turn_index, *, tail, lines=None, claims=None): + return TurnRow( + response_id, + thread_id, + StoredSlice( + input_checkpoint_id=f"in-{turn_index}", + tail_checkpoint_id=tail, + slice_key="k", + codec="msgpack", + data=b"\x80", + lines_key="lk" if lines is not None else None, + lines=lines, + claims=claims, + ), + ) + + +async def _thread(seed_workspace, seed_response, turns): + from src.server.database.conversation import create_thread + + thread_id = str(uuid.uuid4()) + await create_thread( + conversation_thread_id=thread_id, + workspace_id=str(seed_workspace["workspace_id"]), + current_status="completed", + msg_type="ptc", + ) + response_ids = [] + for turn_index in range(turns): + response_id = str(uuid.uuid4()) + await seed_response(response_id, thread_id, turn_index) + response_ids.append(response_id) + return thread_id, response_ids + + +async def test_a_batch_stores_every_row_but_an_orphan( + seed_workspace, seed_response, patched_get_db_connection +): + from src.server.database.conversation import turn_slices as slices_db + + thread_id, (r0, r1) = await _thread(seed_workspace, seed_response, 2) + gone = str(uuid.uuid4()) + + await slices_db.upsert_turn_slices( + [ + _turn_row(r0, thread_id, 0, tail="t0", lines="a\t\t{}", claims=["c0"]), + _turn_row(gone, thread_id, 1, tail="tx"), + _turn_row(r1, thread_id, 1, tail="t1"), + ] + ) + + rows = await slices_db.get_turn_slices([r0, r1, gone]) + assert sorted(rows) == sorted([r0, r1]) + assert rows[r0].lines == "a\t\t{}" + assert rows[r0].claims == ["c0"] + assert rows[r1].data == b"\x80" + assert rows[r1].lines is None + + +async def test_an_update_lands_only_on_the_slice_it_was_projected_from( + seed_workspace, seed_response, patched_get_db_connection +): + from src.server.database.conversation import turn_slices as slices_db + + thread_id, (r0, r1) = await _thread(seed_workspace, seed_response, 2) + await slices_db.upsert_turn_slices( + [ + _turn_row(r0, thread_id, 0, tail="t0", claims=["c0"]), + _turn_row(r1, thread_id, 1, tail="t1-new"), + ] + ) + + await slices_db.update_turns( + [ + TurnUpdate(r0, "t0", "lk", "x\t\t{}", None), + TurnUpdate(r1, "t1-old", "lk", "stale\t\t{}", ["c9"]), + ] + ) + + rows = await slices_db.get_turn_lines([r0, r1]) + assert rows[r0].lines == "x\t\t{}" + assert rows[r0].claims == ["c0"] + assert rows[r0].data is None + assert rows[r1].lines is None + assert rows[r1].claims is None + + # A slice cut again replaces its claims rather than keeping the old ones. + await slices_db.upsert_turn_slices([_turn_row(r0, thread_id, 0, tail="t0b")]) + assert (await slices_db.get_turn_lines([r0]))[r0].claims is None + + +async def test_run_slices_store_in_one_batch_and_a_rewrite_moves_the_tail( + seed_workspace, seed_response, patched_get_db_connection +): + from src.server.database.conversation import turn_slices as slices_db + + thread_id, _ = await _thread(seed_workspace, seed_response, 0) + + def run(task_id, ordinal, tail): + return RunRow( + thread_id, + task_id, + None, + StoredSlice(f"{task_id}:in-{ordinal}", tail, "k", "msgpack", b"\x80"), + ) + + await slices_db.upsert_run_slices( + [run("b", 0, "b0"), run("a", 1, "a1"), run("a", 0, "a0")] + ) + await slices_db.upsert_run_slices([run("a", 1, "a1-moved")]) + + rows = await slices_db.get_run_slices(thread_id, ["a", "b"]) + assert sorted((r.slice.input_checkpoint_id, r.slice.tail_checkpoint_id) for r in rows) == [ + ("a:in-0", "a0"), + ("a:in-1", "a1-moved"), + ("b:in-0", "b0"), + ] diff --git a/tests/integration/test_window_identity_db.py b/tests/integration/test_window_identity_db.py new file mode 100644 index 000000000..b819e50e8 --- /dev/null +++ b/tests/integration/test_window_identity_db.py @@ -0,0 +1,103 @@ +"""The window on the Postgres saver: the scenario of ``window_harness`` run with +the trim and without it, compared as the unit suite compares it on the +in-memory saver. Postgres is where the trim's ``Overwrite`` has to start a +new snapshot for loading the tip to decode only the window, and where the +history reader walks the thread's turns from table metadata. +""" + +from __future__ import annotations + +import pytest +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + +import src.server.utils.checkpointer # noqa: F401 (installs the walk guard) +from tests.unit.middleware.compaction import window_harness as h +from tests.unit.server.services.history.replay_builders import SliceStore, _cache_probe +from tests.unit.server.services.history.test_window import Thread, _page +from src.server.services.history import window as window_module + +pytestmark = [pytest.mark.integration, pytest.mark.asyncio] + + +async def _run(saver, monkeypatch, *, window: bool): + run = await h.run_scenario(saver, window=window) + thread = Thread(run, saver) + decoded = await h.decoded_messages(saver, run.graph) + store = SliceStore() + page = await _page(monkeypatch, thread, store) + head = await thread.trimmed() if window else thread.head(1) + covered = await window_module.runs_held(h.THREAD, head) + scenario_calls = len(run.model.requests) + await h.run_edit(run, saver, 3, window=window) + await saver.adelete_thread(h.THREAD) + return run, decoded, page, covered, scenario_calls + + +async def test_the_window_on_postgres(test_db_pool, monkeypatch) -> None: + saver = AsyncPostgresSaver(test_db_pool) + _cache_probe(monkeypatch) + off, off_decoded, off_page, _, calls = await _run(saver, monkeypatch, window=False) + on, on_decoded, on_page, covered, _ = await _run(saver, monkeypatch, window=True) + + assert h.digests(on.model.requests) == h.digests(off.model.requests) + assert h.digests(on.summary_model.requests) == h.digests(off.summary_model.requests) + assert len(on.model.requests) > calls + assert on_page == off_page + assert covered + assert on.asked and all(whole for _, whole, _ in on.asked) + assert h.transcript(on.store) == h.transcript(off.store) + + scenario = len(h.SCENARIO) - 1 + held_on, held_off = on.held[scenario], off.held[scenario] + assert held_on[1] > 0 and held_off[1] == 0 + assert on_decoded <= held_on[0] + 2 + assert off_decoded >= held_off[0] + + +async def test_the_window_reads_the_threads_slices( + seed_workspace, seed_user, seed_response, patched_get_db_connection +) -> None: + """The two reads the window makes, against the real table: whether a + thread holds slices under one slice key, and the rows of those that start + at given checkpoints.""" + import uuid + + from src.server.database.conversation import ensure_thread_exists + from src.server.database.conversation import turn_slices as slices_db + + threads = [str(uuid.uuid4()), str(uuid.uuid4())] + for thread_id in threads: + await ensure_thread_exists( + workspace_id=str(seed_workspace["workspace_id"]), + conversation_thread_id=thread_id, + user_id=seed_user["user_id"], + initial_query="q", + ) + rows = [ + (threads[0], 0, "key"), + (threads[0], 1, "key"), + (threads[0], 2, "older-key"), + (threads[1], 0, "key"), + ] + stored = [] + for thread_id, turn, key in rows: + response_id = str(uuid.uuid4()) + await seed_response(response_id, thread_id, turn) + stored.append( + slices_db.TurnRow( + response_id, + thread_id, + slices_db.StoredSlice( + f"in-{turn}", f"tail-{turn}", key, "msgpack", f"slice-{turn}".encode() + ), + ) + ) + await slices_db.upsert_turn_slices(stored) + + assert await slices_db.has_turn_slices(threads[0], "key") + assert await slices_db.has_turn_slices(threads[0], "older-key") + assert not await slices_db.has_turn_slices(threads[1], "older-key") + assert await slices_db.get_slices_at(threads[0], "key", ["in-1", "in-2"]) == [ + slices_db.StoredSlice("in-1", "tail-1", "key", "msgpack", b"slice-1") + ] + assert await slices_db.get_slices_at(threads[0], "key", []) == [] diff --git a/tests/unit/middleware/compaction/test_offload_view.py b/tests/unit/middleware/compaction/test_offload_view.py index fd8ecbbbd..1e0b92297 100644 --- a/tests/unit/middleware/compaction/test_offload_view.py +++ b/tests/unit/middleware/compaction/test_offload_view.py @@ -24,7 +24,7 @@ record_offloads, ) from ptc_agent.agent.middleware.compaction.types import OffloadSettings -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.pointer import TranscriptTurns from ptc_agent.config.agent import CompactionConfig @@ -55,7 +55,7 @@ def __init__(self, saves: bool = True): self.saves = saves self.saved: list[int] = [] - async def save_transcript(self, target, messages): + async def save_transcript(self, target, messages, *, window): self.saved.append(len(messages)) return self.saves @@ -286,7 +286,7 @@ def test_legacy_merged_id_sets_only_hit_their_own_tool(): # Older /offload merged arg and read ids into both sets. ids = {"w0", "r0"} msgs = _history(1) - turns = TranscriptTurns.of(TranscriptTarget(THREAD), msgs) + turns = TranscriptTurns.of(TranscriptTarget(THREAD), msgs, window=Window()) out = {m.id: m for m in apply_recorded_offloads(msgs, ids, ids, 200, turns)} assert out["tw0"].content == "wrote" diff --git a/tests/unit/middleware/compaction/test_scratchpad_notes.py b/tests/unit/middleware/compaction/test_scratchpad_notes.py index ee4e02af6..5561da253 100644 --- a/tests/unit/middleware/compaction/test_scratchpad_notes.py +++ b/tests/unit/middleware/compaction/test_scratchpad_notes.py @@ -43,7 +43,7 @@ TURN_ROW_KIND, ) from ptc_agent.agent.prompts import init_loader -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.core.paths import WorkspaceLayout THREAD = "a1b2c3d4-0000-4000-8000-000000000000" @@ -192,6 +192,7 @@ async def _compact(backend): transcript=None, fallback=None, notes_dir=NOTES, + window=Window(), ) diff --git a/tests/unit/middleware/compaction/test_window.py b/tests/unit/middleware/compaction/test_window.py new file mode 100644 index 000000000..035e007da --- /dev/null +++ b/tests/unit/middleware/compaction/test_window.py @@ -0,0 +1,375 @@ +"""The window's pieces: where it may cut, what a trim writes, the carries that +keep counts over trimmed runs, the transcript's numbering base, and when the +middleware stands down. ``test_window_identity`` proves the whole.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import patch + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langgraph.types import Overwrite + +from ptc_agent.agent.middleware.compaction import middleware as middleware_module +from ptc_agent.agent.middleware.compaction.compact import Summarizer +from ptc_agent.agent.middleware.compaction.middleware import CompactionMiddleware +from ptc_agent.agent.middleware.compaction.notes import notes_window_carry +from ptc_agent.agent.middleware.compaction.window import trim, window_cut +from ptc_agent.agent.middleware.runtime_context.durable import ( + DurableUpdate, + build_update_message, +) +from ptc_agent.agent.middleware.subagent_switch import ( + SUBAGENTS_ROW_KIND, + SUBAGENTS_SCHEMA_VERSION, + last_announced, + subagents_window_carry, +) +from ptc_agent.agent.transcript import ( + EarlierTurnsMissing, + Window, + build_directory, + load_manifest, +) +from ptc_agent.agent.transcript import identity +from tests.unit.middleware.compaction.window_harness import NOTES_DIR, THREAD + +INDEX_LIMIT = identity._INDEX_LIMIT + + +def _run(n: int, *, calls: int = 1, write_to: str | None = None) -> list: + out = [HumanMessage(f"Question {n}", id=f"h-{n}")] + for k in range(calls): + call_id = f"c-{n}-{k}" + path = write_to if write_to and k == 0 else f"/home/workspace/work/{n}-{k}.py" + out.append( + AIMessage( + "", + id=f"ai-{n}-{k}", + tool_calls=[{"name": "Write", "id": call_id, "args": {"file_path": path}}], + ) + ) + out.append(ToolMessage("ok", tool_call_id=call_id, id=f"t-{n}-{k}")) + out.append(AIMessage(f"Answer {n}", id=f"ai-{n}-end")) + return out + + +def _thread(runs: int, first: int = 1) -> list: + return [m for n in range(first, first + runs) for m in _run(n)] + + +def _event(messages: list, anchor_id: str | None) -> dict: + index = next((i for i, m in enumerate(messages) if m.id == anchor_id), 0) + event = {"cutoff_index": index, "summary_message": HumanMessage("summary", id="s-1")} + if anchor_id is not None: + event["anchor_message_id"] = anchor_id + return event + + +def _index(messages: list, message_id: str) -> int: + return next(i for i, m in enumerate(messages) if m.id == message_id) + + +def _trim(state: dict, cut: int, carries=()) -> dict: + """``trim`` with the head the middleware hands it.""" + return trim(state, cut, carries, head=Window.of(state).extend(state["messages"][:cut])) + + +# ------------------------------------------------------------ window_cut + + +def test_the_cut_is_the_start_of_the_run_holding_the_boundary() -> None: + messages = _thread(4) + assert window_cut(messages, _event(messages, "t-3-0")) == _index(messages, "h-3") + assert window_cut(messages, _event(messages, "h-3")) == _index(messages, "h-3") + + +def test_nothing_is_cut_inside_the_first_run() -> None: + messages = _thread(3) + assert window_cut(messages, _event(messages, "ai-1-end")) == 0 + + +def test_an_unanchored_or_unresolvable_boundary_cuts_nothing() -> None: + messages = _thread(3) + assert window_cut(messages, None) == 0 + assert window_cut(messages, {"cutoff_index": 9}) == 0 + assert window_cut(messages, {"cutoff_index": 9, "anchor_message_id": "gone"}) == 0 + + +def test_a_stamped_injection_opens_no_run() -> None: + messages = _thread(2) + notice = HumanMessage( + "Background task finished.", id="orch", name="orchestrator", + additional_kwargs={"lc_source": "orchestrator"}, + ) + messages = [*messages, notice, AIMessage("more", id="ai-more")] + assert window_cut(messages, _event(messages, "ai-more")) == _index(messages, "h-2") + + +def test_messages_before_the_first_question_belong_to_the_first_run() -> None: + messages = [AIMessage("hello", id="greeting"), *_thread(3)] + assert window_cut(messages, _event(messages, "t-2-0")) == _index(messages, "h-2") + assert Window().extend(messages[: _index(messages, "h-2")]).runs == 1 + + +# ------------------------------------------------------------ trim + + +def test_a_trim_keeps_the_window_and_counts_what_went() -> None: + messages = _thread(4) + cut = _index(messages, "h-3") + update = _trim({"messages": messages}, cut) + assert isinstance(update["messages"], Overwrite) + assert [m.id for m in update["messages"].value] == [m.id for m in messages[cut:]] + window = Window.of(update) + assert window.runs == 2 + assert window.earlier == {1: "Question 1", 2: "Question 2"} + assert window.holds(messages[:cut]) + + +def test_trims_count_on_from_the_last() -> None: + first = _trim({"messages": _thread(4)}, _index(_thread(4), "h-3")) + state = {"messages": first["messages"].value, **{k: v for k, v in first.items() if k != "messages"}} + second = Window.of(_trim(state, _index(state["messages"], "h-4"))) + assert second.runs == 3 + assert second.earlier == {1: "Question 1", 2: "Question 2", 3: "Question 3"} + # The digest runs on over both trims, as one trim of both would leave it. + whole = _thread(4) + assert second == Window().extend(whole[: _index(whole, "h-4")]) + + +# ------------------------------------------------------------ the head + + +def test_the_head_is_its_messages_in_order_whatever_their_ids() -> None: + """A message saved before ids were stable loads with another id, or + none, from one checkpoint to the next; its words and calls do not.""" + messages = _thread(2) + head = Window().extend(messages) + assert (head.runs, head.messages) == (2, len(messages)) + reloaded = [m.model_copy(update={"id": None}) for m in messages] + assert head.holds(reloaded) + assert not head.holds(messages[:-1]) + assert not head.holds([*messages[:-2], messages[-1], messages[-2]]) + edited = [*messages[:-1], AIMessage("Answer 2, edited", id="ai-2-end")] + assert not head.holds(edited) + + +def test_the_earlier_index_keeps_only_what_a_summary_lists() -> None: + messages = _thread(INDEX_LIMIT + 5) + update = _trim({"messages": messages}, _index(messages, f"h-{INDEX_LIMIT + 4}")) + window = Window.of(update) + assert window.runs == INDEX_LIMIT + 3 + assert sorted(window.earlier) == list(range(4, INDEX_LIMIT + 4)) + + +def test_kept_messages_get_ids_the_overwrite_would_not_stamp() -> None: + messages = [*_thread(1), HumanMessage("Question 2"), AIMessage("Answer 2")] + update = _trim({"messages": messages}, len(_run(1))) + assert all(m.id for m in update["messages"].value) + + +def test_the_recorded_offloads_shrink_to_the_window() -> None: + messages = _thread(3) + state = { + "messages": messages, + "_offloaded_tool_call_ids": {"c-1-0", "c-3-0"}, + "_offloaded_read_result_ids": {"c-3-0"}, + } + update = _trim(state, _index(messages, "h-2")) + assert update["_offloaded_tool_call_ids"] == {"c-3-0"} + assert "_offloaded_read_result_ids" not in update + + +def test_each_carry_reads_the_trimmed_runs() -> None: + seen = [] + + def carry(trimmed, state): + seen.append([m.id for m in trimmed]) + return {"_carried": len(trimmed)} + + messages = _thread(3) + update = _trim({"messages": messages}, _index(messages, "h-2"), [carry]) + assert seen == [[m.id for m in _run(1)]] + assert update["_carried"] == len(_run(1)) + + +# ------------------------------------------------------------ the carries + + +def test_the_notes_count_runs_back_through_trimmed_runs() -> None: + carry = notes_window_carry(THREAD) + wrote = [*_run(1, calls=3, write_to=f"{NOTES_DIR}/plan.md")] + # A write in the trimmed runs: the count restarts after it. + assert carry(wrote, {}) == {"_notes_calls_trimmed": 2} + # No write: the count adds to what earlier trims left. + assert carry(_run(2, calls=2), {"_notes_calls_trimmed": 2}) == {"_notes_calls_trimmed": 4} + assert carry([HumanMessage("hi", id="h")], {"_notes_calls_trimmed": 2}) == {} + + +def _notice(allowed: bool, message_id: str) -> HumanMessage: + row = DurableUpdate( + kind=SUBAGENTS_ROW_KIND, + schema_version=SUBAGENTS_SCHEMA_VERSION, + text="", + provenance={"source": "user", "allowed": allowed}, + ) + message = build_update_message(row) + message.id = message_id + return message + + +def test_a_trimmed_notice_is_carried() -> None: + trimmed = [*_run(1), _notice(False, "n-1"), *_run(2)] + assert subagents_window_carry(trimmed, {}) == {"_subagents_trimmed": True} + assert subagents_window_carry(trimmed, {"_subagents_trimmed": True}) == {} + assert subagents_window_carry(_run(1), {}) == {} + + +def test_a_trimmed_notice_is_restated() -> None: + """The model may still believe what a summarized notice said, so a + switch with no notice in view is restated whichever side it is.""" + assert last_announced({"messages": _run(3)}) is True + assert last_announced({"messages": _run(3), "_subagents_trimmed": True}) is None + # A notice summarized but not yet trimmed. + messages = [*_run(1), _notice(False, "n-1"), *_run(2)] + assert last_announced({"messages": messages}) is False + summarized = {"messages": messages, "_summarization_event": _event(messages, "h-2")} + assert last_announced(summarized) is None + + +# ------------------------------------------------------------ the transcript + + +def test_a_trimmed_render_carries_the_stored_files() -> None: + whole = _thread(5) + stored = build_directory(whole[: _index(whole, "h-5")], window=Window()) + cut = _index(whole, "h-3") + windowed = build_directory( + whole[cut:], + previous=load_manifest(stored.manifest), + window=Window().extend(whole[:cut]), + ) + assert windowed.manifest == build_directory(whole, window=Window()).manifest + assert sorted(windowed.rendered) == ["turn-0005.jsonl"] + + +def test_a_trimmed_render_without_the_stored_files_asks_for_them() -> None: + whole = _thread(4) + cut = _index(whole, "h-3") + window = Window().extend(whole[:cut]) + with pytest.raises(EarlierTurnsMissing): + build_directory(whole[cut:], window=window) + short = build_directory(whole[: _index(whole, "h-2")], window=Window()) + with pytest.raises(EarlierTurnsMissing): + build_directory(whole[cut:], previous=load_manifest(short.manifest), window=window) + + +def test_a_copy_from_another_branch_is_not_carried() -> None: + """A live save mid-turn is labelled with the commit pointer, which after + an edit is the fork point both branches share, so a stored copy can be + another branch's under the same turn numbers.""" + whole = _thread(4) + cut = _index(whole, "h-3") + other = [*_run(1), *_run(2)[:-1], AIMessage("Answer 2, regenerated", id="ai-2-end"), *_run(3)] + stored = load_manifest(build_directory(other, window=Window()).manifest) + with pytest.raises(EarlierTurnsMissing): + build_directory(whole[cut:], previous=stored, window=Window().extend(whole[:cut])) + # Its own branch's copy is carried. + stored = load_manifest(build_directory(whole[: _index(whole, "h-4")], window=Window()).manifest) + windowed = build_directory( + whole[cut:], previous=stored, window=Window().extend(whole[:cut]) + ) + assert windowed.manifest == build_directory(whole, window=Window()).manifest + + +def test_a_copy_from_before_entries_named_their_messages_is_not_carried() -> None: + whole = _thread(4) + cut = _index(whole, "h-3") + stored = load_manifest(build_directory(whole[: _index(whole, "h-4")], window=Window()).manifest) + for entry in stored["segments"]: + del entry["through"] + with pytest.raises(EarlierTurnsMissing): + build_directory(whole[cut:], previous=stored, window=Window().extend(whole[:cut])) + # A full render reuses its files and names them. + again = build_directory(whole, previous=stored, window=Window()) + assert again.manifest == build_directory(whole, window=Window()).manifest + assert sorted(again.rendered) == ["turn-0004.jsonl"] + + +# ------------------------------------------------------------ the middleware + + +def _middleware(coverage=None) -> CompactionMiddleware: + middleware = CompactionMiddleware( + Summarizer(None, limit=10**9, counter=len), token_threshold=10**9, keep_messages=4 + ) + return middleware.with_window(coverage) if coverage else middleware + + +def _summarized(runs: int = 4) -> dict: + messages = _thread(runs) + return {"messages": messages, "_summarization_event": _event(messages, f"t-{runs - 1}-0")} + + +def test_covered_runs_are_trimmed() -> None: + asked = [] + + async def coverage(head): + asked.append(head) + return True + + state = _summarized() + update = asyncio.run(_middleware(coverage)._trim(state)) + head = Window().extend(state["messages"][: _index(state["messages"], "h-3")]) + assert asked == [head] and head.runs == 2 + assert update is not None and Window.of(update) == head + + +def test_a_refusal_is_kept_until_a_summary_moves_the_cut() -> None: + """Kept in the checkpoint, so every worker reads it; a thread whose + slices are backfilled since trims at its next summary.""" + answers = [False, True] + asked = [] + + async def coverage(head): + asked.append(head.runs) + return answers.pop(0) + + middleware = _middleware(coverage) + state = _summarized() + refused = asyncio.run(middleware._trim(state)) + cut = _index(state["messages"], "h-3") + assert refused == {"_window_refused": {"anchor": "t-3-0", "cut": cut}} + state = {**state, **refused} + assert asyncio.run(middleware._trim(state)) is None + assert asked == [2] + + messages = state["messages"] + state = {**state, "_summarization_event": _event(messages, "t-4-0")} + update = asyncio.run(middleware._trim(state)) + assert asked == [2, 3] + assert update is not None and Window.of(update).runs == 3 + + +@pytest.mark.parametrize("answer", [False, RuntimeError("db down"), "slow"]) +def test_nothing_is_trimmed_the_server_does_not_confirm(answer) -> None: + async def coverage(runs): + if answer == "slow": + await asyncio.sleep(1) + if isinstance(answer, Exception): + raise answer + return answer + + with patch.object(middleware_module, "_COVERAGE_TIMEOUT_S", 0.01): + update = asyncio.run(_middleware(coverage)._trim(_summarized())) + # Only an answer is kept as a refusal; a failed or slow check asks again. + if answer is False: + assert set(update) == {"_window_refused"} + else: + assert update is None + + +def test_without_the_server_nothing_is_trimmed() -> None: + assert asyncio.run(_middleware()._trim(_summarized())) is None diff --git a/tests/unit/middleware/compaction/test_window_identity.py b/tests/unit/middleware/compaction/test_window_identity.py new file mode 100644 index 000000000..efd655ff3 --- /dev/null +++ b/tests/unit/middleware/compaction/test_window_identity.py @@ -0,0 +1,275 @@ +"""The window trims the main agent's checkpoint without changing anything the +model, the transcript or an edit reads. + +One scenario (``window_harness.SCENARIO``) runs twice through the real graph +and middleware, once with the trim and once without, and the two are compared +request by request: every model and summary request must be byte-identical +once runtime-minted ids are normalized. The scenario then edits a turn the +window had trimmed and runs on from there, compared the same way. +""" + +from __future__ import annotations + +import asyncio +import re +from dataclasses import dataclass + +import pytest +from langgraph.checkpoint.memory import InMemorySaver + +from ptc_agent.agent.middleware.compaction import middleware as middleware_module +from ptc_agent.agent.state import DeltaAgentState +from ptc_agent.agent.transcript import Window, build_directory +from tests.unit.middleware.compaction import window_harness as h + + +@dataclass +class Phase: + """The two runs' counts at the end of a phase.""" + + calls: int + summaries: int + exports: int + transcripts: tuple[tuple[list[str], dict[str, str]], ...] + decoded: tuple[int, int] + + +@dataclass +class Compared: + off: h.ThreadRun + on: h.ThreadRun + scenario: Phase + edited: Phase + + +def _phase(off: h.ThreadRun, on: h.ThreadRun, decoded: tuple[int, int]) -> Phase: + assert len(off.model.requests) == len(on.model.requests) + return Phase( + calls=len(on.model.requests), + summaries=len(on.summary_model.requests), + exports=len(on.exports), + transcripts=(h.transcript(off.store), h.transcript(on.store)), + decoded=decoded, + ) + + +@pytest.fixture(scope="module") +def compared() -> Compared: + async def run() -> Compared: + saver_off, saver_on = InMemorySaver(), InMemorySaver() + off = await h.run_scenario(saver_off, window=False) + on = await h.run_scenario(saver_on, window=True) + + async def decoded() -> tuple[int, int]: + return ( + await h.decoded_messages(saver_off, off.graph), + await h.decoded_messages(saver_on, on.graph), + ) + + scenario = _phase(off, on, await decoded()) + await h.run_edit(off, saver_off, 3, window=False) + await h.run_edit(on, saver_on, 3, window=True) + return Compared(off, on, scenario, _phase(off, on, await decoded())) + + return asyncio.run(run()) + + +def _texts(compared: Compared) -> list[str]: + return h.digests(compared.on.model.requests[: compared.scenario.calls]) + + +# ------------------------------------------------------------ the requests + + +def test_the_model_is_sent_the_same_requests(compared: Compared) -> None: + off, on = compared.off, compared.on + assert compared.scenario.calls > 40 + assert h.digests(on.model.requests) == h.digests(off.model.requests) + assert h.digests(on.summary_model.requests) == h.digests(off.summary_model.requests) + + +def test_the_scenario_reaches_every_reading_the_trim_could_change(compared: Compared) -> None: + """Each count a trim could move has to show in the requests compared, or + their equality proves nothing about it.""" + texts = "\n".join(_texts(compared)) + on = compared.on + # Several trims, each moving the base on, and each dropping exactly the + # start of the whole thread, as the server checks against its slices. + assert len(on.asked) >= 3 + assert all(whole for _, whole, _ in on.asked) + assert len({base for _, base in on.held[: len(h.SCENARIO)]} - {0}) >= 3 + # Tier 1 cuts and summaries, whose turn numbers and index lines run past + # the trimmed runs (the last test breaks each). + assert texts.count("argument cut here") > 10 + assert len(re.findall(r"Summary \d+:", texts)) > 5 + assert compared.scenario.summaries >= 4 + # The notes check-in, which counts back to the last notes write, made in + # turn 2 and trimmed before the check-in ran. + assert re.search(r"(\d+) tool calls have run since", texts) + assert on.held[1][1] == 0 and on.held[6][1] > 1 + # The subagent notices: off in turns 3-4, back on in 5, after trims. + assert "turned subagents off" in texts + assert "turned subagents back on" in texts + + +def test_the_checkpoint_keeps_only_the_window(compared: Compared) -> None: + last = len(h.SCENARIO) - 1 + off_held, on_held = compared.off.held[last], compared.on.held[last] + assert off_held[1] == 0 + assert on_held[1] > 0 + assert on_held[0] * 3 < off_held[0] + # Loading the tip decodes the window, not the thread. + decoded_off, decoded_on = compared.scenario.decoded + assert decoded_on <= on_held[0] + 2 + assert decoded_off >= off_held[0] + + +# ------------------------------------------------------------ the re-entry + + +def test_a_reentry_trims_without_touching_its_turn(compared: Compared) -> None: + on = compared.on + assert on.reentry is not None and on.reentry_base is not None + before, after = on.reentry + assert after[: len(before)] == before + assert len(after) > len(before) + base_before, base_after = on.reentry_base + assert base_after > base_before + # The turn's own messages were never trimmed. + assert after[0] == "h-8" + + +# ------------------------------------------------------------ a refusal + + +def test_a_refusal_is_asked_once_per_summary_and_sends_the_same(compared: Compared) -> None: + """A thread whose slices do not hold its head keeps the server's refusal + in its checkpoint: it asks again only once a new summary moves the + anchor, and the write changes nothing the model is sent.""" + refused = asyncio.run(h.run_scenario(InMemorySaver(), window=True, covered=False)) + off = compared.off + assert h.digests(refused.model.requests) == h.digests( + off.model.requests[: compared.scenario.calls] + ) + assert h.digests(refused.summary_model.requests) == h.digests( + off.summary_model.requests[: compared.scenario.summaries] + ) + assert {base for _, base in refused.held} == {0} + summaries = [count for _, _, count in refused.asked] + assert len(summaries) >= 3 + assert len(set(summaries)) == len(summaries) + + +# ------------------------------------------------------------ a notification turn + + +def _notified() -> tuple[h.ThreadRun, h.ThreadRun]: + async def run() -> tuple[h.ThreadRun, h.ThreadRun]: + off = await h.run_scenario(InMemorySaver(), window=False, turns=h.NOTIFIED) + on = await h.run_scenario(InMemorySaver(), window=True, turns=h.NOTIFIED) + return off, on + + return asyncio.run(run()) + + +def test_a_build_without_the_switch_trims_the_same() -> None: + """A notification turn's build has no subagent switch, and its trims + still carry the notices they drop, into a field ``MainAgentState`` + declares: the next turn with the switch restates it as the untrimmed + thread does.""" + off, on = _notified() + assert h.digests(on.model.requests) == h.digests(off.model.requests) + assert h.digests(on.summary_model.requests) == h.digests(off.summary_model.requests) + # Turns 7 and 8, the ones without the switch, trimmed. + bases = [base for _, base in on.held] + assert bases[5] < bases[6] < bases[7] + + +def test_a_schema_without_the_field_would_lose_the_notice( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Why every build declares ``MainAgentState``: without the switch's + field in the schema, a trim in a turn without the switch drops what its + carry kept, and the comparison sees it.""" + monkeypatch.setattr(h, "MainAgentState", DeltaAgentState) + off, on = _notified() + assert h.digests(on.model.requests) != h.digests(off.model.requests) + + +# ------------------------------------------------------------ the transcript + + +def test_the_transcript_is_the_whole_threads(compared: Compared) -> None: + (off_saves, off_files), (on_saves, on_files) = compared.scenario.transcripts + assert on_files == off_files + assert on_saves == off_saves + assert sorted(on_files) == [f"turn-{n:04d}.jsonl" for n in range(1, len(h.SCENARIO) + 1)] + + +def test_each_turn_end_renders_as_the_whole_thread_would(compared: Compared) -> None: + """Digests and shape keys included: the carried entries and the chain + through them are what a render of every message produces.""" + for run in (compared.off, compared.on): + exports = run.exports[: compared.scenario.exports] + wholes = h.whole_threads(run)[: compared.scenario.exports] + for (_, base, manifest), whole in zip(exports, wholes): + assert manifest == build_directory(whole, window=Window()).manifest, base + assert any(base for _, base, _ in compared.on.exports[: compared.scenario.exports]) + + +# ------------------------------------------------------------ an edit + + +def test_an_edit_into_a_trimmed_turn(compared: Compared) -> None: + """Turn 3 was trimmed long before; editing it forks from the end of + turn 2, whose checkpoint still holds the whole list, and the thread runs + on from there exactly as the untrimmed one does, trimming again.""" + on, off = compared.on, compared.off + assert compared.edited.calls > compared.scenario.calls + assert h.digests(on.model.requests) == h.digests(off.model.requests) + assert h.digests(on.summary_model.requests) == h.digests(off.summary_model.requests) + edited = on.held[len(h.SCENARIO) :] + assert edited[0][1] == 0 + assert edited[-1][1] > 0 + first = on.model.requests[compared.scenario.calls] + humans = [m for m in first if m.type == "human" and m.id and m.id.startswith("h-")] + assert humans[-1].id == "h-3" + assert "edited" in str(humans[-1].content) + (off_saves, off_files), (on_saves, on_files) = compared.edited.transcripts + assert on_files == off_files and on_saves == off_saves + + +# ------------------------------------------------------------ the comparison + + +def _without(key: str): + """``trim``, its window kept without ``key``.""" + real = middleware_module.trim + + def trimmed(state, cut, carries=(), **kwargs): + update = real(state, cut, carries, **kwargs) + update["_window"] = {**update["_window"], key: 0 if key == "runs" else []} + return update + + return trimmed + + +@pytest.mark.parametrize( + "broken", ["numbering", "earlier index", "notes count", "subagent notice"] +) +def test_the_comparison_sees_each_reading_the_trim_keeps( + compared: Compared, monkeypatch: pytest.MonkeyPatch, broken: str +) -> None: + """A window that lost any one of them sends different requests, so the + identity above is not one the scenario would show anyway.""" + if broken == "numbering": + monkeypatch.setattr(middleware_module, "trim", _without("runs")) + elif broken == "earlier index": + monkeypatch.setattr(middleware_module, "trim", _without("requests")) + elif broken == "notes count": + monkeypatch.setattr(h, "notes_window_carry", lambda thread: lambda trimmed, state: {}) + else: + monkeypatch.setattr(h, "subagents_window_carry", lambda trimmed, state: {}) + on = asyncio.run(h.run_scenario(InMemorySaver(), window=True)) + expected = h.digests(compared.off.model.requests[: compared.scenario.calls]) + assert h.digests(on.model.requests) != expected diff --git a/tests/unit/middleware/compaction/window_harness.py b/tests/unit/middleware/compaction/window_harness.py new file mode 100644 index 000000000..f14edaf4a --- /dev/null +++ b/tests/unit/middleware/compaction/window_harness.py @@ -0,0 +1,520 @@ +"""A deterministic thread run through the main agent's context middleware, with +the window trimming its checkpoint or not, for the suites that compare the two. + +The graph is langchain's ``create_agent`` over the middleware that read or +write the history the model is sent, in the main stack's order: the subagent +switch, compaction (summaries, Tier 1 cuts, the window), tool-call patching, +the turn and baseline rows, the notes reminders and the tail envelope. The +model and the summary model are fakes that record every request; the +transcript store is ``build_directory`` kept in memory, saved live by +compaction and again at each turn end from the checkpoint, as the server's +export does. Everything an id or a clock could vary is fixed or normalized, +so two runs of the same scenario compare as strings. +""" + +from __future__ import annotations + +import json +import re +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from typing import Any + +from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware +from langchain.agents import create_agent +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_core.tools import tool +from langgraph.checkpoint.serde.types import _DeltaSnapshot +from langgraph.types import Overwrite +from langsmith import tracing_context +from pydantic import Field + +from ptc_agent.agent.context_stack import build_context_middleware +from ptc_agent.agent.main_state import MainAgentState +from ptc_agent.agent.middleware.runtime_context import durable +from ptc_agent.agent.middleware.compaction import notes as notes_module +from ptc_agent.agent.middleware.compaction.compact import Summarizer +from ptc_agent.agent.middleware.compaction.middleware import CompactionMiddleware +from ptc_agent.agent.middleware.compaction.notes import ( + NotesDueMiddleware, + ThreadScratchpad, + notes_window_carry, + thread_notes_subdir, +) +from ptc_agent.agent.middleware.compaction.types import OffloadSettings +from ptc_agent.agent.middleware.subagent_switch import ( + SubagentSwitchMiddleware, + subagents_window_carry, +) +from ptc_agent.agent.transcript import TranscriptTarget, Window, build_directory, load_manifest +from ptc_agent.agent.transcript.render import split_runs + +THREAD = "abcd1234-5678-4000-8000-00000000c0de" +WORKSPACE = "/home/workspace" +NOTES_SUBDIR = thread_notes_subdir(THREAD) +NOTES_DIR = f"{WORKSPACE}/{NOTES_SUBDIR}" + +#: Messages in view, the system prompt counted, at which a summary runs. +THRESHOLD = 26 +KEEP = 6 +#: Tool calls since the notes were last written at which a check-in is due. +CHECK_IN_CALLS = 4 + +_START = datetime(2026, 9, 1, 9, 0, tzinfo=UTC) + + +class _Clock(datetime): + """The wall clock rows are stamped from, stopped at the turn's start.""" + + at = _START + + @classmethod + def now(cls, tz=None): + return cls.at + + +@dataclass(frozen=True) +class Turn: + """One turn of the scenario.""" + + #: Tool rounds before the answer. + rounds: int = 1 + #: The subagent switch as the turn reads it. + subagents: bool = True + #: Whether the build reads the switch at all; a notification turn's + #: does not. + switch: bool = True + #: Whether the first round writes the notes. + notes: bool = False + #: Tool rounds an orchestrator re-entry runs after the turn, 0 for none. + reentry: int = 0 + + +#: Long enough for several summaries and trims, with a flip of the subagent +#: switch and a notes write the trims drop, and a long turn that compacts +#: mid-turn before a background task's report re-enters it. +SCENARIO: tuple[Turn, ...] = ( + Turn(), + Turn(notes=True), + Turn(subagents=False), + Turn(subagents=False), + Turn(), + Turn(rounds=2), + Turn(), + Turn(rounds=6, reentry=2), + Turn(), + Turn(rounds=2), + Turn(), + Turn(), +) + +#: The scenario with the two turns whose trims drop the subagent notices +#: built without the switch. +NOTIFIED: tuple[Turn, ...] = tuple( + replace(turn, switch=False) if number in (7, 8) else turn + for number, turn in enumerate(SCENARIO, start=1) +) + + +class RecordingModel(BaseChatModel): + """Answers from a script and records every request it is sent.""" + + script: list[AIMessage] = Field(default_factory=list) + requests: list[list[AnyMessage]] = Field(default_factory=list) + + @property + def _llm_type(self) -> str: + return "recording" + + def bind_tools(self, tools: Any, **kwargs: Any) -> RecordingModel: + return self + + def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult: + self.requests.append(list(messages)) + if not self.script: + raise AssertionError("the model was called more often than scripted") + return ChatResult(generations=[ChatGeneration(message=self.script.pop(0))]) + + +class SummaryModel(BaseChatModel): + """Writes numbered summaries and records what it is asked.""" + + requests: list[list[AnyMessage]] = Field(default_factory=list) + + @property + def _llm_type(self) -> str: + return "summary" + + def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult: + self.requests.append(list(messages)) + text = f"Summary {len(self.requests)}: the user asked for quotes and files." + return ChatResult(generations=[ChatGeneration(message=AIMessage(text))]) + + +class TranscriptStore: + """The thread's stored transcript, as the server's store keeps it: each + save renders against the manifest the last one left.""" + + def __init__(self) -> None: + self.manifest: str | None = None + self.files: dict[str, bytes] = {} + self.saves: list[str] = [] + + async def save_transcript(self, target, messages, *, window: Window) -> bool: + directory = build_directory( + list(messages), + unit=target.unit, + previous=load_manifest(self.manifest), + window=window, + ) + self.manifest = directory.manifest + self.files.update(directory.rendered) + self.saves.append(directory.manifest) + return True + + +def _backend(store: TranscriptStore) -> Any: + async def settled_livefs(workspace_id=None): + return store + + async def aglob(pattern, path=None): + return SimpleNamespace(error=None, matches=[{"path": f"{NOTES_DIR}/plan.md"}]) + + return SimpleNamespace( + livefs=store, + settled_livefs=settled_livefs, + aglob=aglob, + normalize_path=lambda path: path, + ) + + +@tool("Write") +def write_file(file_path: str, content: str) -> str: + """Write a file.""" + return f"wrote {file_path}" + + +@tool("fetch_quote") +def fetch_quote(symbol: str) -> str: + """A quote.""" + return f"{symbol} last 101.25, volume 2.1M, range 99.80-102.40" + + +def _round(turn: int, round_: str, *, notes: bool) -> list[AIMessage]: + """A file write, then a quote: one call per message, so the tool results + land in one order whichever tool finishes first.""" + tag = f"t{turn}r{round_}" + path = f"{NOTES_DIR}/plan.md" if notes else f"{WORKSPACE}/work/{tag}.py" + write = { + "name": "Write", + "id": f"call-w-{tag}", + "args": {"file_path": path, "content": f"# {tag}\n" + "x = 1\n" * 120}, + } + quote = {"name": "fetch_quote", "id": f"call-q-{tag}", "args": {"symbol": f"S{tag}"}} + return [ + AIMessage("", id=f"ai-{tag}-w", tool_calls=[write]), + AIMessage(f"Wrote it; quoting {tag}.", id=f"ai-{tag}-q", tool_calls=[quote]), + ] + + +def _script(turn: int, rounds: int, *, notes: bool = False, label: str = "") -> list[AIMessage]: + out = [ + message + for j in range(rounds) + for message in _round(turn, f"{label}{j}", notes=notes and j == 0) + ] + out.append(AIMessage(f"Answer to turn {turn}{label}.", id=f"ai-t{turn}{label}-end")) + return out + + +Coverage = Callable[[Window], Any] + + +@dataclass +class ThreadRun: + """What one run of the scenario left behind.""" + + graph: Any = None + model: RecordingModel = field(default_factory=RecordingModel) + summary_model: SummaryModel = field(default_factory=SummaryModel) + store: TranscriptStore = field(default_factory=TranscriptStore) + #: Per turn, after it (and its re-entry): messages held and window base. + held: list[tuple[int, int]] = field(default_factory=list) + #: Each trim's question to the coverage: the runs it would leave trimmed, + #: whether its head is exactly the start of the whole thread, and how + #: many summaries had been written. + asked: list[tuple[int, bool, int]] = field(default_factory=list) + #: The messages of the turn the re-entry ran in, before and after it. + reentry: tuple[list[str], list[str]] | None = None + #: The window base before and after the re-entry. + reentry_base: tuple[int, int] | None = None + #: Per turn end: the messages exported, their base and the manifest saved. + exports: list[tuple[list[AnyMessage], int, str]] = field(default_factory=list) + #: Per turn, the checkpoint its end left. + tails: list[str] = field(default_factory=list) + + +def build_graph( + run: ThreadRun, + saver: Any, + *, + now: datetime, + turn: Turn, + window: bool, + covered: bool = True, +) -> Any: + """The main agent's context stack for ``turn`` around ``run``'s fakes, as of ``now``, + the server answering every coverage check with ``covered``.""" + + async def switch() -> bool: + return turn.subagents + + compaction = CompactionMiddleware( + Summarizer(run.summary_model, limit=10**9, counter=len), + token_threshold=THRESHOLD, + keep_messages=KEEP, + offload=OffloadSettings(keep_messages=KEEP, max_length=200, idle_seconds=0.0), + backend=_backend(run.store), + workspace_id="ws-1", + ).with_scratchpad_notes(NOTES_DIR) + scratchpad = ThreadScratchpad( + folder=f"{WORKSPACE}/.agents/scratchpad/", notes_dir=NOTES_DIR, notes_subdir=NOTES_SUBDIR + ) + notes_rows = NotesDueMiddleware(scratchpad, compaction) + if window: + + async def coverage(head: Window) -> bool: + # Not checked against the head, so a test that breaks the window + # still sees it trim; the answer it should have had is kept. + whole = whole_threads(run)[-1] if run.exports else [] + run.asked.append( + ( + head.runs, + head.holds(whole[: head.messages]), + len(run.summary_model.requests), + ) + ) + return covered + + compaction = compaction.with_window( + coverage, [notes_window_carry(THREAD), subagents_window_carry] + ) + context = build_context_middleware( + now=now, guidance=None, model_name=None, sandbox_enabled=False + ) + return create_agent( + run.model, + system_prompt="You are the test analyst.", + tools=[write_file, fetch_quote], + middleware=[ + *([SubagentSwitchMiddleware(switch)] if turn.switch else []), + compaction, + PatchToolCallsMiddleware(), + context.turn, + context.baseline, + notes_rows, + context.tail, + ], + checkpointer=saver, + state_schema=MainAgentState, + ) + + +def config(checkpoint_id: str | None = None, *, turn: int | None = None) -> dict[str, Any]: + """The thread's config; with ``turn``, stamped as the server stamps a + turn's input, which pairs it with its query row.""" + configurable: dict[str, Any] = {"thread_id": THREAD} + if checkpoint_id: + configurable["checkpoint_ns"] = "" + configurable["checkpoint_id"] = checkpoint_id + out: dict[str, Any] = {"configurable": configurable} + if turn is not None: + out["metadata"] = {"turn_index": turn - 1, "run_id": f"run-{turn}"} + return out + + +async def export_turn_end(run: ThreadRun, graph: Any) -> None: + """The server's turn-end export: the thread's state, rendered over the + stored copy from the window's base.""" + values = (await graph.aget_state(config())).values + messages, window = list(values["messages"]), Window.of(values) + await run.store.save_transcript(TranscriptTarget(THREAD), messages, window=window) + run.exports.append((messages, window.runs, run.store.manifest or "")) + + +async def run_turn( + run: ThreadRun, saver: Any, number: int, turn: Turn, *, window: bool, text: str | None = None, + checkpoint_id: str | None = None, covered: bool = True, +) -> Any: + _Clock.at = _START + timedelta(hours=number) + graph = build_graph( + run, + saver, + now=_Clock.at, + turn=turn, + window=window, + covered=covered, + ) + run.model.script = _script(number, turn.rounds, notes=turn.notes) + question = HumanMessage(text or f"Question {number}: quote and file it.", id=f"h-{number}") + await graph.ainvoke({"messages": [question]}, config(checkpoint_id, turn=number)) + if turn.reentry: + values = (await graph.aget_state(config())).values + before = [m.id for m in _turn_messages(values, number)] + base_before = Window.of(values).runs + notice = HumanMessage( + "Background task finished; collect its result.", + id=f"orch-{number}", + name="orchestrator", + additional_kwargs={"lc_source": "orchestrator"}, + ) + await graph.aupdate_state(config(), {"messages": [notice]}, as_node="__start__") + run.model.script = _script(number, turn.reentry, label="b") + await graph.ainvoke(None, config()) + values = (await graph.aget_state(config())).values + run.reentry = (before, [m.id for m in _turn_messages(values, number)]) + run.reentry_base = (base_before, Window.of(values).runs) + await export_turn_end(run, graph) + snapshot = await graph.aget_state(config()) + run.held.append((len(snapshot.values["messages"]), Window.of(snapshot.values).runs)) + run.tails.append(snapshot.config["configurable"]["checkpoint_id"]) + run.graph = graph + return graph + + +def _turn_messages(values: Mapping[str, Any], number: int) -> list[AnyMessage]: + messages = values["messages"] + start = next(i for i, m in enumerate(messages) if m.id == f"h-{number}") + return messages[start:] + + +async def run_scenario( + saver: Any, *, window: bool, turns=SCENARIO, covered: bool = True +) -> ThreadRun: + run = ThreadRun() + with frozen(): + for number, turn in enumerate(turns, start=1): + await run_turn(run, saver, number, turn, window=window, covered=covered) + return run + + +async def run_edit( + run: ThreadRun, saver: Any, number: int, turns=SCENARIO, *, window: bool +) -> None: + """An edit of turn ``number``: the thread forked from the end of the turn + before it, then the scenario's turns from there run again, the edited + one asked differently.""" + fork = run.tails[number - 2] + with frozen(): + for offset, turn in enumerate(turns[number - 1 :]): + at = number + offset + await run_turn( + run, + saver, + at, + turn, + window=window, + text=f"Question {at}, edited: quote it again." if offset == 0 else None, + checkpoint_id=fork if offset == 0 else None, + ) + + +def whole_threads(run: ThreadRun) -> list[list[AnyMessage]]: + """Each turn end's export with the runs trimmed before it put back, from + the exports before: what the transcript renders when nothing is trimmed.""" + whole: list[AnyMessage] = [] + out = [] + for messages, base, _ in run.exports: + whole = [m for trimmed in split_runs(whole)[:base] for m in trimmed] + messages + out.append(whole) + return out + + +@contextmanager +def frozen(): + """The row clock stopped and the check-in brought within the scenario. + Tracing is off: ``load_dotenv`` may have found a checkout's ``.env``.""" + saved = notes_module.NOTES_CHECK_IN_CALLS, durable.datetime + notes_module.NOTES_CHECK_IN_CALLS, durable.datetime = CHECK_IN_CALLS, _Clock + try: + with tracing_context(enabled=False): + yield + finally: + notes_module.NOTES_CHECK_IN_CALLS, durable.datetime = saved + + +# ---------------------------------------------------------------- comparing + +_ID = re.compile( + r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}|lc_run--[0-9a-f-]+" +) + + +class Normalizer: + """Ids minted at run time, renamed in order of first appearance, so two + runs that mint them in the same places read the same.""" + + def __init__(self) -> None: + self.names: dict[str, str] = {} + + def __call__(self, text: str) -> str: + def name(match: re.Match[str]) -> str: + value = match.group(0) + if value == THREAD: + return value + return self.names.setdefault(value, f"") + + return _ID.sub(name, text) + + +def dump(messages: list[AnyMessage]) -> str: + """A request as text, every field the message carries.""" + return json.dumps( + [message.model_dump() for message in messages], sort_keys=True, default=str + ) + + +def digests(requests: list[list[AnyMessage]]) -> list[str]: + normalize = Normalizer() + return [normalize(dump(request)) for request in requests] + + +def transcript(store: TranscriptStore) -> tuple[list[str], dict[str, str]]: + """The saved manifests and the files, ids normalized. A manifest's + digests hash the minted ids, so they are left out: the files they digest + are compared instead, and a run checks its digests against a whole render + (``whole_threads``).""" + normalize = Normalizer() + return ( + [normalize(_without_digests(save)) for save in store.saves], + {name: normalize(data.decode()) for name, data in sorted(store.files.items())}, + ) + + +def _without_digests(manifest: str) -> str: + data = json.loads(manifest) + for entry in data.get("segments") or (): + entry.pop("sha256", None) + entry.pop("shape", None) + entry.pop("through", None) + return json.dumps(data, sort_keys=True) + + +async def decoded_messages(saver: Any, graph: Any) -> int: + """How many messages loading the thread's latest checkpoint decodes: the + seed it starts from plus every write replayed onto it.""" + tip = (await graph.aget_state(config())).config + history = (await saver.aget_delta_channel_history(config=tip, channels=["messages"]))[ + "messages" + ] + seed = history.get("seed") + count = len(seed.value if isinstance(seed, _DeltaSnapshot) else seed or []) + for _, _, value in history["writes"]: + value = value.value if isinstance(value, Overwrite) else value + count += len(value) if isinstance(value, list) else 1 + return count diff --git a/tests/unit/middleware/skills/test_skill_content.py b/tests/unit/middleware/skills/test_skill_content.py index 1eb62fee3..616d5ddc5 100644 --- a/tests/unit/middleware/skills/test_skill_content.py +++ b/tests/unit/middleware/skills/test_skill_content.py @@ -268,12 +268,15 @@ def test_compute_no_compaction_trusts_loaded_set(): assert result == {"chart-annotation"} -def test_compute_event_without_cutoff_trusts_loaded_set(): - """A summarization event with no/zero cutoff_index means nothing was actually - summarized away, so trust the loaded set as-is.""" +def test_compute_event_at_index_zero_still_reads_the_view(): + """A summary whose boundary sits at index 0 still stands in for history: + once the window trims the messages before its run, a later summary's + boundary can be the first message left. The body has to be in view.""" assert compute_already_loaded( - ["chart-annotation"], [{"content": "hi"}], {"cutoff_index": 0} - ) == {"chart-annotation"} + ["chart-annotation"], + [{"content": "hi", "id": "m0"}], + {"cutoff_index": 0, "summary_message": {"content": "compacted summary"}}, + ) == set() def test_compute_compaction_drops_body_before_cutoff(): diff --git a/tests/unit/middleware/test_compaction_request_shape.py b/tests/unit/middleware/test_compaction_request_shape.py index 047b4a1db..16a752c78 100644 --- a/tests/unit/middleware/test_compaction_request_shape.py +++ b/tests/unit/middleware/test_compaction_request_shape.py @@ -193,14 +193,14 @@ def _raw(turns: int) -> list: def test_turns_are_numbered_over_the_full_checkpoint(self): from ptc_agent.agent.middleware.compaction.utils import build_summary_message - from ptc_agent.agent.transcript import TranscriptTarget + from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.pointer import SummarySpan, TranscriptTurns raw = self._raw(5) target = TranscriptTarget("abcd1234-0000") prior = build_summary_message("turns 1-2", target, span=SummarySpan(1, 2)) sent = [prior, *raw[4:8]] - turns = TranscriptTurns.of(target, raw) + turns = TranscriptTurns.of(target, raw, window=Window()) system, human = _build_summary_request("sys", sent, turns) @@ -218,12 +218,13 @@ def test_turns_are_numbered_over_the_full_checkpoint(self): assert "(turn-0007.jsonl)" in system.content def test_a_subagent_cites_its_own_runs(self): - from ptc_agent.agent.transcript import TranscriptTarget + from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.pointer import TranscriptTurns raw = self._raw(2) target = TranscriptTarget.for_agent("abcd1234-0000", "task:t1") - system, human = _build_summary_request("sys", raw, TranscriptTurns.of(target, raw)) + turns = TranscriptTurns.of(target, raw, window=Window()) + system, human = _build_summary_request("sys", raw, turns) assert "[transcript: run-0002.jsonl]\nHuman: q2" in human.content assert "(run-0007.jsonl)" in system.content @@ -524,7 +525,7 @@ async def _passthrough(backend, messages, **kwargs): async def _summarize(self, mw): messages = [HumanMessage(content="hi", id="h")] - request = SimpleNamespace(messages=messages, model=None) + request = SimpleNamespace(messages=messages, model=None, state={}) return (await mw._compact(request, messages, 1, None)).summary @pytest.mark.asyncio @@ -631,6 +632,7 @@ def test_a_turn_cut_by_the_earlier_summary_keeps_its_own_reply(self): from ptc_agent.agent.middleware.compaction.summarize import server_summary from ptc_agent.agent.middleware.compaction.utils import build_summary_message + from ptc_agent.agent.transcript import Window raw = [ HumanMessage(content="first request", id="h1"), @@ -643,7 +645,9 @@ def test_a_turn_cut_by_the_earlier_summary_keeps_its_own_reply(self): earlier = build_summary_message("First and second requests.", None) to_summarize = [earlier, raw[3], raw[4]] - text = server_summary(to_summarize, [raw[5]], raw_messages=raw, turns=None).text + text = server_summary( + to_summarize, [raw[5]], raw_messages=raw, turns=None, window=Window() + ).text second, third = text.split("## Turn 2\n", 1)[1].split("## Turn 3\n", 1) assert second.strip() == "Last reply: second reply" diff --git a/tests/unit/middleware/test_credit_gate.py b/tests/unit/middleware/test_credit_gate.py index 3c035b42a..838728163 100644 --- a/tests/unit/middleware/test_credit_gate.py +++ b/tests/unit/middleware/test_credit_gate.py @@ -37,7 +37,7 @@ ) from langgraph.types import Interrupt -from src.server.contracts.status import classify_interrupt_reason +from src.server.contracts.status import classify_interrupts # Copy the platform would author; langalpha only ever relays it. _DENIAL_COPY = "Out of credits. Top up to keep this run going." @@ -523,8 +523,10 @@ async def test_pause_payload_classifies_as_a_credit_pause(): lease.denial = {"message": _DENIAL_COPY} payload = build_pause_payload(lease) - assert classify_interrupt_reason([Interrupt(value=payload)]) == "credit_pause" - assert payload["action_requests"][0]["message"] == _DENIAL_COPY + assert classify_interrupts([Interrupt(value=payload)]) == ( + "credit_pause", + _DENIAL_COPY, + ) @pytest.mark.asyncio @@ -535,8 +537,9 @@ async def test_a_denial_without_copy_still_classifies_and_still_explains(): lease.denial = {"message": None} payload = build_pause_payload(lease) - assert classify_interrupt_reason([Interrupt(value=payload)]) == "credit_pause" - assert payload["action_requests"][0]["message"] + reason, message = classify_interrupts([Interrupt(value=payload)]) + assert reason == "credit_pause" + assert message # -- what the family reserves against ------------------------------------- diff --git a/tests/unit/middleware/test_returned_steering_facts.py b/tests/unit/middleware/test_returned_steering_facts.py new file mode 100644 index 000000000..4f1782851 --- /dev/null +++ b/tests/unit/middleware/test_returned_steering_facts.py @@ -0,0 +1,224 @@ +"""Steering handed back to the user is kept on the run's replay facts. + +Replay reads a run's transcript from its checkpoint, where returned steering +never lands, so each sweep records what it returned on the run's ledger row. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langchain_core.messages import AIMessage, HumanMessage + +from ptc_agent.agent.middleware.background_subagent.middleware import ( + current_background_tool_call_id, +) +from ptc_agent.agent.middleware.background_subagent.redis_stream import ( + record_returned, +) +from ptc_agent.agent.middleware.background_subagent.registry import BackgroundTask +from ptc_agent.agent.middleware.background_subagent.run_executor import ( + _return_unconsumed_steering, +) +from ptc_agent.agent.middleware.background_subagent.steering import ( + SubagentSteeringMiddleware, +) +from tests.unit.server.handlers.chat.redis_fakes import FakeCache + + +def _task(run_id: str | None = "run-9") -> BackgroundTask: + return BackgroundTask( + tool_call_id="tc-9", + task_id="abc123", + description="d", + prompt="p", + subagent_type="general-purpose", + task_run_id=run_id, + ) + + +_RETURNED = [{"content": "c", "input_id": "i1", "reason": "run_ended"}] + + +@pytest.mark.asyncio +async def test_record_writes_the_returns_to_the_runs_row(): + ledger = SimpleNamespace(record_replay_facts=AsyncMock(return_value=True)) + + assert await record_returned( + SimpleNamespace(run_ledger=ledger), "run-9", _RETURNED + ) + + ledger.record_replay_facts.assert_awaited_once_with( + "run-9", steering_returned=_RETURNED + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "outcome", + [{"side_effect": RuntimeError("db down")}, {"return_value": False}], + ids=["raises", "no_row"], +) +async def test_record_reports_a_fact_that_did_not_land_without_raising(outcome): + ledger = SimpleNamespace(record_replay_facts=AsyncMock(**outcome)) + + assert not await record_returned( + SimpleNamespace(run_ledger=ledger), "run-9", _RETURNED + ) + + +@pytest.mark.asyncio +async def test_record_skips_when_nothing_was_returned_or_no_run(): + ledger = SimpleNamespace(record_replay_facts=AsyncMock()) + registry = SimpleNamespace(run_ledger=ledger) + + # Nothing is owed a row, so nothing holds a caller's queue back. + assert await record_returned(registry, "run-9", []) + assert await record_returned(registry, None, _RETURNED) + assert await record_returned(SimpleNamespace(), "run-9", _RETURNED) + + ledger.record_replay_facts.assert_not_awaited() + + +def _sweep(task, ledger): + """A run-end sweep over one queued entry whose append lands.""" + payload = json.dumps( + {"content": "late note", "expected_task_run_id": "run-9", "input_id": "i1"} + ) + + def _surface(*_a, **_k): + task.captured_event_seq += 1 + + client = SimpleNamespace(lrange=AsyncMock(return_value=[payload]), lrem=AsyncMock()) + registry = SimpleNamespace( + append_event_for_task=AsyncMock(side_effect=_surface), run_ledger=ledger + ) + return client, registry + + +@pytest.mark.asyncio +async def test_run_end_sweep_records_what_it_returns(): + task = _task("run-9") + ledger = SimpleNamespace(record_replay_facts=AsyncMock(return_value=True)) + client, registry = _sweep(task, ledger) + cache = SimpleNamespace(enabled=True, client=client) + + with patch("src.utils.cache.redis_cache.get_cache_client", return_value=cache): + await _return_unconsumed_steering(registry, task) + + returned = {"content": "late note", "input_id": "i1", "reason": "run_ended"} + ledger.record_replay_facts.assert_awaited_once_with( + "run-9", steering_returned=[returned] + ) + # The event names the return exactly as the fact does, plus its lane. + event = registry.append_event_for_task.await_args.args[1] + assert event["event"] == "steering_returned" + assert event["data"] == {"agent": "task:abc123", **returned} + client.lrem.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "outcome", + [{"side_effect": RuntimeError("db down")}, {"return_value": False}], + ids=["raises", "no_row"], +) +async def test_run_end_sweep_keeps_the_queue_when_the_returns_are_not_recorded( + outcome, +): + """Replay reads returned steering from the run's row, so a queue entry + whose return never landed there stays in Redis until its TTL.""" + task = _task("run-9") + ledger = SimpleNamespace(record_replay_facts=AsyncMock(**outcome)) + client, registry = _sweep(task, ledger) + cache = SimpleNamespace(enabled=True, client=client) + + with patch("src.utils.cache.redis_cache.get_cache_client", return_value=cache): + await _return_unconsumed_steering(registry, task) + + ledger.record_replay_facts.assert_awaited_once() + client.lrem.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_mismatch_return_records_after_the_last_message(): + call = "call-1" + foreign = json.dumps( + {"content": "stale", "expected_task_run_id": "run-OLD", "input_id": "i2"} + ) + cache = FakeCache() + from ptc_agent.agent.middleware.background_subagent.redis_stream import ( + steering_queue_key, + ) + + cache.client.lists[steering_queue_key(call, "run-A")] = [foreign] + ledger = SimpleNamespace(record_replay_facts=AsyncMock()) + registry = SimpleNamespace( + _tasks={call: SimpleNamespace(task_id="k7Xm2p", task_run_id="run-A")}, + append_captured_event=AsyncMock(), + run_ledger=ledger, + ) + state = {"messages": [HumanMessage("hi", id="m1"), AIMessage("ok", id="m2")]} + + token = current_background_tool_call_id.set(call) + try: + with patch("src.utils.cache.redis_cache.get_cache_client", return_value=cache): + await SubagentSteeringMiddleware(registry).abefore_model(state, MagicMock()) + finally: + current_background_tool_call_id.reset(token) + + returned = {"content": "stale", "input_id": "i2", "reason": "run_mismatch"} + ledger.record_replay_facts.assert_awaited_once_with( + "run-A", steering_returned=[{**returned, "after": "m2"}] + ) + # ``after`` is the fact's alone; the event places itself by stream order. + _, event = registry.append_captured_event.await_args.args + assert event["event"] == "steering_returned" + assert event["data"] == {"agent": "task:k7Xm2p", **returned} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("row", "landed"), + [({"thread_id": "t-1", "parent_run_id": "resp-1"}, True), (None, False)], +) +async def test_the_ledger_reports_whether_the_returns_landed(row, landed): + from src.server.services import subagent_run_coordinator as coordinator + + refresh = MagicMock() + with ( + patch.object( + coordinator.facts_db, + "append_returned_steering", + AsyncMock(return_value=row), + ), + patch( + "src.server.services.history.replay.refresh.schedule_replay_refresh", + refresh, + ), + ): + ok = await coordinator.SubagentRunCoordinator("t-1").record_replay_facts( + "run-9", steering_returned=_RETURNED + ) + + assert ok is landed + assert refresh.call_args_list == ([(("t-1", "resp-1"),)] if landed else []) + + +@pytest.mark.asyncio +async def test_a_ledger_write_that_fails_reports_not_landed(): + from src.server.services import subagent_run_coordinator as coordinator + + with patch.object( + coordinator.facts_db, + "append_returned_steering", + AsyncMock(side_effect=RuntimeError("pool closed")), + ): + ok = await coordinator.SubagentRunCoordinator("t-1").record_replay_facts( + "run-9", steering_returned=_RETURNED + ) + + assert ok is False diff --git a/tests/unit/ptc_agent/agent/middleware/compaction/test_summary_message.py b/tests/unit/ptc_agent/agent/middleware/compaction/test_summary_message.py index cf26ad4ca..84a819212 100644 --- a/tests/unit/ptc_agent/agent/middleware/compaction/test_summary_message.py +++ b/tests/unit/ptc_agent/agent/middleware/compaction/test_summary_message.py @@ -16,7 +16,7 @@ build_summary_message, parse_summary_message, ) -from ptc_agent.agent.transcript import TranscriptTarget +from ptc_agent.agent.transcript import TranscriptTarget, Window from ptc_agent.agent.transcript.classify import LEGACY_SUMMARY_PREFIX from ptc_agent.agent.transcript.pointer import ( SummarySpan, @@ -91,7 +91,7 @@ def _turns(n): def test_span_names_the_turns_the_summary_covers(): raw = _turns(5) - turns = message_turns(raw) + turns = message_turns(raw, base=0) assert summary_span(turns, raw[:8], raw[:8]) == SummarySpan(1, 4) assert summary_span(turns, raw[:8], raw[4:8]) == SummarySpan(3, 4, (1, 2)) # An earlier summary heads the summarized list and is not in raw. Kept, @@ -106,7 +106,7 @@ def test_span_names_the_turns_the_summary_covers(): def test_span_flags_a_message_cut_partway(): # One huge message dominates: trimming keeps its tail under the same id. raw = _turns(3) - turns = message_turns(raw) + turns = message_turns(raw, base=0) tail = raw[2].model_copy(update={"content": "q2 (tail)"}) assert summary_span(turns, raw[:4], [tail, raw[3]]).gap == (1, 2) head_cut = raw[0].model_copy(update={"content": "1 (tail)"}) @@ -125,7 +125,7 @@ def test_a_second_summary_keeps_the_turns_the_first_left_out(): stretch = [first, raw[6]] event = build_summary_event( Summary("second", "model", stretch), - TranscriptTurns.of(TRANSCRIPT, raw), + TranscriptTurns.of(TRANSCRIPT, raw, window=Window()), raw_messages=raw, to_summarize=stretch, preserved=raw[7:], @@ -150,7 +150,7 @@ def test_event_hands_back_skills_whose_body_was_summarized(): ): event = build_summary_event( Summary("the summary", "model", raw[:4]), - TranscriptTurns.of(TRANSCRIPT, raw), + TranscriptTurns.of(TRANSCRIPT, raw, window=Window()), raw_messages=raw, to_summarize=raw[:4], preserved=raw[4:], diff --git a/tests/unit/ptc_agent/agent/middleware/compaction/test_transcript_pointer.py b/tests/unit/ptc_agent/agent/middleware/compaction/test_transcript_pointer.py index 8b18e0dfa..6eed9653d 100644 --- a/tests/unit/ptc_agent/agent/middleware/compaction/test_transcript_pointer.py +++ b/tests/unit/ptc_agent/agent/middleware/compaction/test_transcript_pointer.py @@ -30,7 +30,7 @@ class _Mount: def __init__(self, outcome): self.outcome = outcome - async def save_transcript(self, target, messages): + async def save_transcript(self, target, messages, *, window): if self.outcome == "raises": raise ConnectionError("store down") if self.outcome == "hangs": diff --git a/tests/unit/ptc_agent/agent/test_delta_channel_determinism.py b/tests/unit/ptc_agent/agent/test_delta_channel_determinism.py index 0b26e5127..0a972dd46 100644 --- a/tests/unit/ptc_agent/agent/test_delta_channel_determinism.py +++ b/tests/unit/ptc_agent/agent/test_delta_channel_determinism.py @@ -115,7 +115,7 @@ def _simulate_hard_stop_flush(graph, config) -> None: def test_head_checkpoint_omits_messages_on_non_snapshot_step(): """Head ``channel_values`` omits ``messages`` (sentinel / non-snapshot step). - With ``snapshot_frequency=50`` (DeltaAgentState's default) a thread of + With DeltaAgentState's ``MESSAGES_SNAPSHOT_FREQUENCY`` a thread of ~18 messages never hits a snapshot step, so the latest checkpoint blob is a sentinel and the raw ``channel_values`` dict has no ``messages`` key. This is the precondition that makes the determinism property meaningful — without it, diff --git a/tests/unit/ptc_agent/agent/test_main_state.py b/tests/unit/ptc_agent/agent/test_main_state.py new file mode 100644 index 000000000..8f687dc80 --- /dev/null +++ b/tests/unit/ptc_agent/agent/test_main_state.py @@ -0,0 +1,57 @@ +"""The main agent's one schema declares what every middleware keeps, so a +graph built on it (a turn, thread maintenance, the history reader) writes no +checkpoint without one of them. ``test_main_state_db`` shows the erasure on +Postgres.""" + +from __future__ import annotations + +from typing import Annotated, get_args, get_origin, get_type_hints + +from langchain.agents.middleware.types import AgentMiddleware, PrivateStateAttr +from langgraph.channels.delta import DeltaChannel +from langgraph.checkpoint.memory import InMemorySaver +from typing_extensions import NotRequired, Required + +import ptc_agent.agent.middleware # noqa: F401 (imports every middleware) +from ptc_agent.agent.main_state import MainAgentState +from src.server.services.history.reader import CheckpointHistoryReader + + +def _subclasses(cls: type) -> set[type]: + out = set() + for sub in cls.__subclasses__(): + out |= {sub, *_subclasses(sub)} + return out + + +def _private(annotation: object) -> bool: + if get_origin(annotation) in (Required, NotRequired): + (annotation,) = get_args(annotation) + return get_origin(annotation) is Annotated and PrivateStateAttr in annotation.__metadata__ + + +def test_it_declares_every_private_field_a_middleware_keeps() -> None: + declared = get_type_hints(MainAgentState, include_extras=True) + kept = { + (middleware.__name__, name): annotation + for middleware in _subclasses(AgentMiddleware) + if middleware.__module__.startswith("ptc_agent.") + for name, annotation in get_type_hints( + middleware.state_schema, include_extras=True + ).items() + if _private(annotation) + } + assert {"_subagents_trimmed", "_notes_calls_trimmed"} <= {name for _, name in kept} + for (middleware, name), annotation in kept.items(): + assert declared.get(name) == annotation, (middleware, name) + + +def test_messages_keep_their_delta_channel() -> None: + (annotated,) = get_args(get_type_hints(MainAgentState, include_extras=True)["messages"]) + assert any(isinstance(meta, DeltaChannel) for meta in annotated.__metadata__) + + +def test_the_reader_writes_with_the_main_agents_schema() -> None: + reader = CheckpointHistoryReader(InMemorySaver()) + for graph in (reader.graph, reader._updater): + assert graph.builder.state_schema is MainAgentState diff --git a/tests/unit/ptc_agent/agent/test_subagent_switch_wiring.py b/tests/unit/ptc_agent/agent/test_subagent_switch_wiring.py index 08a65c2de..90375044e 100644 --- a/tests/unit/ptc_agent/agent/test_subagent_switch_wiring.py +++ b/tests/unit/ptc_agent/agent/test_subagent_switch_wiring.py @@ -72,6 +72,14 @@ def test_no_reader_builds_without_it(machine): assert "SubagentSwitchMiddleware" not in _names(_build(**machine)) +def test_a_build_without_it_still_declares_its_state(machine): + """Thread maintenance builds without a reader; a write through that + build must not erase the last notice's value (``main_state``).""" + from ptc_agent.agent.main_state import MainAgentState + + assert _build(**machine)["state_schema"] is MainAgentState + + def test_a_build_without_subagents_has_nothing_to_hold(machine): names = _names( _build(subagent_switch=AsyncMock(), disable_subagents=True, **machine) @@ -143,3 +151,47 @@ def test_without_the_scratchpad_its_rows_stay_out_of_the_calls(machine): at = names.index("NotesOffMiddleware") assert names[at - 1] == "BaselineContextMiddleware" assert names[at + 1] == "TailEnvelopeMiddleware" + + +def _compactions(built: dict, handed: list | None = None): + from ptc_agent.agent.middleware import CompactionMiddleware + + stacks = [built["middleware"], *(handed or [])] + return [[m for m in stack if isinstance(m, CompactionMiddleware)] for stack in stacks] + + +def test_the_window_trims_the_main_agent_only(machine): + """A subagent's run is bounded, and its transcript reads its namespace's + whole history, so only the main stack's compaction gets the window.""" + from ptc_agent.agent.middleware.subagent_switch import subagents_window_carry + + real = SubAgentMiddleware + handed: list = [] + + def spy(*args, **kwargs): + handed.append(kwargs["default_middleware"]) + return real(*args, **kwargs) + + coverage = AsyncMock(return_value=True) + with patch("ptc_agent.agent.agent.SubAgentMiddleware", side_effect=spy): + built = _build(subagent_switch=AsyncMock(), window_coverage=coverage, **machine) + (main,), (sub,) = _compactions(built, handed) + assert main._coverage is coverage + assert subagents_window_carry in main._carries and len(main._carries) == 2 + assert sub._coverage is None + + +def test_the_window_needs_only_the_server(machine): + """A build without the switch's middleware (a notification turn) trims + too, the carry keeping a field ``MainAgentState`` declares; without the + server to ask there is nothing to trim against.""" + coverage = AsyncMock(return_value=True) + for kwargs in ( + {"subagent_switch": AsyncMock(), "disable_subagents": True, "window_coverage": coverage}, + {"window_coverage": coverage}, + ): + ((main,),) = _compactions(_build(**kwargs, **machine)) + assert main._coverage is coverage, kwargs + assert len(main._carries) == 2, kwargs + ((main,),) = _compactions(_build(subagent_switch=AsyncMock(), **machine)) + assert main._coverage is None diff --git a/tests/unit/ptc_agent/agent/test_transcript_render.py b/tests/unit/ptc_agent/agent/test_transcript_render.py index 62c1d98f0..789501553 100644 --- a/tests/unit/ptc_agent/agent/test_transcript_render.py +++ b/tests/unit/ptc_agent/agent/test_transcript_render.py @@ -20,7 +20,7 @@ build_update_message, ) from ptc_agent.agent.middleware.runtime_context.turn import TURN_ROW_KIND -from ptc_agent.agent.transcript import build_directory, load_manifest +from ptc_agent.agent.transcript import Window, build_directory, load_manifest from ptc_agent.agent.transcript.render import message_turns @@ -43,53 +43,53 @@ def _thread(turns: int) -> list: def test_a_turn_end_renders_only_the_turn_that_grew(): - before = build_directory(_thread(3)) + before = build_directory(_thread(3), window=Window()) steering = HumanMessage( content="steer", id="h-3b", additional_kwargs={"lc_source": "steering"} ) grown = [*_thread(3), steering] - after = build_directory(grown, previous=load_manifest(before.manifest)) + after = build_directory(grown, previous=load_manifest(before.manifest), window=Window()) assert set(after.rendered) == {"turn-0003.jsonl"} - full = build_directory(grown) + full = build_directory(grown, window=Window()) assert after.manifest == full.manifest assert after.files == full.files def test_a_new_turn_renders_alone(): - before = build_directory(_thread(3)) - after = build_directory(_thread(4), previous=load_manifest(before.manifest)) + before = build_directory(_thread(3), window=Window()) + after = build_directory(_thread(4), previous=load_manifest(before.manifest), window=Window()) assert set(after.rendered) == {"turn-0004.jsonl"} - assert after.files == build_directory(_thread(4)).files + assert after.files == build_directory(_thread(4), window=Window()).files def test_a_rewritten_turn_renders_again_with_every_turn_after_it(): - before = build_directory(_thread(4)) + before = build_directory(_thread(4), window=Window()) rewritten = _thread(4) rewritten[7] = AIMessage(content="regenerated", id="r-2-new") - after = build_directory(rewritten, previous=load_manifest(before.manifest)) + after = build_directory(rewritten, previous=load_manifest(before.manifest), window=Window()) assert set(after.rendered) == {"turn-0002.jsonl", "turn-0003.jsonl", "turn-0004.jsonl"} - assert after.files == build_directory(rewritten).files + assert after.files == build_directory(rewritten, window=Window()).files def test_a_message_replaced_under_its_id_with_same_length_text_renders_again(): - before = build_directory(_thread(3)) + before = build_directory(_thread(3), window=Window()) edited = _thread(3) # "question 2" -> "question 9": the reducer replaces by id, the length holds. edited[4] = HumanMessage(content="question 9", id="h-2") - after = build_directory(edited, previous=load_manifest(before.manifest)) + after = build_directory(edited, previous=load_manifest(before.manifest), window=Window()) assert set(after.rendered) == {"turn-0002.jsonl", "turn-0003.jsonl"} assert b"question 9" in after.rendered["turn-0002.jsonl"] - assert after.files == build_directory(edited).files + assert after.files == build_directory(edited, window=Window()).files def test_a_tool_call_whose_args_changed_renders_again(): - before = build_directory(_thread(2)) + before = build_directory(_thread(2), window=Window()) edited = _thread(2) edited[1] = AIMessage( content="", @@ -97,17 +97,17 @@ def test_a_tool_call_whose_args_changed_renders_again(): tool_calls=[{"id": "call-1", "name": "execute_code", "args": {"code": "2"}}], ) - after = build_directory(edited, previous=load_manifest(before.manifest)) + after = build_directory(edited, previous=load_manifest(before.manifest), window=Window()) assert set(after.rendered) == {"turn-0001.jsonl", "turn-0002.jsonl"} assert b'"code": "2"' in after.rendered["turn-0001.jsonl"] - assert after.files == build_directory(edited).files + assert after.files == build_directory(edited, window=Window()).files def test_a_manifest_of_another_schema_renders_everything(): - before = load_manifest(build_directory(_thread(2)).manifest) + before = load_manifest(build_directory(_thread(2), window=Window()).manifest) before["schema"] = 1 - after = build_directory(_thread(2), previous=before) + after = build_directory(_thread(2), previous=before, window=Window()) assert set(after.rendered) == {"turn-0001.jsonl", "turn-0002.jsonl"} @@ -117,12 +117,12 @@ def test_a_message_ahead_of_the_first_user_message_stays_with_the_first_turn(): and number every real turn one too high, the compaction pointer's too.""" thread = [SystemMessage(content="Answer tersely.", id="s-0"), *_thread(2)] - directory = build_directory(thread) + directory = build_directory(thread, window=Window()) assert set(directory.files) == {"turn-0001.jsonl", "turn-0002.jsonl"} first = json.loads(directory.rendered["turn-0001.jsonl"].splitlines()[0]) assert (first["type"], first["text"]) == ("user", "question 1") - turns = message_turns(thread) + turns = message_turns(thread, base=0) assert (turns["s-0"], turns["h-1"], turns["h-2"]) == (1, 1, 2) @@ -132,12 +132,12 @@ def test_a_lone_surrogate_renders_as_the_checkpoint_stores_it(): thread = _thread(2) thread[2] = ToolMessage(content="bad \ud800 byte", tool_call_id="call-1", id="t-1") - before = build_directory(thread[:4]) - after = build_directory(thread, previous=load_manifest(before.manifest)) + before = build_directory(thread[:4], window=Window()) + after = build_directory(thread, previous=load_manifest(before.manifest), window=Window()) assert b"bad ? byte" in before.rendered["turn-0001.jsonl"] assert set(after.rendered) == {"turn-0002.jsonl"} - full = build_directory(thread) + full = build_directory(thread, window=Window()) assert after.manifest == full.manifest assert after.files == full.files @@ -163,7 +163,7 @@ def test_the_turn_row_and_the_eviction_pointer_read_as_their_producers_write_the ToolMessage(content=pointer, tool_call_id="call-1", id="t-1"), ] - data = build_directory(messages).rendered["turn-0001.jsonl"] + data = build_directory(messages, window=Window()).rendered["turn-0001.jsonl"] events = [json.loads(line) for line in data.splitlines()] assert events[0]["type"] == "user" diff --git a/tests/unit/scripts/utils/test_replay_parity.py b/tests/unit/scripts/utils/test_replay_parity.py new file mode 100644 index 000000000..6084bdbdd --- /dev/null +++ b/tests/unit/scripts/utils/test_replay_parity.py @@ -0,0 +1,47 @@ +"""A stopped turn is excused only for output streamed past its last +checkpoint, never for a change to what the checkpoints committed.""" + +import pytest + +from scripts.utils.replay_parity import _streamed_past_commit + +TEXT = ("main", "text") +CALL = ("search", '{"q":"x"}') + + +@pytest.mark.parametrize( + ("key", "committed", "streamed"), + [ + ("text", {TEXT: "The answer"}, {TEXT: "The answer is 4"}), + ("text", {}, {TEXT: "Started"}), + ("signals", {"main": 1}, {"main": 3}), + ("elapsed_ms", {"main": [120]}, {"main": [120, 80]}), + ("tool_calls", {CALL}, {CALL, ("fetch", "{}")}), + ( + "results", + {"call-1": ("text", "ok")}, + {"call-1": ("text", "ok"), "call-2": ("text", "")}, + ), + ("results", None, {"call-1": ("text", "ok")}), + ], +) +def test_output_streamed_after_the_last_checkpoint_is_excused(key, committed, streamed): + assert _streamed_past_commit(key, committed, streamed) + + +@pytest.mark.parametrize( + ("key", "committed", "streamed"), + [ + ("text", {TEXT: "The answer is 4"}, {TEXT: "The answer"}), + ("text", {TEXT: "The answer"}, {TEXT: "An answer is 4"}), + ("text", {TEXT: "The answer"}, {}), + ("signals", {"main": 2}, {"main": 1}), + ("elapsed_ms", {"main": [120]}, {"main": [90, 80]}), + ("tool_calls", {CALL}, set()), + ("results", {"call-1": ("text", "ok")}, {"call-1": ("text", "changed")}), + ("results", {"call-1": ("text", "ok")}, None), + ("artifacts", [], ['{"id":"a-1"}']), + ], +) +def test_a_change_to_committed_output_is_a_diff(key, committed, streamed): + assert not _streamed_past_commit(key, committed, streamed) diff --git a/tests/unit/server/app/test_public_order_privacy.py b/tests/unit/server/app/test_public_order_privacy.py index 84ef793c2..08810519d 100644 --- a/tests/unit/server/app/test_public_order_privacy.py +++ b/tests/unit/server/app/test_public_order_privacy.py @@ -21,6 +21,7 @@ from httpx import ASGITransport, AsyncClient from tests.conftest import create_test_app +from tests.unit.server.services.history.replay_builders import replay_rows pytestmark = pytest.mark.asyncio @@ -28,8 +29,7 @@ _THREAD_ID = "44444444-4444-4444-8444-444444444444" _THREAD_BY_TOKEN = "src.server.app.share_access.get_thread_by_share_token" -_QUERIES = "src.server.app.public.get_queries_for_thread" -_RESPONSES = "src.server.app.public.get_responses_for_thread" +_REPLAY_DATA = "src.server.app.public.get_replay_thread_data" _TASK_DETAILS = "src.server.services.history.task_status.resolve_task_details" _RECEIPT = { @@ -182,8 +182,19 @@ async def _replay( ) with ( patch(_THREAD_BY_TOKEN, new=AsyncMock(return_value=thread)), - patch(_QUERIES, new=AsyncMock(return_value=(queries, None))), - patch(_RESPONSES, new=AsyncMock(return_value=(responses, None))), + patch( + _REPLAY_DATA, + new=AsyncMock( + return_value=replay_rows( + { + "conversation_thread_id": _THREAD_ID, + "latest_checkpoint_id": None, + }, + queries, + responses, + ) + ), + ), patch(_TASK_DETAILS, new=AsyncMock(return_value={})), ): resp = await client.get(f"/api/v1/public/shared/{_SHARE_TOKEN}/replay") diff --git a/tests/unit/server/app/test_public_redaction.py b/tests/unit/server/app/test_public_redaction.py index 55aef4b4c..e1cb61901 100644 --- a/tests/unit/server/app/test_public_redaction.py +++ b/tests/unit/server/app/test_public_redaction.py @@ -21,6 +21,7 @@ from tests.conftest import create_test_app from src.server.utils.secret_redactor import SecretRedactor +from tests.unit.server.services.history.replay_builders import replay_rows pytestmark = pytest.mark.asyncio @@ -797,8 +798,7 @@ async def test_truncated_false_when_fully_read(self, public_client): # TestReplayStripsWorkspaceId — workspace_id is a bearer credential # --------------------------------------------------------------------------- -_QUERIES = "src.server.app.public.get_queries_for_thread" -_RESPONSES = "src.server.app.public.get_responses_for_thread" +_REPLAY_DATA = "src.server.app.public.get_replay_thread_data" def _parse_sse_events(body: str): @@ -849,8 +849,19 @@ async def test_replay_strips_workspace_id_from_stored_events(self, public_client query = {"turn_index": 0, "content": "hi", "created_at": "2026-01-01T00:00:00Z"} with ( patch(_THREAD_BY_TOKEN, AsyncMock(return_value=_make_thread())), - patch(_QUERIES, AsyncMock(return_value=([query], 1))), - patch(_RESPONSES, AsyncMock(return_value=([stored], 1))), + patch( + _REPLAY_DATA, + AsyncMock( + return_value=replay_rows( + { + "conversation_thread_id": _THREAD_ID, + "latest_checkpoint_id": None, + }, + [query], + [stored], + ) + ), + ), ): resp = await public_client.get( f"/api/v1/public/shared/{_SHARE_TOKEN}/replay" diff --git a/tests/unit/server/app/test_public_replay_assembly.py b/tests/unit/server/app/test_public_replay_assembly.py new file mode 100644 index 000000000..40e693be8 --- /dev/null +++ b/tests/unit/server/app/test_public_replay_assembly.py @@ -0,0 +1,282 @@ +"""The public replay is the owner's page assembly plus redaction on the way out.""" + +from __future__ import annotations + +import json +from contextlib import ExitStack +from unittest.mock import AsyncMock, patch + +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient + +from src.server.services.history.replay import ReplayPage +from src.server.services.history.replay.lines import line_of +from tests.conftest import create_test_app +from tests.unit.server.services.history.replay_builders import replay_rows + +pytestmark = pytest.mark.asyncio + +_SHARE_TOKEN = "share_abc123" +_THREAD_ID = "44444444-4444-4444-8444-444444444444" + +_THREAD_BY_TOKEN = "src.server.app.share_access.get_thread_by_share_token" +_REPLAY_DATA = "src.server.app.public.get_replay_thread_data" +_READ_PAGE = "src.server.services.history.replay.read_replay_page" +_TASK_DETAILS = "src.server.services.history.task_status.resolve_task_details" + + +@pytest_asyncio.fixture +async def client(): + from src.server.app.public import router + + app = create_test_app(router) + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://test" + ) as c: + yield c + + +async def test_public_stream_is_the_shared_page_through_redaction(client): + thread = { + "conversation_thread_id": _THREAD_ID, + "workspace_id": "ws-1", + "share_permissions": {}, + } + row = {"conversation_thread_id": _THREAD_ID, "latest_checkpoint_id": None} + queries = [{"turn_index": 0, "content": "hi", "metadata": {}}] + responses = [{"turn_index": 0, "conversation_response_id": "r-1"}] + items = [ + { + "event": "user_message", + "data": { + "turn_index": 0, + "run_id": "run-secret", + "content": "hi", + "metadata": {"workspace_id": "ws-secret", "keep": "yes"}, + }, + }, + { + "event": "provenance", + "data": {"tool_call_id": "c1", "args": {"acc_id": "1234567"}}, + }, + {"event": "interrupt", "data": {"interrupt_id": "int-1"}}, + { + "event": "message_chunk", + "data": {"content_type": "text", "content": "hello", "id": "m1"}, + }, + ] + page = ReplayPage( + lines=[line_of(item) for item in items], first_turn_index=0, has_more=False + ) + read_page = AsyncMock(return_value=(page, "checkpoint")) + rows = replay_rows(row, queries, responses) + with ( + patch(_THREAD_BY_TOKEN, new=AsyncMock(return_value=thread)), + patch(_REPLAY_DATA, new=AsyncMock(return_value=rows)), + patch(_READ_PAGE, new=read_page), + patch(_TASK_DETAILS, new=AsyncMock(return_value={})), + ): + resp = await client.get(f"/api/v1/public/shared/{_SHARE_TOKEN}/replay") + assert resp.status_code == 200 + + events = [] + for block in resp.text.split("\n\n"): + kind = payload = None + for line in block.splitlines(): + if line.startswith("event: "): + kind = line[len("event: ") :] + elif line.startswith("data: "): + payload = json.loads(line[len("data: ") :]) + if kind: + events.append((kind, payload)) + + kinds = [k for k, _ in events] + assert kinds == ["user_message", "message_chunk", "replay_done"] + user = events[0][1] + assert "run_id" not in user + assert "workspace_id" not in (user.get("metadata") or {}) + assert events[1][1]["content"] == "hello" + + read_page.assert_awaited_once() + assert read_page.await_args.args == (rows,) + assert read_page.await_args.args[0] is rows + + +# Exception text as a provider or the database raises it: credential scrubbing +# leaves every part of it in place. +_FAULT = "connection to 10.0.3.7:5432 refused for org-7Qh2x in sandbox sbx-41c9" + + +def _events(body: str) -> list[tuple[str, dict]]: + out = [] + for block in body.split("\n\n"): + fields = dict( + line.split(": ", 1) for line in block.splitlines() if ": " in line + ) + if "event" in fields: + out.append((fields["event"], json.loads(fields["data"]))) + return out + + +async def _shared(client, row: dict, responses: list[dict], read_page=None) -> str: + thread = { + "conversation_thread_id": _THREAD_ID, + "workspace_id": "ws-1", + "share_permissions": {}, + } + queries = [{"turn_index": 0, "content": "hi", "metadata": {}}] + patches = [ + patch(_THREAD_BY_TOKEN, new=AsyncMock(return_value=thread)), + patch( + _REPLAY_DATA, + new=AsyncMock(return_value=replay_rows(row, queries, responses)), + ), + patch(_TASK_DETAILS, new=AsyncMock(return_value={})), + ] + if read_page is not None: + patches.append(patch(_READ_PAGE, new=read_page)) + with ExitStack() as stack: + for p in patches: + stack.enter_context(p) + resp = await client.get(f"/api/v1/public/shared/{_SHARE_TOKEN}/replay") + assert resp.status_code == 200 + return resp.text + + +_FAILED_TURN = { + "turn_index": 0, + "conversation_response_id": "r-1", + "status": "error", + "errors": [_FAULT], + "metadata": {"error_type": "connection", "error_class": "OperationalError"}, +} + + +async def test_a_projected_failure_reaches_a_share_without_its_text(client): + """The owner's assembly closes a failed turn with its error and settles a + failed workflow run from its ledger row; a viewer gets the statuses.""" + from src.server.services.history import projector + from src.server.services.history.slices import SpanDelta + from src.server.services.history.replay import items as replay_items + + stamped = SpanDelta( + new_ui_records=[ + { + "name": projector.WORKFLOW_RUN_UI_NAME, + "props": { + "frames": [ + { + "agent": "task:wf-1", + "run_id": "wf-1", + "phase": "run_completed", + "status": "failed", + "error": _FAULT, + } + ] + }, + }, + { + "name": projector.MODEL_FALLBACK_UI_NAME, + "props": {"from_model": "a", "to_model": "b", "error": _FAULT}, + }, + ] + ) + page_items = [ + *projector.workflow_run_items( + SpanDelta(), task_id="wf-2", ledger_status="failed", ledger_failure=_FAULT + ), + *projector.workflow_run_items(stamped, task_id="wf-1", ledger_status="failed"), + *projector.model_fallback_items(_THREAD_ID, stamped), + *replay_items.terminal_items(_THREAD_ID, _FAILED_TURN, []), + ] + assert sum(_FAULT in json.dumps(item) for item in page_items) == 4 + page = ReplayPage( + lines=[line_of(item) for item in page_items], + first_turn_index=0, + has_more=False, + ) + + body = await _shared( + client, + {"conversation_thread_id": _THREAD_ID, "latest_checkpoint_id": "cp-1"}, + [_FAILED_TURN], + read_page=AsyncMock(return_value=(page, "checkpoint")), + ) + + assert _FAULT not in body + events = _events(body) + assert [k for k, _ in events] == [ + "workflow_lifecycle", + "workflow_lifecycle", + "model_fallback", + "replay_done", + ] + assert [d["status"] for k, d in events if k == "workflow_lifecycle"] == [ + "failed", + "failed", + ] + + +async def test_a_stored_failure_reaches_a_share_without_its_text(client): + """The same for a turn replayed from its stored events.""" + response = { + **_FAILED_TURN, + "sse_events": [ + { + "event": "workflow_lifecycle", + "data": { + "agent": "task:wf-1", + "run_id": "wf-1", + "phase": "run_completed", + "status": "failed", + "error": _FAULT, + }, + }, + {"event": "model_fallback", "data": {"to_model": "b", "error": _FAULT}}, + {"event": "error", "data": {"error": _FAULT}}, + ], + } + + body = await _shared( + client, + {"conversation_thread_id": _THREAD_ID, "latest_checkpoint_id": None}, + [response], + ) + + assert _FAULT not in body + kinds = [k for k, _ in _events(body)] + assert "error" not in kinds + assert "workflow_lifecycle" in kinds + + +async def test_a_shared_task_card_gets_its_status_and_never_the_failure(client): + """A viewer's task card settles from current liveness like the owner's, + but a failed task's reason is exception text and stays off.""" + card = { + "event": "artifact", + "data": { + "artifact_type": "task", + "payload": {"task_id": "tsk1", "status": "running", "workspace_id": "w"}, + }, + } + page = ReplayPage(lines=[line_of(card)], first_turn_index=0, has_more=False) + details = {"tsk1": {"status": "failed", "error": _FAULT, "error_type": "X"}} + thread = { + "conversation_thread_id": _THREAD_ID, + "workspace_id": "ws-1", + "share_permissions": {}, + } + row = {"conversation_thread_id": _THREAD_ID, "latest_checkpoint_id": None} + with ( + patch(_THREAD_BY_TOKEN, new=AsyncMock(return_value=thread)), + patch(_REPLAY_DATA, new=AsyncMock(return_value=replay_rows(row, [], []))), + patch(_READ_PAGE, new=AsyncMock(return_value=(page, "checkpoint"))), + patch(_TASK_DETAILS, new=AsyncMock(return_value=details)), + ): + resp = await client.get(f"/api/v1/public/shared/{_SHARE_TOKEN}/replay") + + (kind, data), _done = _events(resp.text) + assert kind == "artifact" + assert data["payload"] == {"task_id": "tsk1", "status": "failed"} + assert _FAULT not in resp.text diff --git a/tests/unit/server/app/test_public_replay_duration.py b/tests/unit/server/app/test_public_replay_duration.py index 45ea1e2c2..e5419d091 100644 --- a/tests/unit/server/app/test_public_replay_duration.py +++ b/tests/unit/server/app/test_public_replay_duration.py @@ -1,9 +1,8 @@ """The shared replay has to carry the turn's end, or its fold row has no duration. -``build_sse_replay_items`` is not on this path: ``public.py`` hand-builds its own -``user_message`` payload so it can drop the keys a viewer must not read. That -makes it a second construction site for the same event, and the completion stamp -the fold row pairs with the query timestamp has to be mirrored into it. +The public route builds its pages with the owner's assembly and then drops the +keys a viewer must not read from each ``user_message``; the completion stamp the +fold row pairs with the query timestamp has to survive that copy. """ from __future__ import annotations @@ -17,6 +16,7 @@ from httpx import ASGITransport, AsyncClient from tests.conftest import create_test_app +from tests.unit.server.services.history.replay_builders import replay_rows pytestmark = pytest.mark.asyncio @@ -24,8 +24,7 @@ _THREAD_ID = "44444444-4444-4444-8444-444444444444" _THREAD_BY_TOKEN = "src.server.app.share_access.get_thread_by_share_token" -_QUERIES = "src.server.app.public.get_queries_for_thread" -_RESPONSES = "src.server.app.public.get_responses_for_thread" +_REPLAY_DATA = "src.server.app.public.get_replay_thread_data" _TASK_DETAILS = "src.server.services.history.task_status.resolve_task_details" _STARTED = datetime(2026, 9, 21, 15, 4, 0, tzinfo=timezone.utc) @@ -59,8 +58,19 @@ async def _replay(client, response_row: dict) -> list[dict]: ] with ( patch(_THREAD_BY_TOKEN, new=AsyncMock(return_value=thread)), - patch(_QUERIES, new=AsyncMock(return_value=(queries, None))), - patch(_RESPONSES, new=AsyncMock(return_value=([response_row], None))), + patch( + _REPLAY_DATA, + new=AsyncMock( + return_value=replay_rows( + { + "conversation_thread_id": _THREAD_ID, + "latest_checkpoint_id": None, + }, + queries, + [response_row], + ) + ), + ), patch(_TASK_DETAILS, new=AsyncMock(return_value={})), ): resp = await client.get(f"/api/v1/public/shared/{_SHARE_TOKEN}/replay") @@ -99,14 +109,13 @@ def _user_message(events: list[dict]) -> dict: def test_the_public_query_actually_selects_the_settle_instant(): """The mock above can hand the route any shape; production cannot. - The owner's replay reads ``_SETTLED_ATTEMPTS`` directly and gets the column - from its ``*``. The public route goes through ``get_responses_for_thread``, - which projects ``_RESPONSE_COLUMNS``, so the preferred branch is only - reachable while that list carries the column. + The public route reads the light response columns through + ``get_replay_thread_data``, so the preferred branch is only reachable + while that list carries the column. """ from src.server.database.conversation import _sql - assert "usage_settled_at" in _sql._RESPONSE_COLUMNS + assert "usage_settled_at" in _sql._LIGHT_RESPONSE_COLUMNS async def test_settled_turn_carries_its_end(client): diff --git a/tests/unit/server/app/test_public_replay_stopped.py b/tests/unit/server/app/test_public_replay_stopped.py index 337494bbd..4d1fe2e8d 100644 --- a/tests/unit/server/app/test_public_replay_stopped.py +++ b/tests/unit/server/app/test_public_replay_stopped.py @@ -1,7 +1,7 @@ """A shared transcript keeps the Stopped chip of a turn its owner stopped. -``public.py`` streams stored events itself, so the stop close the owner's -replay derives from the response row has to be added on this path too. +The public route shares the owner's replay assembly, so the stop close derived +from the response row has to reach the shared stream too. """ from __future__ import annotations @@ -14,6 +14,7 @@ from httpx import ASGITransport, AsyncClient from tests.conftest import create_test_app +from tests.unit.server.services.history.replay_builders import replay_rows pytestmark = pytest.mark.asyncio @@ -22,8 +23,7 @@ _RESPONSE_ID = "55555555-5555-4555-8555-555555555555" _THREAD_BY_TOKEN = "src.server.app.share_access.get_thread_by_share_token" -_QUERIES = "src.server.app.public.get_queries_for_thread" -_RESPONSES = "src.server.app.public.get_responses_for_thread" +_REPLAY_DATA = "src.server.app.public.get_replay_thread_data" _TASK_DETAILS = "src.server.services.history.task_status.resolve_task_details" _TEXT = { @@ -83,8 +83,19 @@ async def _closes(client, sse_events: list[dict], metadata: dict) -> list[dict]: } with ( patch(_THREAD_BY_TOKEN, new=AsyncMock(return_value=thread)), - patch(_QUERIES, new=AsyncMock(return_value=(queries, None))), - patch(_RESPONSES, new=AsyncMock(return_value=([row], None))), + patch( + _REPLAY_DATA, + new=AsyncMock( + return_value=replay_rows( + { + "conversation_thread_id": _THREAD_ID, + "latest_checkpoint_id": None, + }, + queries, + [row], + ) + ), + ), patch(_TASK_DETAILS, new=AsyncMock(return_value={})), ): resp = await client.get(f"/api/v1/public/shared/{_SHARE_TOKEN}/replay") diff --git a/tests/unit/server/app/test_replay_snapshot_frame.py b/tests/unit/server/app/test_replay_snapshot_frame.py index bab7b6613..0beba8ee9 100644 --- a/tests/unit/server/app/test_replay_snapshot_frame.py +++ b/tests/unit/server/app/test_replay_snapshot_frame.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, patch import pytest +from tests.unit.server.services.history.replay_builders import replay_rows THREADS_MOD = "src.server.app.threads.messaging" @@ -18,7 +19,7 @@ def _replay_data(): """Owner matches create_test_app's override; no checkpoint pointer, so replay takes the (empty) sse path without touching a DB.""" thread = {"conversation_thread_id": THREAD, "latest_checkpoint_id": None} - return ("test-user-123", thread, [], [], {}, {}) + return replay_rows(thread, [], [], owner_id="test-user-123") def _parse(body: str) -> list[tuple[str, str]]: diff --git a/tests/unit/server/contracts/test_status.py b/tests/unit/server/contracts/test_status.py index 60d00bcf4..12d98651a 100644 --- a/tests/unit/server/contracts/test_status.py +++ b/tests/unit/server/contracts/test_status.py @@ -20,7 +20,7 @@ RAW_TERMINAL_SNAPSHOT_STATUSES, TERMINAL_PUBLIC_STATUSES, TERMINAL_STATUSES, - classify_interrupt_reason, + classify_interrupts, is_live, is_terminal, to_public, @@ -122,6 +122,10 @@ def _pause(*requests): return [Interrupt(value={"action_requests": list(requests)})] +def _reason(interrupts): + return classify_interrupts(interrupts)[0] + + class TestInterruptReason: """The reason column is the run ledger's classification of a pause. @@ -133,35 +137,35 @@ class TestInterruptReason: def test_credit_pause_keeps_its_exact_spelling(self): # Load-bearing, not cosmetic: the subagent resume query selects on this # literal, so a rename here silently strands stopped tasks. - assert classify_interrupt_reason( + assert _reason( _pause({"type": "credit_pause", "message": "out of credits"}) ) == "credit_pause" def test_a_question_reads_as_a_question(self): - assert classify_interrupt_reason( + assert _reason( _pause({"type": "ask_user_question", "question": "which ticker?"}) ) == "user_question" def test_a_proposal_is_an_approval(self): for kind in ("create_workspace", "delete_workspace", "delete_thread"): - assert classify_interrupt_reason( + assert _reason( _pause({"type": kind, "workspace_id": "w-1"}) ) == "approval_required" def test_an_unknown_action_generalizes_instead_of_guessing(self): - assert classify_interrupt_reason( + assert _reason( _pause({"type": "some_future_action", "detail": "x"}) ) == "approval_required" def test_an_unreadable_payload_is_unclassified_rather_than_labelled(self): - assert classify_interrupt_reason(_pause({"detail": "no discriminator"})) is None - assert classify_interrupt_reason([Interrupt(value="not a dict")]) is None - assert classify_interrupt_reason([]) is None + assert _reason(_pause({"detail": "no discriminator"})) is None + assert _reason([Interrupt(value="not a dict")]) is None + assert _reason([]) is None def test_credit_pause_outranks_a_proposal_buffered_ahead_of_it(self): # A pause can carry several payloads; the one with behaviour attached # has to win regardless of the order they were buffered in. - assert classify_interrupt_reason( + assert _reason( _pause({"type": "create_workspace", "workspace_name": "scratch"}) + _pause({"type": "credit_pause", "message": "out of credits"}) ) == "credit_pause" @@ -171,7 +175,7 @@ def test_a_payload_is_read_whole(self): # payload, so precedence has to hold within a payload and not just # across them: reading only the leading request would make the answer # depend on which request happened to come first. - assert classify_interrupt_reason( + assert _reason( _pause({"name": "run_backtest", "args": {}}, {"type": "ask_user_question", "question": "which ticker?"}) ) == "user_question" @@ -203,7 +207,7 @@ def test_the_classifier_only_ever_returns_a_member(self): {"type": "some_future_action"}, {"name": "any_approved_tool", "args": {}}, ): - assert classify_interrupt_reason(_pause(request)) in INTERRUPT_REASONS + assert _reason(_pause(request)) in INTERRUPT_REASONS def test_the_resume_query_binds_the_constant_not_a_literal(self): from src.server.database.runs import credit_ledger diff --git a/tests/unit/server/database/test_conversation_db.py b/tests/unit/server/database/test_conversation_db.py index c2bb81cbd..b92bd6d7a 100644 --- a/tests/unit/server/database/test_conversation_db.py +++ b/tests/unit/server/database/test_conversation_db.py @@ -773,7 +773,7 @@ async def test_replay_data_excludes_subagent_usage_rows( "latest_checkpoint_id": "cp-1", }, ] - mock_cursor.fetchall.side_effect = [[], [], [], []] + mock_cursor.fetchall.side_effect = [[], [], [], [], []] await get_replay_thread_data(thread_id) @@ -790,7 +790,7 @@ async def test_get_replay_thread_data_malformed_id_skips_database( result = await get_replay_thread_data("results") - assert result == (None, None, [], [], [], []) + assert result is None mock_cursor.execute.assert_not_awaited() diff --git a/tests/unit/server/database/test_provenance_insert.py b/tests/unit/server/database/test_provenance_insert.py index 03124b736..d667e7c54 100644 --- a/tests/unit/server/database/test_provenance_insert.py +++ b/tests/unit/server/database/test_provenance_insert.py @@ -4,7 +4,7 @@ mocked psycopg3 cursor (the ``mock_connection`` / ``mock_cursor`` fixtures) and asserts: -* the write is delete-then-insert keyed by conversation_response_id, batched via +* the write is delete-then-insert keyed by conversation_response_id and lane, batched via a single multi-row INSERT, guarded by a per-response advisory lock inside a transaction (savepoint) so a failed write can't poison the turn-persist commit, * every TEXT bind is NUL-stripped and JSONB ``args_fingerprint`` wraps in SafeJson, @@ -19,6 +19,7 @@ from src.server.database.provenance import ( _INSERT_COLUMNS, + MAIN_PROVENANCE_LANE, insert_provenance_records, sync_provenance_for_response, ) @@ -66,6 +67,7 @@ async def test_lock_delete_then_batched_insert(self, mock_connection, mock_curso conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[_record(), _record(identifier="https://example.test/b")], ) assert n == 2 @@ -73,10 +75,10 @@ async def test_lock_delete_then_batched_insert(self, mock_connection, mock_curso sqls = _execute_sqls(mock_cursor) # Advisory lock first (serializes concurrent drains for this response). assert any("pg_advisory_xact_lock" in s for s in sqls) - # Then the DELETE keyed by response_id. + # Then the DELETE keyed by response_id and the lane written. delete = next(c for c in mock_cursor.execute.call_args_list if "DELETE FROM provenance_records" in c.args[0]) - assert delete.args[1] == (RESPONSE_ID,) + assert delete.args[1] == (RESPONSE_ID, [MAIN_PROVENANCE_LANE]) # Rows go in one multi-row INSERT, not per-row executemany. mock_cursor.executemany.assert_not_awaited() @@ -96,6 +98,7 @@ async def test_detail_bound_at_its_column(self, mock_connection, mock_cursor): conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[_record(detail="daily_prices")], ) _, params = _insert_call(mock_cursor).args @@ -108,6 +111,7 @@ async def test_empty_records_only_deletes(self, mock_connection, mock_cursor): conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[], ) assert n == 0 @@ -138,6 +142,7 @@ async def test_nul_bytes_stripped_from_text_binds( conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[dirty], ) _, params = _insert_call(mock_cursor).args @@ -154,6 +159,7 @@ async def test_args_fingerprint_wrapped_in_safejson( conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[_record(args_fingerprint={"sha256": "te\x00st"})], ) _, params = _insert_call(mock_cursor).args @@ -171,6 +177,7 @@ async def test_null_args_fingerprint_bound_as_none( conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[_record(args_fingerprint=None)], ) _, params = _insert_call(mock_cursor).args @@ -184,6 +191,7 @@ async def test_args_wrapped_in_safejson(self, mock_connection, mock_cursor): conversation_response_id=RESPONSE_ID, conversation_thread_id=THREAD_ID, turn_index=0, + lanes=(MAIN_PROVENANCE_LANE,), records=[_record(args_fingerprint=None, args={"symbol": "AA\x00PL"})], ) _, params = _insert_call(mock_cursor).args @@ -209,6 +217,7 @@ async def test_returns_inserted_count(self, mock_connection, mock_cursor): conversation_thread_id=THREAD_ID, turn_index=0, sse_events=events, + lanes=(MAIN_PROVENANCE_LANE,), ) assert n == 1 assert any("INSERT INTO provenance_records" in s @@ -236,5 +245,6 @@ async def test_swallows_insert_failure_returns_zero(self, mock_connection): conversation_thread_id=THREAD_ID, turn_index=0, sse_events=events, + lanes=(MAIN_PROVENANCE_LANE,), ) assert n == 0 # swallowed, never re-raised diff --git a/tests/unit/server/database/test_provenance_lane_scope.py b/tests/unit/server/database/test_provenance_lane_scope.py new file mode 100644 index 000000000..035f8f226 --- /dev/null +++ b/tests/unit/server/database/test_provenance_lane_scope.py @@ -0,0 +1,152 @@ +"""Lane-scoped provenance rewrites. + +A response's provenance rows are rewritten per lane: the finalize replaces the +main lane from its own record, a subagent archive replaces only the lanes it +holds. The delete is therefore scoped to the given lanes plus every inserted +record's own agent, so a write always replaces what it inserts and never +touches a lane it does not name. +""" + +import pytest + +from src.server.database.provenance import ( + MAIN_PROVENANCE_LANE, + insert_provenance_records, + sync_provenance_for_response, +) + +RESPONSE_ID = "resp-1" +THREAD_ID = "thread-1" + +_SCOPED_DELETE = ( + "DELETE FROM provenance_records " + "WHERE conversation_response_id = %s " + "AND COALESCE(agent, '') = ANY(%s)" +) + + +def _record(agent, identifier="https://example.test/a"): + return { + "source_type": "web_search", + "identifier": identifier, + "result_sha256": f"sha-{identifier}", + "agent": agent, + } + + +def _statements(mock_cursor): + return [" ".join(c.args[0].split()) for c in mock_cursor.execute.call_args_list] + + +def _delete_call(mock_cursor): + return next( + c + for c in mock_cursor.execute.call_args_list + if "DELETE FROM provenance_records" in c.args[0] + ) + + +async def _insert(conn, records, lanes): + return await insert_provenance_records( + conn, + conversation_response_id=RESPONSE_ID, + conversation_thread_id=THREAD_ID, + turn_index=0, + records=records, + lanes=lanes, + ) + + +@pytest.mark.asyncio +async def test_the_main_lane_replaces_only_its_own_rows(mock_connection, mock_cursor): + n = await _insert( + mock_connection, + [_record("main"), _record("main", "https://example.test/b")], + (MAIN_PROVENANCE_LANE,), + ) + assert n == 2 + + statements = _statements(mock_cursor) + assert [s.split()[0] for s in statements] == ["SELECT", "DELETE", "INSERT"] + assert "pg_advisory_xact_lock" in statements[0] + sql, params = _delete_call(mock_cursor).args + assert " ".join(sql.split()) == _SCOPED_DELETE + assert params == (RESPONSE_ID, ["main"]) + + +@pytest.mark.asyncio +async def test_every_records_own_agent_joins_the_scope(mock_connection, mock_cursor): + await _insert( + mock_connection, + [ + _record("task:k7"), + _record(None, "https://example.test/b"), + _record("main", "https://example.test/c"), + ], + (MAIN_PROVENANCE_LANE,), + ) + _, params = _delete_call(mock_cursor).args + # An agentless record is lane "", which COALESCE(agent, '') matches. + assert params == (RESPONSE_ID, ["", "main", "task:k7"]) + + +@pytest.mark.asyncio +async def test_a_record_without_an_agent_key_is_lane_empty( + mock_connection, mock_cursor +): + record = _record(None) + del record["agent"] + await _insert(mock_connection, [record], ()) + _, params = _delete_call(mock_cursor).args + assert params == (RESPONSE_ID, [""]) + + +@pytest.mark.asyncio +async def test_the_scope_is_the_sorted_union_whatever_the_lanes_order( + mock_connection, mock_cursor +): + await _insert(mock_connection, [_record("task:a")], {"task:c", "task:b", "task:a"}) + _, params = _delete_call(mock_cursor).args + assert params == (RESPONSE_ID, ["task:a", "task:b", "task:c"]) + + +@pytest.mark.asyncio +async def test_an_empty_lane_still_clears_its_rows(mock_connection, mock_cursor): + """A lane rewritten with no records (an archive that emptied it) loses + the rows it had, and only those.""" + n = await _insert(mock_connection, [], (MAIN_PROVENANCE_LANE,)) + assert n == 0 + + sql, params = _delete_call(mock_cursor).args + assert " ".join(sql.split()) == _SCOPED_DELETE + assert params == (RESPONSE_ID, ["main"]) + assert not any(s.startswith("INSERT") for s in _statements(mock_cursor)) + + +@pytest.mark.asyncio +async def test_the_sync_entry_point_passes_its_lanes_through( + mock_connection, mock_cursor +): + events = [ + { + "event": "provenance", + "data": { + "source_type": "web_search", + "identifier": "https://example.test/a", + "result_sha256": "sha-a", + "agent": "task:k7", + }, + } + ] + n = await sync_provenance_for_response( + mock_connection, + conversation_response_id=RESPONSE_ID, + conversation_thread_id=THREAD_ID, + turn_index=0, + sse_events=events, + strict=True, + lanes=("task:k7",), + ) + assert n == 1 + _, params = _delete_call(mock_cursor).args + assert params == (RESPONSE_ID, ["task:k7"]) diff --git a/tests/unit/server/database/test_provenance_sync_skip.py b/tests/unit/server/database/test_provenance_sync_skip.py index 0ee86cb08..bb38733f7 100644 --- a/tests/unit/server/database/test_provenance_sync_skip.py +++ b/tests/unit/server/database/test_provenance_sync_skip.py @@ -28,6 +28,7 @@ async def test_skips_when_no_provenance_event(): conversation_thread_id="t", turn_index=0, sse_events=events, + lanes=(), ) inner.assert_not_awaited() # no extract + delete-then-insert on context_window @@ -42,6 +43,7 @@ async def test_skips_on_empty_or_none(): conversation_thread_id="t", turn_index=0, sse_events=events, + lanes=(), ) inner.assert_not_awaited() @@ -64,5 +66,6 @@ async def test_runs_when_provenance_event_present(): conversation_thread_id="t", turn_index=0, sse_events=events, + lanes=(), ) inner.assert_awaited_once() diff --git a/tests/unit/server/database/test_start_run_fork_order_refusal.py b/tests/unit/server/database/test_start_run_fork_order_refusal.py index 94f8e0e77..669de8574 100644 --- a/tests/unit/server/database/test_start_run_fork_order_refusal.py +++ b/tests/unit/server/database/test_start_run_fork_order_refusal.py @@ -25,6 +25,7 @@ OA_DB = "src.server.database.runs.lifecycle.oa_db" QR_DB = "src.server.database.runs.lifecycle.qr_db" SR_DB = "src.server.database.runs.lifecycle.sr_db" +REPAIR = "src.server.database.runs.lifecycle.subagent_repair" THREAD_ID = "11111111-1111-1111-1111-111111111111" RUN_ID = "22222222-2222-2222-2222-222222222222" @@ -79,7 +80,7 @@ async def _truncate(*args, **kwargs): patch(f"{SR_DB}.find_open_run_from_turn", new=_task_guard), patch(f"{OA_DB}.refuse_attempts_on_fork", new=refuse), patch(f"{QR_DB}.truncate_thread_from_turn", new=_truncate), - patch(f"{SR_DB}.repair_task_chains", new=AsyncMock()), + patch(f"{REPAIR}.repair_task_chains", new=AsyncMock()), # Aborts the START past the truncation, inside the same transaction. patch(f"{QR_DB}.update_thread_checkpoint_id", new=AsyncMock(return_value=False)), ): diff --git a/tests/unit/server/database/test_start_run_fork_task_guard.py b/tests/unit/server/database/test_start_run_fork_task_guard.py index ef3058025..45f109adc 100644 --- a/tests/unit/server/database/test_start_run_fork_task_guard.py +++ b/tests/unit/server/database/test_start_run_fork_task_guard.py @@ -22,6 +22,7 @@ QR_DB = "src.server.database.runs.lifecycle.qr_db" SR_DB = "src.server.database.runs.lifecycle.sr_db" +REPAIR = "src.server.database.runs.lifecycle.subagent_repair" THREAD_ID = "11111111-1111-1111-1111-111111111111" RUN_ID = "22222222-2222-2222-2222-222222222222" @@ -78,7 +79,7 @@ async def test_guard_is_scoped_to_the_truncated_turns(mock_connection, mock_curs with ( patch(f"{SR_DB}.find_open_run_from_turn", new=probe), patch(f"{QR_DB}.truncate_thread_from_turn", new=AsyncMock(return_value=3)), - patch(f"{SR_DB}.repair_task_chains", new=AsyncMock()), + patch(f"{REPAIR}.repair_task_chains", new=AsyncMock()), patch(f"{QR_DB}.update_thread_checkpoint_id", new=AsyncMock(return_value=False)), ): # Fails at the checkpoint pin, well past the guard. @@ -110,7 +111,7 @@ async def _pin(*args, **kwargs): with ( patch(f"{SR_DB}.find_open_run_from_turn", new=AsyncMock(return_value=None)), patch(f"{QR_DB}.truncate_thread_from_turn", new=_truncate), - patch(f"{SR_DB}.repair_task_chains", new=_repair), + patch(f"{REPAIR}.repair_task_chains", new=_repair), patch(f"{QR_DB}.update_thread_checkpoint_id", new=_pin), ): with pytest.raises(TurnLifecycleError): @@ -129,7 +130,7 @@ async def test_repair_runs_on_the_forks_own_session(mock_connection, mock_cursor with ( patch(f"{SR_DB}.find_open_run_from_turn", new=AsyncMock(return_value=None)), patch(f"{QR_DB}.truncate_thread_from_turn", new=AsyncMock(return_value=1)), - patch(f"{SR_DB}.repair_task_chains", new=repair), + patch(f"{REPAIR}.repair_task_chains", new=repair), patch(f"{QR_DB}.update_thread_checkpoint_id", new=AsyncMock(return_value=False)), ): with pytest.raises(TurnLifecycleError): diff --git a/tests/unit/server/database/test_subagent_runs_repair.py b/tests/unit/server/database/test_subagent_runs_repair.py index 7abb970b0..448d8ba28 100644 --- a/tests/unit/server/database/test_subagent_runs_repair.py +++ b/tests/unit/server/database/test_subagent_runs_repair.py @@ -23,11 +23,13 @@ import pytest +from src.server.database.runs.subagent_repair import ( + repair_dangling_task_chains, + repair_task_chains, +) from src.server.database.runs.subagent_runs import ( count_open_runs_for_thread, find_open_run_from_turn, - repair_dangling_task_chains, - repair_task_chains, ) THREAD_ID = "11111111-1111-1111-1111-111111111111" @@ -148,9 +150,10 @@ async def test_repair_runs_on_the_pinned_session(mock_connection, mock_cursor): _rowcounts(mock_cursor, [1, 1]) @asynccontextmanager - async def _explode(): # pragma: no cover - must never be entered - raise AssertionError("repair_task_chains took a pool connection") - yield + async def _explode(conn=None): + if conn is None: + raise AssertionError("repair_task_chains took a pool connection") + yield conn with patch(POOL, new=_explode): result = await repair_task_chains(THREAD_ID, conn=mock_connection) diff --git a/tests/unit/server/handlers/test_checkpoint_handler.py b/tests/unit/server/handlers/test_checkpoint_handler.py index c13528f54..13a489044 100644 --- a/tests/unit/server/handlers/test_checkpoint_handler.py +++ b/tests/unit/server/handlers/test_checkpoint_handler.py @@ -8,7 +8,7 @@ """ from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest from fastapi import HTTPException diff --git a/tests/unit/server/handlers/test_streaming_handler.py b/tests/unit/server/handlers/test_streaming_handler.py index 961d59a82..503244cc4 100644 --- a/tests/unit/server/handlers/test_streaming_handler.py +++ b/tests/unit/server/handlers/test_streaming_handler.py @@ -1,5 +1,6 @@ """ -Tests for src/server/services/runs/sse_producer.py +Tests for the run SSE producer (src/server/services/runs/sse_producer.py), its +event archive (event_archive.py) and its error classification (stream_errors.py). Covers: - StreamEventAccumulator: accumulation, merging, max buffer size @@ -21,7 +22,7 @@ class TestStreamEventAccumulator: """Tests for the StreamEventAccumulator class.""" def _make_accumulator(self, max_bytes=16 * 1024): - from src.server.services.runs.sse_producer import StreamEventAccumulator + from src.server.services.runs.event_archive import StreamEventAccumulator return StreamEventAccumulator(max_merged_bytes=max_bytes) @@ -471,7 +472,7 @@ def test_format_error_event_upstream_wrapped_by_internal(self): def test_classify_stream_exception_parses_status_from_message(self): """When status_code isn't on the exception, parse it from the message.""" - from src.server.services.runs.sse_producer import classify_stream_exception + from src.server.services.runs.stream_errors import classify_stream_exception import httpx exc = httpx.HTTPError("got HTTP 429 from upstream, backing off") diff --git a/tests/unit/server/handlers/test_trigger_compaction.py b/tests/unit/server/handlers/test_trigger_compaction.py index c293a7ea3..b18994515 100644 --- a/tests/unit/server/handlers/test_trigger_compaction.py +++ b/tests/unit/server/handlers/test_trigger_compaction.py @@ -101,13 +101,12 @@ def _stub_resolve_graph_and_state(): state.values = {"_summarization_event": None} messages = [MagicMock(id="m1"), MagicMock(id="m2")] backend = None - lg_config = {"configurable": {"thread_id": "thread-1"}} async def _stub(thread_id, verb, config=None, checkpointer=None, user_id=None, held=None): _stub.captured_config = config _stub.captured_checkpointer = checkpointer _stub.captured_user_id = user_id - return graph, lg_config, state, messages, "ws-1", backend + return graph, state, messages, "ws-1", backend _stub.captured_config = None _stub.captured_checkpointer = None diff --git a/tests/unit/server/services/history/reader_graphs.py b/tests/unit/server/services/history/reader_graphs.py new file mode 100644 index 000000000..2a04266c3 --- /dev/null +++ b/tests/unit/server/services/history/reader_graphs.py @@ -0,0 +1,68 @@ +"""Graphs and reads shared by the checkpoint-reader suites: a real +DeltaAgentState graph writes the turns, and ``_history`` reads them back the +way replay does.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.graph import START, StateGraph + +from ptc_agent.agent.state import DeltaAgentState + +THREAD = "thread-1" + + +def _echo_graph(checkpointer): + """One turn = reply 'echo: ' with a stable per-turn id.""" + + def agent(state): + humans = [m for m in state["messages"] if isinstance(m, HumanMessage)] + last = humans[-1].content if humans else "?" + return { + "messages": [AIMessage(content=f"echo: {last}", id=f"ai-{len(state['messages'])}")] + } + + return ( + StateGraph(DeltaAgentState) + .add_node("agent", agent) + .add_edge(START, "agent") + .compile(checkpointer=checkpointer) + ) + + +def _cfg(thread_id=THREAD, *, turn_index=None, run_id=None, checkpoint_id=None): + cfg = {"configurable": {"thread_id": thread_id}} + if checkpoint_id: + cfg["configurable"]["checkpoint_id"] = checkpoint_id + metadata = {} + if run_id is not None: + metadata["run_id"] = run_id + if turn_index is not None: + metadata["turn_index"] = turn_index + if metadata: + cfg["metadata"] = metadata + return cfg + + +async def _history(reader, thread_id=THREAD, tip=None): + """Every turn on the branch, the interrupts pending at its tip and the + tip's ui records, read the way replay reads them.""" + anchors, tip_id = await reader.aget_turn_anchors(thread_id, tip) + if not anchors: + return SimpleNamespace(turns=[], interrupts=[], ui=[]) + tip_state = await reader.aget_state(thread_id, tip_id) + return SimpleNamespace( + turns=await reader.aget_turn_slices(thread_id, anchors), + interrupts=await reader.aget_tip_interrupts(thread_id, tip_id), + ui=list(tip_state.values.get("ui") or []), + ) + + +async def _run_turns(graph, n, start=0): + for i in range(start, start + n): + await graph.ainvoke( + {"messages": [HumanMessage(content=f"q{i}", id=f"h-{i}")]}, + _cfg(turn_index=i, run_id=f"run-{i}"), + ) diff --git a/tests/unit/server/services/history/replay_builders.py b/tests/unit/server/services/history/replay_builders.py index 606bd3269..67d2db4c8 100644 --- a/tests/unit/server/services/history/replay_builders.py +++ b/tests/unit/server/services/history/replay_builders.py @@ -1,24 +1,44 @@ -"""Builders shared by the replay suites: checkpoint turns, stored rows, a mocked reader.""" +"""Builders shared by the replay suites: checkpoint turns, stored rows, a mocked +reader, and an in-memory slice store.""" from __future__ import annotations +from dataclasses import dataclass, field, replace +from typing import Any from unittest.mock import AsyncMock, MagicMock -from langchain_core.messages import HumanMessage +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from src.server.services.history.reader import TaskHistory, TurnSlice +from ptc_agent.agent.transcript.classify import is_run_boundary_message +from src.server.database.conversation.replay_rows import ThreadRows +from src.server.database.conversation.turn_slices import RunRow, StoredSlice +from src.server.services.history.reader import TaskHistory, TaskRun, TurnAnchor +from src.server.services.history.slices import SpanDelta, TurnSlice THREAD = "thread-r" +TIP = "cp-tip" -def _turn(ordinal, messages, user="hello", turn_index=None, new_ui_records=None): +@dataclass +class ThreadHistory: + """A thread's turns and the interrupts pending at its tip: what + ``_mock_reader`` serves through the anchor and slice reads.""" + + thread_id: str + turns: list[TurnSlice] = field(default_factory=list) + interrupts: list[dict[str, Any]] = field(default_factory=list) + + +def _turn(ordinal, messages, turn_index=None, new_ui_records=None): return TurnSlice( - turn_ordinal=ordinal, - input_checkpoint_id=f"cp-in-{ordinal}", - end_checkpoint_id=f"cp-end-{ordinal}", - user_message=HumanMessage(content=user, id=f"h-{ordinal}"), + anchor=TurnAnchor( + turn_ordinal=ordinal, + input_checkpoint_id=f"cp-in-{ordinal}", + tail_checkpoint_id=None, + end_checkpoint_id=f"cp-end-{ordinal}", + turn_index=turn_index, + ), messages=messages, - turn_index=turn_index, new_ui_records=new_ui_records or [], ) @@ -35,12 +55,282 @@ def _response(turn_index, sse_events=None, status="completed"): } +def thread_rows( + queries, + responses_by_turn, + *, + usages=None, + provenance=None, + run_facts=None, + thread_id=THREAD, + tip=TIP, +) -> ThreadRows: + """A thread's replay rows, as ``get_replay_thread_data`` reads them.""" + return ThreadRows( + thread_id=thread_id, + owner_id="user-1", + thread={"conversation_thread_id": thread_id, "latest_checkpoint_id": tip}, + queries=list(queries), + responses_by_turn=dict(responses_by_turn), + usages=list(usages or []), + provenance=list(provenance or []), + run_facts=list(run_facts or []), + ) + + +def replay_rows( + thread, queries, responses, *, owner_id="owner-1", usages=(), provenance=() +) -> ThreadRows: + """``get_replay_thread_data``'s result for a thread row and the row lists + an endpoint test hands it.""" + return ThreadRows( + thread_id=str(thread["conversation_thread_id"]), + owner_id=owner_id, + thread=dict(thread), + queries=list(queries), + responses_by_turn={r.get("turn_index"): r for r in responses}, + usages=list(usages), + provenance=list(provenance), + run_facts=[], + ) + + +class SliceStore: + """In-memory stand-in for the ``turn_slices`` / ``task_run_slices`` tables.""" + + def __init__(self) -> None: + self.turns: dict[str, dict[str, Any]] = {} + self.runs: dict[tuple[str, str], dict[str, Any]] = {} + # Tail checkpoint of every turn whose projected lines were stored. + self.lined: list[str] = [] + # Rows per write statement, in order. + self.statements: list[tuple[str, int]] = [] + # Response id of every turn row written, once per write. + self.written: list[str] = [] + + @staticmethod + def _stored(row, *, data=True): + return StoredSlice( + input_checkpoint_id=row["input_checkpoint_id"], + tail_checkpoint_id=row["tail_checkpoint_id"], + slice_key=row["slice_key"], + codec=row["slice_codec"] if data else None, + data=row["slice"] if data else None, + lines_key=row.get("lines_key"), + lines=row.get("lines"), + claims=row.get("run_claims"), + ) + + async def get_turn_lines(self, response_ids): + return { + r: self._stored(self.turns[r], data=False) + for r in response_ids + if r in self.turns + } + + async def get_turn_slices(self, response_ids): + return {r: self._stored(self.turns[r]) for r in response_ids if r in self.turns} + + async def has_turn_slices(self, thread_id, slice_key): + return any(row["slice_key"] == slice_key for row in self.turns.values()) + + async def get_slices_at(self, thread_id, slice_key, input_checkpoint_ids): + return [ + self._stored(row) + for row in self.turns.values() + if row["slice_key"] == slice_key + and row["input_checkpoint_id"] in input_checkpoint_ids + ] + + async def upsert_turn_slices(self, rows): + self.statements.append(("upsert_turn_slices", len(rows))) + for row in rows: + stored = row.slice + self.written.append(row.response_id) + self.turns[row.response_id] = { + "response_id": row.response_id, + "input_checkpoint_id": stored.input_checkpoint_id, + "tail_checkpoint_id": stored.tail_checkpoint_id, + "slice_key": stored.slice_key, + "slice_codec": stored.codec, + "slice": stored.data, + "lines_key": stored.lines_key, + "lines": stored.lines, + "run_claims": stored.claims, + } + if stored.lines is not None: + self.lined.append(stored.tail_checkpoint_id) + + async def update_turns(self, rows): + self.statements.append(("update_turns", len(rows))) + for update in rows: + row = self.turns.get(update.response_id) + if not row or row["tail_checkpoint_id"] != update.tail_checkpoint_id: + continue + self.written.append(update.response_id) + for column, value in ( + ("lines_key", update.lines_key), + ("lines", update.lines), + ("run_claims", update.claims), + ): + if value is not None: + row[column] = value + if update.lines is not None: + self.lined.append(update.tail_checkpoint_id) + + async def get_run_slices(self, thread_id, task_ids): + return [ + RunRow(thread_id, r["task_id"], r["task_run_id"], self._stored(r)) + for (t, _), r in self.runs.items() + if t in task_ids + ] + + async def upsert_run_slices(self, rows): + self.statements.append(("upsert_run_slices", len(rows))) + for row in rows: + stored = row.slice + self.runs[(row.task_id, stored.input_checkpoint_id)] = { + "task_id": row.task_id, + "input_checkpoint_id": stored.input_checkpoint_id, + "task_run_id": row.task_run_id, + "tail_checkpoint_id": stored.tail_checkpoint_id, + "slice_key": stored.slice_key, + "slice_codec": stored.codec, + "slice": stored.data, + } + + +def _slice_store(monkeypatch) -> SliceStore: + """The test's slice store, installed over the DB module on first use.""" + store = getattr(monkeypatch, "_replay_slice_store", None) + if store is None: + from src.server.database.conversation import turn_slices as slices_db + + store = SliceStore() + for name in ( + "get_turn_lines", + "get_turn_slices", + "upsert_turn_slices", + "update_turns", + "get_run_slices", + "upsert_run_slices", + ): + monkeypatch.setattr(slices_db, name, getattr(store, name)) + monkeypatch._replay_slice_store = store + return store + + +def _reset_slice_store(monkeypatch) -> SliceStore: + """A fresh store, for a test that replays an unrelated second thread.""" + monkeypatch._replay_slice_store = None + return _slice_store(monkeypatch) + + +def _split_runs(messages: list[Any]) -> list[list[Any]]: + """A namespace's transcript cut at each run's input message. A leading + slice with no input attaches to the first run.""" + segments: list[list[Any]] = [] + current: list[Any] = [] + saw_boundary = False + for message in messages: + if is_run_boundary_message(message): + if saw_boundary: + segments.append(current) + current = [] + saw_boundary = True + current.append(message) + if current: + segments.append(current) + return segments + + def _mock_reader(monkeypatch, history, task_messages=None, task_history=None): + """A reader serving ``history`` through the anchor/slice API. + + Turns come from ``thread_history`` and a task's runs from + ``aget_task_history`` (cut at run inputs, stamped from + ``task_run_stamps`` when one stamp per run), each read when the replay + asks, so a test may replace any of these mocks afterwards. + Namespace-wide signals ride the task's first run. + """ + _slice_store(monkeypatch) reader = MagicMock() - reader.aget_thread_history = AsyncMock(return_value=history) + reader.serde = JsonPlusSerializer() + reader.thread_history = AsyncMock(return_value=history) reader.aget_task_history = AsyncMock( return_value=task_history or TaskHistory(messages=task_messages or []) ) + reader.task_run_stamps = AsyncMock(return_value=[]) + seen: dict[str, Any] = {"turns": {}, "interrupts": [], "tasks": {}} + + async def turn_anchors(thread_id, branch_tip_checkpoint_id=None): + current = await reader.thread_history(thread_id, branch_tip_checkpoint_id) + seen["interrupts"] = list(current.interrupts) + anchors = [] + for turn in current.turns: + anchor = turn.anchor + seen["turns"][anchor.input_checkpoint_id] = turn + anchors.append( + replace( + anchor, + tail_checkpoint_id=anchor.tail_checkpoint_id + or f"cp-tail-{anchor.turn_ordinal}", + ending_interrupts=list(anchor.ending_interrupts), + ) + ) + return anchors, TIP + + async def turn_slices(thread_id, anchors): + # Placed on the anchor asked for, as the reader places a slice. + return [replace(seen["turns"][a.input_checkpoint_id], anchor=a) for a in anchors] + + async def tip_interrupts(thread_id, tip_checkpoint_id): + return seen["interrupts"] + + async def task_runs(thread_id, task_id): + namespace = await reader.aget_task_history(thread_id, task_id) + segments = _split_runs(namespace.messages) + stamps = list(await reader.task_run_stamps(thread_id, task_id) or []) + if len(stamps) != len(segments): + stamps = [None] * len(segments) + seen["tasks"][task_id] = (namespace, segments) + return [ + TaskRun( + task_id=task_id, + ordinal=k, + input_checkpoint_id=f"{task_id}:in-{k}", + end_checkpoint_id=f"{task_id}:end-{k}", + tail_checkpoint_id=f"{task_id}:tail-{k}", + task_run_id=stamps[k], + ) + for k in range(len(segments)) + ] + + async def run_slices(thread_id, task_id, runs): + namespace, segments = seen["tasks"][task_id] + out = [] + for run in runs: + first = run.ordinal == 0 + out.append( + SpanDelta( + messages=list(segments[run.ordinal]), + new_summarization_event=namespace.new_summarization_event + if first + else None, + newly_offloaded_args=namespace.newly_offloaded_args if first else 0, + newly_offloaded_reads=namespace.newly_offloaded_reads + if first + else 0, + new_ui_records=list(namespace.new_ui_records) if first else [], + ) + ) + return out + + reader.aget_turn_anchors = AsyncMock(side_effect=turn_anchors) + reader.aget_turn_slices = AsyncMock(side_effect=turn_slices) + reader.aget_tip_interrupts = AsyncMock(side_effect=tip_interrupts) + reader.aget_task_runs = AsyncMock(side_effect=task_runs) + reader.aget_run_slices = AsyncMock(side_effect=run_slices) monkeypatch.setattr( "src.server.services.history.replay.CheckpointHistoryReader.get_instance", lambda: reader, @@ -49,28 +339,17 @@ def _mock_reader(monkeypatch, history, task_messages=None, task_history=None): def _cache_probe(monkeypatch): - """Absorb cache writes; return the list of cached tail checkpoint ids.""" - from src.server.services.history import replay as replay_module - from src.server.services.history.replay import task_lane as task_lane_module + """Seal every unledgered task by its stream; return the list of tail + checkpoint ids whose projected lines were stored.""" + from src.server.services.history import task_streams + from src.server.services.history.replay import run_lane async def fake_details(thread_id, task_ids): return {} - async def fake_live(thread_id, task_ids): - return set() + async def fake_streams(thread_id, task_ids): + return dict.fromkeys(task_ids, task_streams.STREAM_SEALED) - cached: list[str] = [] - - async def fake_store(thread_id, tail_checkpoint_id, fingerprint, items): - cached.append(tail_checkpoint_id) - - async def fake_delete(thread_id, turn_keys): - pass - - monkeypatch.setattr(task_lane_module, "resolve_task_details", fake_details) - monkeypatch.setattr( - replay_module.projection_cache, "live_task_streams", fake_live - ) - monkeypatch.setattr(replay_module.projection_cache, "store_turn", fake_store) - monkeypatch.setattr(replay_module.projection_cache, "delete_turns", fake_delete) - return cached + monkeypatch.setattr(run_lane, "resolve_task_details", fake_details) + monkeypatch.setattr(task_streams, "task_stream_states", fake_streams) + return _slice_store(monkeypatch).lined diff --git a/tests/unit/server/services/history/test_committed.py b/tests/unit/server/services/history/test_committed.py new file mode 100644 index 000000000..f91bfbc91 --- /dev/null +++ b/tests/unit/server/services/history/test_committed.py @@ -0,0 +1,94 @@ +"""A turn's committed text is read from its stored slice only while replay +would serve that slice: cut on the branch the thread now ends in, by the code +running now. Any other turn is projected first, as its settle would have.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + +from src.server.database.conversation.turn_slices import StoredSlice +from src.server.services.history import committed, slices +from src.server.services.history.reader import TurnAnchor + +pytestmark = pytest.mark.asyncio + +_SERDE = JsonPlusSerializer() +_ANCHOR = TurnAnchor( + turn_ordinal=0, input_checkpoint_id="in-1", tail_checkpoint_id="tail-1", turn_index=1 +) +_REFRESH = "src.server.services.history.replay.refresh.refresh_thread" + + +def _row(text: str, *, tail: str = "tail-1") -> StoredSlice: + codec, data = _SERDE.dumps_typed( + {"messages": [HumanMessage(content="q", id="h1"), AIMessage(content=text, id="a1")]} + ) + return StoredSlice("in-1", tail, slices.SLICE_KEY, codec, data) + + +class _Reader: + serde = _SERDE + + def __init__(self, anchors): + self.aget_turn_anchors = AsyncMock(return_value=(anchors, "tip")) + + +def _branch(stored, *, anchors=(_ANCHOR,), persisted=(1,)): + """The thread's pointer, its branch, and its stored slices (a list is + one read after another).""" + reads = stored if isinstance(stored, list) else [stored] + return ( + patch.object( + committed, + "get_branch_pointer", + AsyncMock(return_value=("tip", list(persisted))), + ), + patch.object( + committed.CheckpointHistoryReader, + "get_instance", + return_value=_Reader(list(anchors)), + ), + patch.object(committed.slices_db, "get_turn_slices", AsyncMock(side_effect=reads)), + ) + + +async def test_a_turn_reads_its_text_from_the_slice_replay_would_serve(): + refresh = AsyncMock() + pointer, reader, rows = _branch({"r-1": _row("from the slice")}) + with pointer, reader, rows, patch(_REFRESH, refresh): + texts = await committed.committed_texts("t-1", {"r-1": 1}) + assert texts == {"r-1": "from the slice"} + refresh.assert_not_awaited() + + +async def test_a_slice_cut_on_another_branch_is_projected_again_first(): + refresh = AsyncMock() + pointer, reader, rows = _branch( + [{"r-1": _row("old branch", tail="tail-0")}, {"r-1": _row("this branch")}] + ) + with pointer, reader, rows, patch(_REFRESH, refresh): + texts = await committed.committed_texts("t-1", {"r-1": 1}) + assert texts == {"r-1": "this branch"} + refresh.assert_awaited_once_with("t-1", {"r-1"}, claims_pass=False) + + +async def test_a_turn_still_without_a_slice_is_left_to_the_caller(): + pointer, reader, rows = _branch([{}, {}]) + with pointer, reader, rows, patch(_REFRESH, AsyncMock()): + assert await committed.committed_texts("t-1", {"r-1": 1}) == {} + + +async def test_a_branch_that_does_not_pair_reads_and_projects_nothing(): + refresh = AsyncMock() + stamped_elsewhere = TurnAnchor( + turn_ordinal=0, input_checkpoint_id="in-1", tail_checkpoint_id="tail-1", turn_index=7 + ) + pointer, reader, rows = _branch({}, anchors=(stamped_elsewhere,)) + with pointer, reader, rows as read, patch(_REFRESH, refresh): + assert await committed.committed_texts("t-1", {"r-1": 1}) == {} + read.assert_not_awaited() + refresh.assert_not_awaited() diff --git a/tests/unit/server/services/history/test_event_ledger.py b/tests/unit/server/services/history/test_event_ledger.py index 2d739e01f..0a7ed234c 100644 --- a/tests/unit/server/services/history/test_event_ledger.py +++ b/tests/unit/server/services/history/test_event_ledger.py @@ -53,6 +53,10 @@ "provenance", # provenance_records "credit_usage", # conversation_usages "error", # conversation_responses.errors + metadata (terminal event) + "steering_returned", # subagent_runs.replay_facts: undelivered steering + # has no checkpoint twin by design (nothing was injected into the graph); + # a run settled before the facts keeps its stored copy, which the legacy + # facts backfill carries into the response row } LIVE_ONLY = { @@ -68,6 +72,7 @@ "stream_gap", # reconnect trim notice (Last-Event-ID < stream head) "run_end", # post-CAS terminal frame closing the run stream (I6) "replay_done", # replay sentinel + "history_page", # paged replay header: oldest turn sent, older remain "snapshot", # read-time cursors for the runs replay could not project; # recomputed from ledger+Redis on every load, never replayed from a store "market_watch_update", # transient stamp notice, accumulate=False; the @@ -94,11 +99,7 @@ # KNOWN GAP: survives replay only through persisted sse_events. Before # sse_events writes stop (cutover step 5) each entry here must move to a # category above or be explicitly accepted as not replayed. -STORED_EVENTS_ONLY: set[str] = { - "steering_returned", # undelivered-steering drain at run end; no - # checkpoint twin by design (nothing was injected into the graph), so - # the collected stored copy is its only durable home -} +STORED_EVENTS_ONLY: set[str] = set() _CATEGORIES = { "checkpoint": CHECKPOINT_PROJECTED, @@ -176,12 +177,12 @@ def test_replay_transition_rules_cover_the_replayable_set(): """The replay module's passthrough/stored-preferred tuples must stay inside the replayable ledger (a stored-preferred type outside the ledger would silently vanish post-cutover).""" - from src.server.services.history.replay.stored_merge import ( - _PASSTHROUGH_EVENTS, - _STORED_PREFERRED_EVENTS, + from src.server.services.history.replay.legacy import ( + PASSTHROUGH_EVENTS, + STORED_PREFERRED_EVENTS, ) replayable = CHECKPOINT_PROJECTED | TABLE_SOURCED - assert set(_STORED_PREFERRED_EVENTS) <= replayable + assert set(STORED_PREFERRED_EVENTS) <= replayable # Passthrough may additionally carry the stored-events-only legacy set. - assert set(_PASSTHROUGH_EVENTS) <= replayable | STORED_EVENTS_ONLY + assert set(PASSTHROUGH_EVENTS) <= replayable | STORED_EVENTS_ONLY diff --git a/tests/unit/server/services/history/test_legacy_facts.py b/tests/unit/server/services/history/test_legacy_facts.py new file mode 100644 index 000000000..7fdc52405 --- /dev/null +++ b/tests/unit/server/services/history/test_legacy_facts.py @@ -0,0 +1,661 @@ +"""A backfilled turn replays from its legacy facts exactly as from its stored +events, and without reading them. + +The backfill derives what a turn's stored events add to its projection once, +into ``replay_facts.legacy``. The facts are stored as jsonb, so they must +survive a JSON round trip, and must place every row where the stored-event +merge placed it. +""" + +from __future__ import annotations + +import copy +import json +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import AIMessage, ToolMessage + +from scripts.utils.backfill_replay_facts import Derivation +from src.server.database.replay_facts import LEGACY_KEY +from src.server.services.history import replay +from src.server.services.history.replay import legacy, stored_merge +from tests.unit.server.services.history.replay_builders import ( + thread_rows, + THREAD, + TIP, + ThreadHistory, + _cache_probe, + _mock_reader, + _query, + _response, + _turn, +) +from tests.unit.server.services.history.test_replay_facts import ( + _RUN_FACTS, + _returned, + _signal, + _task_reader, + _text, +) + +# A pointer as older turns checkpointed it, which is what derive has to +# recognize whatever the agent-facing template reads today. +_EVICTION_POINTER = ( + "Tool result too large, the result of this tool call tc-r was saved in the " + "filesystem at this path: .agents/large_tool_results/tc-r.md" +) + + +def _stored_reads_refused(response): + raise AssertionError("a backfilled turn read its stored events") + + +async def _both_ways(monkeypatch, responses, run_facts=None): + """Each turn's lines from its stored events, the facts they derived, and + the lines again from those facts (through JSON) with the events gone.""" + queries = [_query(ti) for ti in sorted(responses)] + turns = sorted(responses) + derivation = Derivation() + before = await replay.project_turns( + thread_rows(queries, responses, run_facts=run_facts), + TIP, + turn_indexes=turns, + derive=derivation, + ) + derived = derivation.turns + assert all(d.verified for d in derived.values()) + backfilled = {} + for ti, response in responses.items(): + row = {k: v for k, v in response.items() if k != "sse_events"} + row["replay_facts"] = { + **(response.get("replay_facts") or {}), + LEGACY_KEY: json.loads(json.dumps(derived[ti].facts)), + } + backfilled[ti] = row + monkeypatch.setattr(stored_merge, "_stored_events", _stored_reads_refused) + again = Derivation() + after = await replay.project_turns( + thread_rows(queries, backfilled, run_facts=run_facts), + TIP, + turn_indexes=turns, + derive=again, + ) + assert again.turns == {} # nothing derived a second time + return before, {ti: d.facts for ti, d in derived.items()}, after + + +def _items(lines): + return {ti: [line.item() for line in turn] for ti, turn in lines.items()} + + +# ------------------------------------------------------------ through replay + + +@pytest.mark.asyncio +async def test_widgets_results_and_signals_replay_from_facts(monkeypatch): + messages = [ + AIMessage( + content="looking", + id="ai-1", + tool_calls=[ + {"name": "ShowWidget", "args": {}, "id": "tc-s"}, + {"name": "get_sec_filing", "args": {}, "id": "tc-r"}, + ], + ), + ToolMessage( + content="shown", + tool_call_id="tc-s", + name="ShowWidget", + id="tm-1", + artifact={"html": "
", "title": "W"}, + ), + ToolMessage( + content=_EVICTION_POINTER, tool_call_id="tc-r", name="get_sec_filing", id="tm-2" + ), + AIMessage(content="done", id="ai-2"), + ] + + def widget(artifact_id, data): + payload = {"html": "
", "title": "W", "data": data} + return { + "event": "artifact", + "data": { + "artifact_type": "html_widget", + "artifact_id": artifact_id, + "payload": payload, + }, + } + + stored = [ + {"event": "message_chunk", "data": {"id": "lc-1", "content_type": "text", "content": "looking"}}, + widget("widget_1", {"a.csv": "1"}), + # Beyond what the checkpoint projected: placed like a signal row. + widget("widget_2", {"b.csv": "2"}), + {"event": "tool_call_result", "data": {"tool_call_id": "tc-r", "content": "full", "content_type": "text"}}, + {"event": "context_window", "data": {"action": "offload", "kind": "tool_args"}}, + {"event": "provenance", "data": {"source_type": "web", "identifier": "https://a.example"}}, + {"event": "message_chunk", "data": {"id": "lc-2", "content_type": "text", "content": "done"}}, + {"event": "credit_usage", "data": {"total_credits": 1.0}}, + ] + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, messages)])) + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + merge = facts[0]["merge"] + # Named by the projected widget it upgraded; the extra one by nothing. + assert [artifact_id for artifact_id, _data in merge["widgets"]] == ["tc-s", None] + assert merge["results"] == {"tc-r": {"content": "full", "content_type": "text"}} + assert [s.get("event", "widget") for s in merge["signals"]] == [ + "widget", "widget", "context_window", "provenance", "credit_usage" + ] + assert after == before + result = next(i for i in _items(after)[0] if i["event"] == "tool_call_result" + and i["data"]["tool_call_id"] == "tc-r") + assert result["data"]["content"] == "full" + + +@pytest.mark.asyncio +async def test_reasoning_durations_replay_from_facts(monkeypatch): + reasoning = AIMessage( + content=[ + {"type": "thinking", "thinking": "weighing it"}, + {"type": "text", "text": "the answer"}, + ], + id="msg_1", + ) + stored = [ + _signal("lc-1", "start"), + _signal("lc-1", "complete", elapsed_ms=1500), + _signal("lc-1", "start"), + _signal("lc-1", "complete", elapsed_ms=500), + _text("lc-1", "the answer"), + ] + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, [reasoning])])) + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + assert [ms for _key, ms in facts[0]["merge"]["durations"]] == [2000] + assert after == before + + +@pytest.mark.asyncio +async def test_interrupt_and_error_keep_their_places(monkeypatch): + messages = [ + AIMessage( + content="", + id="ai-1", + tool_calls=[{"name": "AskUserQuestion", "args": {}, "id": "tc-q"}], + ), + ToolMessage( + content="answer: blue", tool_call_id="tc-q", name="AskUserQuestion", id="tm-1" + ), + ] + stored = [ + {"event": "tool_calls", "data": {"id": "run-1", "tool_calls": [{"id": "tc-q", "name": "AskUserQuestion"}]}}, + {"event": "interrupt", "data": {"interrupt_id": "int-q", "action_requests": [{"type": "ask_user_question"}]}}, + {"event": "tool_call_result", "data": {"tool_call_id": "tc-q", "content": "answer: blue"}}, + {"event": "error", "data": {"error": "provider exploded"}}, + ] + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, messages)])) + + before, _facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + assert after == before + events = [i["event"] for i in _items(after)[0]] + assert events.index("tool_calls") < events.index("interrupt") < events.index("tool_call_result") + + +@pytest.mark.asyncio +async def test_unresolved_images_resolve_to_the_stored_url_with_their_basename(monkeypatch): + """Capture rewrote the stored copy's paths to keys ending in the file's + basename. The turn keeps its projection, so a stored message the + checkpoint never committed stays out, and a path capture never reached + stays a path.""" + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=[ + _turn( + 0, + [AIMessage(content="![chart](work/chart.png) ![raw](work/raw.png)", id="ai-1")], + ) + ], + ), + ) + stored = [ + _text("lc-1", "![chart](https://cdn/abc/chart.png) ![raw](work/raw.png)"), + _text("lc-2", "a reply the checkpoint never committed"), + ] + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + assert facts[0]["images"] == {"work/chart.png": "https://cdn/abc/chart.png"} + assert "verbatim" not in facts[0] + assert after == before + assert [i["data"]["content"] for i in _items(after)[0] if i["event"] == "message_chunk"] == [ + "![chart](https://cdn/abc/chart.png) ![raw](work/raw.png)" + ] + + +@pytest.mark.asyncio +async def test_a_run_lanes_unresolved_image_resolves_from_the_stored_copy(monkeypatch): + """A run settled before its row kept facts has its image URL only in + the stored copy, which a run lane resolves from like the main lane.""" + _cache_probe(monkeypatch) + _task_reader(monkeypatch) + stored = [_text("lc-s", "see ![c](https://cdn/abc/sub.png)", agent="task:tsk1")] + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + assert facts[0]["images"] == {"work/sub.png": "https://cdn/abc/sub.png"} + assert after == before + assert [ + i["data"]["content"] + for i in _items(after)[0] + if i["event"] == "message_chunk" and i["data"].get("agent") == "task:tsk1" + ] == ["see ![c](https://cdn/abc/sub.png)"] + + +@pytest.mark.asyncio +async def test_a_run_lanes_ambiguous_image_keeps_the_main_lane_projected(monkeypatch): + """Only the main lane can replay verbatim, so a run lane's path that + several stored URLs could match stays a path and sends nothing there.""" + _cache_probe(monkeypatch) + _task_reader(monkeypatch) + stored = [ + _text("lc-s1", "see ![c](https://cdn/1/sub.png)", agent="task:tsk1"), + _text("lc-s2", "or ![c](https://cdn/2/sub.png)", agent="task:tsk1"), + ] + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + assert "verbatim" not in facts[0] + assert "images" not in facts[0] + assert after == before + + +@pytest.mark.asyncio +async def test_images_no_basename_settles_replay_their_main_lane_verbatim(monkeypatch): + """Two stored URLs share the basename of both projected paths, so only + the stored copy knows which image is which.""" + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=[ + _turn( + 0, + [AIMessage(content="![a](work/a/chart.png) ![b](work/b/chart.png)", id="ai-1")], + ) + ], + ), + ) + stored = [_text("lc-1", "![a](https://cdn/1/chart.png) ![b](https://cdn/2/chart.png)")] + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0, stored)}) + + assert facts[0]["verbatim"] == stored + assert "images" not in facts[0] + assert after == before + chunk = next(i for i in _items(after)[0] if i["event"] == "message_chunk") + assert chunk["data"]["content"] == "![a](https://cdn/1/chart.png) ![b](https://cdn/2/chart.png)" + + +@pytest.mark.asyncio +async def test_a_run_lane_keeps_its_own_returns_once(monkeypatch): + """A run with facts places its returns; the stored copy adds only the + ones its facts do not hold.""" + _cache_probe(monkeypatch) + _task_reader(monkeypatch) + stored = [_returned("in-1"), _returned("in-2", content="and NVDA")] + + before, _facts, after = await _both_ways( + monkeypatch, {0: _response(0, stored)}, run_facts=[_RUN_FACTS] + ) + + assert after == before + returns = [i["data"]["input_id"] for i in _items(after)[0] if i["event"] == "steering_returned"] + assert sorted(returns) == ["in-1", "in-2"] + + +@pytest.mark.asyncio +async def test_a_turn_without_stored_events_is_known_empty(monkeypatch): + _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=[_turn(0, [AIMessage(content="hi", id="ai-1")])]), + ) + + before, facts, after = await _both_ways(monkeypatch, {0: _response(0)}) + + assert facts[0] == {"v": legacy.FACTS_VERSION} + assert after == before + + +@pytest.mark.asyncio +async def test_facts_the_row_would_hold_differently_are_not_verified(monkeypatch): + """A jsonb bind drops a NUL, so facts carrying one would replay without + it: the backfill must not write them.""" + _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=[_turn(0, [AIMessage(content="hi", id="ai-1")])]), + ) + stored = [ + {"event": "message_chunk", "data": {"id": "lc-1", "content_type": "text", "content": "hi"}}, + {"event": "provenance", "data": {"source_type": "web", "identifier": "a\x00b"}}, + ] + derivation = Derivation() + await replay.project_turns( + thread_rows([_query(0)], {0: _response(0, stored)}), + TIP, + turn_indexes=[0], + derive=derivation, + ) + + assert derivation.turns[0].verified is False + signal = derivation.turns[0].facts["merge"]["signals"][0] + assert signal["data"]["identifier"] == "ab" # as the row would hold it + + +# ------------------------------------------------------------ facts_of + + +@pytest.mark.parametrize( + ("replay_facts", "expected"), + [ + (None, None), + ({"v": 1, "reasoning_ms": {}}, None), + ({"legacy": "not a dict"}, None), + ({"legacy": {"v": legacy.FACTS_VERSION + 1, "merge": {}}}, None), + ({"legacy": {"v": legacy.FACTS_VERSION}}, {"v": legacy.FACTS_VERSION}), + ], + ids=["none", "run_facts_only", "malformed", "other_version", "known_empty"], +) +def test_only_this_versions_facts_stand_in_for_the_events(replay_facts, expected): + assert legacy.facts_of({"replay_facts": replay_facts}) == expected + + +# ------------------------------------------------------------ apply + + +def _row(event, **data): + return {"event": event, "data": data} + + +def _projection(): + return [ + _row("message_chunk", agent="main", id="m1", content_type="text", content="a"), + _row("tool_calls", agent="main", tool_calls=[{"id": "tc-1"}, {"id": "tc-2"}]), + _row("tool_call_result", agent="main", tool_call_id="tc-1", content="r1"), + _row("tool_call_result", agent="main", tool_call_id="tc-2", content="r2"), + _row("message_chunk", agent="main", id="m2", content_type="text", content="b"), + ] + + +def _events(items): + return [(i["event"], i["data"].get("tool_call_id") or i["data"].get("tag")) for i in items] + + +M1 = ["message_chunk", "main", "m1", "text"] +TC1 = ["tool_call_result", "tc-1"] +TC2 = ["tool_call_result", "tc-2"] +GONE = ["tool_call_result", "gone"] + + +def test_a_signal_without_at_follows_the_one_before_it(): + merge = { + "anchors": [TC1], + "signals": [ + {"event": "provenance", "data": {"tag": "p1"}}, + {"event": "provenance", "data": {"tag": "p2"}, "at": 0}, + {"event": "provenance", "data": {"tag": "p3"}}, + ], + } + assert _events(legacy.apply(_projection(), merge)) == [ + ("provenance", "p1"), + ("message_chunk", None), + ("tool_calls", None), + ("tool_call_result", "tc-1"), + ("provenance", "p2"), + ("provenance", "p3"), + ("tool_call_result", "tc-2"), + ("message_chunk", None), + ] + + +def test_a_lost_anchor_falls_back_to_the_newest_that_remains(): + merge = {"anchors": [M1, TC1, GONE], "signals": [{"event": "provenance", "data": {"tag": "p"}, "at": 2}]} + out = _events(legacy.apply(_projection(), merge)) + assert out[2:4] == [("tool_call_result", "tc-1"), ("provenance", "p")] + + +def test_the_fallback_stops_at_the_previous_signal(): + """Anchors before the previous signal placed that one; a signal whose own + anchors are all gone follows it.""" + merge = { + "anchors": [M1, TC1, GONE], + "signals": [ + {"event": "provenance", "data": {"tag": "p1"}, "at": 1}, + {"event": "provenance", "data": {"tag": "p2"}, "at": 2}, + ], + } + out = _events(legacy.apply(_projection(), merge)) + assert out[2:5] == [("tool_call_result", "tc-1"), ("provenance", "p1"), ("provenance", "p2")] + + +def test_a_signal_whose_anchors_are_all_gone_leads_the_turn(): + merge = {"anchors": [GONE], "signals": [{"event": "credit_usage", "data": {"tag": "c"}, "at": 0}]} + assert _events(legacy.apply(_projection(), merge))[0] == ("credit_usage", "c") + + +def test_each_call_of_a_parallel_round_anchors_on_the_round(): + merge = { + "anchors": [["tool_calls", "tc-2"]], + "signals": [{"event": "provenance", "data": {"tag": "p"}, "at": 0}], + } + out = _events(legacy.apply(_projection(), merge)) + assert out[1:3] == [("tool_calls", None), ("provenance", "p")] + + +def test_an_error_follows_everything_its_lane_committed(): + merge = {"anchors": [TC1], "signals": [{"event": "error", "data": {"agent": "main", "tag": "e"}, "at": 0}]} + assert _events(legacy.apply(_projection(), merge))[-1] == ("error", "e") + + +def test_a_result_is_restored_by_its_call_whatever_the_projection_shows(): + """The pointer text is agent-facing and may be reworded; the facts + already decided which results to restore.""" + projected = _projection() + projected[2]["data"]["content"] = "Saved to .agents/large_tool_results/tc-1.md" + merge = {"results": {"tc-1": {"content": "full", "content_type": "text"}}} + out = legacy.apply(projected, merge) + assert (out[2]["data"]["content"], out[2]["data"]["content_type"]) == ("full", "text") + assert out[3]["data"]["content"] == "r2" + + +def _widget(artifact_id, **payload): + return _row("artifact", artifact_type="html_widget", artifact_id=artifact_id, payload=payload) + + +def _stored_widget(k): + return {"artifact_type": "html_widget", "artifact_id": f"s{k}", "payload": {"k": k}} + + +def test_a_widget_takes_the_payload_that_names_it(): + projected = [_widget("tc-a"), _widget("tc-b")] + merge = {"widgets": [["tc-b", _stored_widget(1)], ["tc-a", _stored_widget(0)]]} + out = legacy.apply(projected, merge) + assert [i["data"]["artifact_id"] for i in out] == ["s0", "s1"] + + +def test_an_unnamed_payload_pairs_in_order_or_is_placed(): + """One the projection did not emit then pairs with a widget no payload + names; one whose widget the projection dropped is placed where it + streamed, and a payload a widget took is not placed again.""" + projected = [ + _widget("tc-new"), + _row("message_chunk", agent="main", id="m1", content_type="text", content="a"), + ] + merge = { + "widgets": [["tc-gone", _stored_widget(0)], [None, _stored_widget(1)]], + "anchors": [M1], + "signals": [{"widget": 0, "at": 0}, {"widget": 1}, {"widget": 5}], + } + out = legacy.apply(projected, merge) + assert [(i["event"], i["data"].get("artifact_id")) for i in out] == [ + ("artifact", "s1"), + ("message_chunk", None), + ("artifact", "s0"), + ] + # The placed copy is the turn's own: enriching it leaves the facts alone. + out[-1]["data"]["turn_index"] = 0 + assert "turn_index" not in merge["widgets"][0][1] + + +def test_no_merge_leaves_the_projected_signals_standing(): + projected = [*_projection(), _row("context_window", action="offload")] + assert legacy.apply(projected, None) is projected + + +def test_a_message_without_an_id_is_named_by_its_occurrence(): + """The projector gives each id-less message the same placeholder.""" + rows = [ + _row("message_chunk", agent="main", id="unknown", content_type="reasoning", content="x"), + _row("message_chunk", agent="main", id="unknown", content_type="text", content="a"), + _row("tool_calls", agent="main", tool_calls=[{"id": "tc-1"}]), + _row("message_chunk", agent="main", id="unknown", content_type="text", content="b"), + ] + assert legacy.row_keys(rows) == [ + ("message_chunk", "main", "unknown", "reasoning"), + ("message_chunk", "main", "unknown", "text"), + ("tool_calls", "tc-1"), + ("message_chunk", "main", "unknown", "text", 1), + ] + + +# ------------------------------------------------------------ projection drift +# +# Facts resolved against one projection and applied to a later one, the way a +# backfilled row meets a projector change: what the facts carry stays with +# the rows it was resolved onto. + + +def _message(message_id, elapsed_ms=None): + close = {"elapsed_ms": elapsed_ms} if elapsed_ms is not None else {} + return [ + _row("message_chunk", agent="main", id=message_id, content_type="reasoning_signal", content="start"), + _row("message_chunk", agent="main", id=message_id, content_type="reasoning", content="r"), + _row("message_chunk", agent="main", id=message_id, content_type="reasoning_signal", content="complete", **close), + _row("message_chunk", agent="main", id=message_id, content_type="text", content="t"), + ] + + +def _call(call_id): + return [ + _row("tool_calls", agent="main", tool_calls=[{"id": call_id}]), + _row("tool_call_result", agent="main", tool_call_id=call_id, content="ok"), + ] + + +def _resolved_then_applied(then, now, stored): + facts = json.loads(json.dumps(stored_merge.derive_turn(then, [], stored))) + return legacy.apply_turn(copy.deepcopy(now), [], facts) + + +def _shape(items): + out = [] + for item in items: + data = item["data"] + if item["event"] == "message_chunk" and data["content_type"] == "reasoning_signal": + if data["content"] == "complete": + out.append(f"close:{data['id']}:{data.get('elapsed_ms')}") + elif item["event"] == "message_chunk" and data["content_type"] == "text": + out.append(f"text:{data['id']}") + elif item["event"] == "artifact": + out.append(f"widget:{data['artifact_id']}") + elif item["event"] not in ("message_chunk", "tool_calls", "tool_call_result"): + out.append(item["event"]) + return out + + +def test_a_message_the_projection_gained_takes_nothing_of_its_neighbours(): + stored = [ + *_message("lc0", 1000), + *_call("A"), + *_message("lc1", 2000), + _row("context_window", agent="main", action="token_usage"), + ] + then = [*_message("p0"), *_call("A"), *_message("p1")] + now = [*_message("p0"), *_call("A"), *_message("px"), *_message("p1")] + + assert _shape(_resolved_then_applied(then, now, stored)) == [ + "close:p0:1000", "text:p0", + "close:px:None", "text:px", + "close:p1:2000", "text:p1", + "context_window", + ] + + +@pytest.mark.parametrize( + ("now", "expected"), + [ + ([*_message("p1")], ["close:p1:None", "text:p1", "credit_usage"]), + ([*_message("p0")], ["close:p0:None", "text:p0", "credit_usage"]), + ], + ids=["first_dropped", "last_dropped"], +) +def test_a_signal_whose_anchor_the_projection_dropped_stays_in_place(now, expected): + stored = [*_message("lc0"), *_message("lc1"), _row("credit_usage", agent="main")] + then = [*_message("p0"), *_message("p1")] + + assert _shape(_resolved_then_applied(then, now, stored)) == expected + + +@pytest.mark.parametrize( + ("now", "expected"), + [ + ( + [_widget("tc-new"), _widget("tc-a"), _widget("tc-b")], + ["widget:tc-new", "widget:s0", "widget:s1", "widget:s2", "provenance"], + ), + ( + [*_call("x"), _widget("tc-b")], + ["widget:s0", "widget:s1", "widget:s2", "provenance"], + ), + ( + [_widget("tc-a"), _widget("tc-b"), _widget("tc-c")], + ["widget:s0", "widget:s1", "widget:s2", "provenance"], + ), + ], + ids=["widget_gained", "widget_dropped", "missed_widget_emitted"], +) +def test_a_widget_payload_stays_with_its_widget(now, expected): + """The third payload streamed beyond what the projection then emitted: + placed where it streamed, or taken by a widget that appears after its + neighbours' widgets, never by one ahead of them.""" + stored = [ + _row("artifact", artifact_type="html_widget", artifact_id=f"s{k}", payload={"k": k}) + for k in range(3) + ] + [_row("provenance", agent="main")] + then = [_widget("tc-a"), _widget("tc-b")] + + assert _shape(_resolved_then_applied(then, now, stored)) == expected + + +# ------------------------------------------------------------ re-read + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("verbatim", "asked"), [(False, legacy.FACTS_VERSION), (True, None)]) +async def test_the_re_read_leaves_out_events_the_facts_stand_in_for( + monkeypatch, verbatim, asked +): + read = AsyncMock(return_value={"resp-0": {"conversation_response_id": "resp-0"}}) + monkeypatch.setattr("src.server.database.conversation.get_replay_responses", read) + + await replay.with_stored_events( + {0: {"conversation_response_id": "resp-0"}}, [0], verbatim=verbatim + ) + + read.assert_awaited_once_with(["resp-0"], legacy_facts=asked) diff --git a/tests/unit/server/services/history/test_pair_turns.py b/tests/unit/server/services/history/test_pair_turns.py new file mode 100644 index 000000000..c2d0aba7f --- /dev/null +++ b/tests/unit/server/services/history/test_pair_turns.py @@ -0,0 +1,73 @@ +"""Pairing a branch's turns with the thread's persisted turns. + +One rule serves replay and the edit/regenerate listing: a stamped turn keeps +its stamp, any other turn names the next persisted turn after the one before. +Replay used to pair an unstamped turn by its position instead, so the rule is +held to agree with that wherever replay accepted the pairing. +""" + +from __future__ import annotations + +from itertools import combinations, product + +from src.server.services.history.reader import TurnAnchor, pair_turns + + +def _anchors(stamps): + return [ + TurnAnchor( + turn_ordinal=k, + input_checkpoint_id=f"cp-{k}", + tail_checkpoint_id=None, + turn_index=stamp, + is_resume=stamp is None, + ) + for k, stamp in enumerate(stamps) + ] + + +def test_a_resume_after_a_dead_turn_names_the_turn_after_the_one_it_answers(): + # Turn 1 died before its first checkpoint; turn 3 resumes turn 2. + assert pair_turns(_anchors([0, 2, None]), [0, 1, 2, 3]) == [0, 2, 3] + # By position the resume named turn 2 again, and replay gave up. + assert _positional([0, 2, None], [0, 1, 2, 3]) is None + + +def test_unstamped_turns_take_the_persisted_turns_in_order(): + assert pair_turns(_anchors([None, None, None]), [0, 1]) == [0, 1, None] + + +def test_a_stamp_is_kept_even_where_no_persisted_turn_has_it(): + assert pair_turns(_anchors([5, None]), [0, 1]) == [5, None] + + +def _positional(stamps, persisted): + """Replay's former rule: a stamp, else the persisted turn at the turn's + position; None where it raised.""" + known = set(persisted) + paired = [] + for k, stamp in enumerate(stamps): + if stamp is not None: + if stamp not in known: + return None + paired.append(stamp) + elif k < len(persisted): + paired.append(persisted[k]) + else: + return None + return paired if paired == sorted(set(paired)) else None + + +def test_every_pairing_replay_accepted_by_position_is_unchanged(): + values = range(5) + checked = 0 + for size in range(len(values) + 1): + for persisted in combinations(values, size): + for length in range(1, 5): + for stamps in product([None, *values, 5], repeat=length): + accepted = _positional(stamps, list(persisted)) + if accepted is None: + continue + checked += 1 + assert pair_turns(_anchors(stamps), list(persisted)) == accepted + assert checked == 584 # every accepted case of up to four turns over five diff --git a/tests/unit/server/services/history/test_projection_cache.py b/tests/unit/server/services/history/test_projection_cache.py deleted file mode 100644 index 819f0e485..000000000 --- a/tests/unit/server/services/history/test_projection_cache.py +++ /dev/null @@ -1,556 +0,0 @@ -"""Per-turn projection cache: primitives, fast-path assembly, backfill.""" - -from __future__ import annotations - -import asyncio -import json -from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock - -import pytest -from langchain_core.messages import AIMessage, HumanMessage, ToolMessage - -from src.server.services.history import projection_cache -from src.server.services.history.reader import ( - TaskHistory, - ThreadHistory, - TurnAnchor, - TurnSlice, -) -from src.server.services.history.replay import build_checkpoint_replay_items - -pytestmark = pytest.mark.asyncio - -THREAD = "thread-pc" - -_END_SENTINEL = json.dumps({"event": "subagent_stream_end"}) - - -class _FakeRedis: - """Raw-client stand-in: per-key subagent stream tails (None = no stream).""" - - def __init__(self): - self.stream_tails: dict[str, str] = {} - - async def xrevrange(self, key, count=1): - raw = self.stream_tails.get(key) - if raw is None: - return [] - return [(b"0-1", {b"event": raw.encode()})] - - -class _FakeCacheClient: - """In-memory stand-in matching the RedisCacheClient surface used here.""" - - def __init__(self): - self.store: dict[str, object] = {} - self.enabled = True - self.client = _FakeRedis() - - async def mget(self, keys): - return [self.store.get(k) for k in keys] - - async def set(self, key, value, ttl=None): - self.store[key] = value - return True - - -@pytest.fixture -def fake_cache(monkeypatch): - client = _FakeCacheClient() - monkeypatch.setattr(projection_cache, "get_cache_client", lambda: client) - return client - - -def _anchor(ordinal, tail, turn_index=None): - return TurnAnchor( - turn_ordinal=ordinal, - input_checkpoint_id=f"cp-in-{ordinal}", - tail_checkpoint_id=tail, - turn_index=turn_index, - ) - - -def _turn(ordinal, messages, tail, turn_index=None): - return TurnSlice( - turn_ordinal=ordinal, - input_checkpoint_id=f"cp-in-{ordinal}", - end_checkpoint_id=f"cp-end-{ordinal}", - user_message=HumanMessage(content="hello", id=f"h-{ordinal}"), - messages=messages, - turn_index=turn_index, - tail_checkpoint_id=tail, - ) - - -def _query(turn_index, content="hello"): - return {"turn_index": turn_index, "content": content, "type": "user", "created_at": "t0"} - - -def _response(turn_index, status="completed"): - return { - "conversation_response_id": f"resp-{turn_index}", - "turn_index": turn_index, - "sse_events": [], - "status": status, - } - - -def _key(tail, response=None, provenance=None, usage=None): - """Production cache key for a turn — keys carry the input fingerprint, so - a test that seeds or asserts an entry must derive it the same way.""" - return projection_cache._key( - THREAD, - tail, - projection_cache.turn_fingerprint(response, provenance or [], usage), - ) - - -def _mock_reader(monkeypatch, *, anchors=None, tip="cp-tip", history=None, - tip_interrupts=None): - reader = MagicMock() - reader.aget_turn_anchors = AsyncMock(return_value=(anchors or [], tip)) - reader.aget_tip_interrupts = AsyncMock(return_value=tip_interrupts or []) - reader.aget_thread_history = AsyncMock( - return_value=history or ThreadHistory(thread_id=THREAD) - ) - reader.aget_recent_history = AsyncMock( - return_value=history or ThreadHistory(thread_id=THREAD) - ) - reader.aget_task_history = AsyncMock(return_value=TaskHistory()) - monkeypatch.setattr( - "src.server.services.history.replay.CheckpointHistoryReader.get_instance", - lambda: reader, - ) - return reader - - -# ---------------------------------------------------------------- primitives - - -async def test_store_and_get_round_trip(fake_cache): - items = [{"event": "user_message", "data": {"content": "hi"}}] - await projection_cache.store_turn(THREAD, "tail-1", "fp-1", items) - cached = await projection_cache.get_cached_turns( - THREAD, [("tail-1", "fp-1"), ("tail-2", "fp-2")] - ) - assert cached["tail-1"] == items - assert cached["tail-2"] is None - - -async def test_superseded_fingerprint_is_unreachable(fake_cache): - """The whole point of the fingerprint: an entry built before a turn's - post-finalize archive drain must not be servable afterwards, without - anyone having to evict it.""" - stale = [{"event": "user_message", "data": {"content": "pre-drain"}}] - await projection_cache.store_turn(THREAD, "tail-1", "fp-before", stale) - - cached = await projection_cache.get_cached_turns(THREAD, [("tail-1", "fp-after")]) - - assert cached["tail-1"] is None - # The old generation is still in Redis (it expires by TTL) — it is - # unreachable, not deleted. - assert projection_cache._key(THREAD, "tail-1", "fp-before") in fake_cache.store - - -async def test_fingerprint_tracks_each_mutable_input(fake_cache): - response = {"conversation_response_id": "r0", "sse_events": [], "status": "completed"} - base = projection_cache.turn_fingerprint(response, [], None) - - drained = dict(response, sse_events=[{"event": "message_chunk", "data": {}}]) - assert projection_cache.turn_fingerprint(drained, [], None) != base - assert projection_cache.turn_fingerprint(response, [{"identifier": "x"}], None) != base - assert projection_cache.turn_fingerprint(response, [], {"total_credits": 1}) != base - # Stable across dict ordering — psycopg row key order is not a contract. - reordered = {k: response[k] for k in reversed(list(response))} - assert projection_cache.turn_fingerprint(reordered, [], None) == base - - -async def test_fingerprint_reads_the_response_row_through_its_row_version(): - """The replay query ships ``xmin``, so the fingerprint never serializes the - transcript: any rewrite of the row lands a new version, and that is what - invalidates.""" - response = { - "conversation_response_id": "r0", - "status": "completed", - "sse_events": [{"event": "message_chunk", "data": {}}], - "xmin": "100", - } - base = projection_cache.turn_fingerprint(response, [], None) - - drained = dict(response, sse_events=[], xmin="101") - assert projection_cache.turn_fingerprint(drained, [], None) != base - # A retry's row is a different response entirely, version or not. - assert ( - projection_cache.turn_fingerprint( - dict(response, conversation_response_id="r1"), [], None - ) - != base - ) - # The columns themselves are not hashed — a rewrite always bumps xmin. - assert ( - projection_cache.turn_fingerprint(dict(response, sse_events=[]), [], None) - == base - ) - - -async def test_store_canonicalizes_to_wire_json(fake_cache): - # The endpoint serializes with json.dumps(default=str); cached entries - # must round-trip datetimes to the identical wire text (str(), not - # isoformat — space separator). - ts = datetime(2026, 7, 7, 12, 0, 30, tzinfo=timezone.utc) - await projection_cache.store_turn( - THREAD, "tail-1", "fp-1", [{"event": "user_message", "data": {"timestamp": ts}}] - ) - cached = await projection_cache.get_cached_turns(THREAD, [("tail-1", "fp-1")]) - assert cached["tail-1"][0]["data"]["timestamp"] == str(ts) - - -async def test_store_skips_oversize_and_missing_tail(fake_cache): - await projection_cache.store_turn(THREAD, None, "fp-1", [{"event": "x", "data": {}}]) - big = [{"event": "x", "data": {"blob": "a" * (projection_cache._MAX_ENTRY_BYTES + 1)}}] - await projection_cache.store_turn(THREAD, "tail-big", "fp-1", big) - assert fake_cache.store == {} - - -async def test_inactive_without_connected_client(fake_cache): - fake_cache.client = None - assert projection_cache.cache_active() is False - - -async def test_zero_ttl_disables(fake_cache, monkeypatch): - monkeypatch.setattr( - projection_cache, "get_replay_projection_cache_ttl", lambda: 0 - ) - assert projection_cache.cache_active() is False - - -# ---------------------------------------------------------- replay fast path - - -async def test_full_hit_assembles_without_materializing(monkeypatch, fake_cache): - entry0 = [{"event": "user_message", "data": {"thread_id": THREAD, "turn_index": 0}}] - entry1 = [ - {"event": "user_message", "data": {"thread_id": THREAD, "turn_index": 1}}, - {"event": "message_chunk", "data": {"content": "hi", "turn_index": 1}}, - ] - fake_cache.store[_key("tail-0", _response(0))] = entry0 - fake_cache.store[_key("tail-1", _response(1))] = entry1 - reader = _mock_reader( - monkeypatch, - anchors=[_anchor(0, "tail-0", 0), _anchor(1, "tail-1", 1)], - tip="tail-1", - tip_interrupts=[{"id": "int-1", "value": {"action_requests": []}}], - ) - - items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1)], - {0: _response(0), 1: _response(1)}, - ) - - assert items[: len(entry0) + len(entry1)] == entry0 + entry1 - assert items[-1]["event"] == "interrupt" - reader.aget_thread_history.assert_not_called() - reader.aget_recent_history.assert_not_called() - - -async def test_fast_path_stubs_inflight_turn_from_rows(monkeypatch, fake_cache): - fake_cache.store[_key("tail-0", _response(0))] = [ - {"event": "user_message", "data": {"turn_index": 0}} - ] - _mock_reader(monkeypatch, anchors=[_anchor(0, "tail-0", 0)], tip="tail-0") - - items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="in flight")], - {0: _response(0), 1: _response(1, status="streaming")}, - ) - - assert [i["event"] for i in items] == ["user_message", "user_message"] - assert items[1]["data"]["content"] == "in flight" - - -async def test_any_miss_rebuilds_and_backfills(monkeypatch, fake_cache): - turns = [ - _turn(0, [HumanMessage(content="hello", id="h-0"), - AIMessage(content="a0", id="ai-0")], "tail-0", 0), - _turn(1, [HumanMessage(content="hello", id="h-1"), - AIMessage(content="a1", id="ai-1")], "tail-1", 1), - ] - fake_cache.store[_key("tail-0", _response(0))] = [ - {"event": "user_message", "data": {"turn_index": 0}} - ] # tail-1 missing → full rebuild - reader = _mock_reader( - monkeypatch, - anchors=[_anchor(0, "tail-0", 0), _anchor(1, "tail-1", 1)], - tip="tail-1", - history=ThreadHistory(thread_id=THREAD, turns=turns), - ) - - items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1)], - {0: _response(0), 1: _response(1)}, - ) - - reader.aget_thread_history.assert_called_once() - assert [i["event"] for i in items] == [ - "user_message", "message_chunk", "user_message", "message_chunk", - ] - # Both turns backfilled under their tail keys, wire-canonical. - for tail, turn_index in (("tail-0", 0), ("tail-1", 1)): - entry = fake_cache.store[_key(tail, _response(turn_index))] - assert [i["event"] for i in entry] == ["user_message", "message_chunk"] - assert entry[-1]["data"]["turn_index"] == turn_index - - -async def test_post_drain_inputs_miss_the_pre_drain_entry(monkeypatch, fake_cache): - """The workflow-child regression. A settled turn's provenance rows are - (re)derived by the subagent collector's archive drain, seconds AFTER the - turn finalized and its projection was cached — and the drain never moves - the tail. Under a tail-only key the pre-drain, provenance-free entry - stayed servable for the full cache TTL; the fingerprint must make it - unreachable so the read rebuilds with the rows now present.""" - turns = [_turn(0, [HumanMessage(content="hello", id="h-0"), - AIMessage(content="a0", id="ai-0")], "tail-0", 0)] - reader = _mock_reader( - monkeypatch, - anchors=[_anchor(0, "tail-0", 0)], - tip="tail-0", - history=ThreadHistory(thread_id=THREAD, turns=turns), - ) - response = _response(0) - # Cached at finalize, before the drain: no provenance in the entry. - fake_cache.store[_key("tail-0", response)] = [ - {"event": "user_message", "data": {"turn_index": 0}} - ] - # The drain lands: rows now exist for the same response, same tail. - provenance = [ - { - "conversation_response_id": "resp-0", - "turn_index": 0, - "source_type": "web_search", - "identifier": "https://example.test/a", - "agent": "task:child1", - } - ] - - items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: response}, provenance=provenance - ) - - reader.aget_thread_history.assert_called_once() # missed → rebuilt - assert [i["event"] for i in items] == [ - "user_message", "message_chunk", "provenance", - ] - assert items[-1]["data"]["agent"] == "task:child1" - # Backfilled under the post-drain fingerprint, leaving the old generation - # stranded rather than overwritten. - assert _key("tail-0", response, provenance) in fake_cache.store - - -async def test_missing_tail_falls_back_to_rebuild(monkeypatch, fake_cache): - turns = [_turn(0, [HumanMessage(content="hello", id="h-0"), - AIMessage(content="a0", id="ai-0")], None, 0)] - reader = _mock_reader( - monkeypatch, - anchors=[_anchor(0, None, 0)], - history=ThreadHistory(thread_id=THREAD, turns=turns), - ) - - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) - - reader.aget_thread_history.assert_called_once() - assert [i["event"] for i in items] == ["user_message", "message_chunk"] - assert fake_cache.store == {} # no tail → nothing cached - - -async def test_refresh_writes_last_two_turn_entries(monkeypatch, fake_cache): - # The post-persist refresh rebuilds a two-turn window: the just-persisted - # turn's entry is new, and the PREVIOUS turn's entry is overwritten under - # its unchanged tail key — that is how an interrupted turn's - # ending_interrupts (known only at the resume boundary) reach the cache. - turns = [ - _turn(0, [HumanMessage(content="hello", id="h-0"), - AIMessage(content="a0", id="ai-0")], "tail-0", 0), - _turn(1, [HumanMessage(content="hello", id="h-1"), - AIMessage(content="a1", id="ai-1")], "tail-1", 1), - ] - _mock_reader( - monkeypatch, - anchors=[_anchor(0, "tail-0", 0), _anchor(1, "tail-1", 1)], - tip="tail-1", - history=ThreadHistory(thread_id=THREAD, turns=turns), - ) - # Stale previous-turn entry that the refresh must overwrite. - fake_cache.store[_key("tail-0", _response(0))] = [ - {"event": "user_message", "data": {"stale": True}} - ] - monkeypatch.setattr( - "src.server.database.conversation.get_replay_thread_data", - AsyncMock(return_value=( - "owner", - {"latest_checkpoint_id": "tail-1"}, - [_query(0), _query(1)], - [_response(0), _response(1)], - [], - [], - )), - ) - - await projection_cache.refresh_thread_projection(THREAD) - - for turn_index, tail in enumerate(("tail-0", "tail-1")): - entry = fake_cache.store[_key(tail, _response(turn_index))] - assert [i["event"] for i in entry] == ["user_message", "message_chunk"] - assert "stale" not in str(fake_cache.store[_key("tail-0", _response(0))]) - - -async def test_refresh_noops_without_commit_pointer(monkeypatch, fake_cache): - monkeypatch.setattr( - "src.server.database.conversation.get_replay_thread_data", - AsyncMock(return_value=("owner", {"latest_checkpoint_id": None}, [], [], [], [])), - ) - await projection_cache.refresh_thread_projection(THREAD) - assert fake_cache.store == {} - - -async def test_schedule_refresh_coalesces_overlapping_requests(monkeypatch, fake_cache): - first_started = asyncio.Event() - release_first = asyncio.Event() - second_started = asyncio.Event() - calls: list[int] = [] - active = 0 - max_active = 0 - - async def fake_refresh(thread_id): - nonlocal active, max_active - assert thread_id == THREAD - call_number = len(calls) + 1 - calls.append(call_number) - active += 1 - max_active = max(max_active, active) - try: - if call_number == 1: - first_started.set() - await release_first.wait() - else: - second_started.set() - finally: - active -= 1 - - monkeypatch.setattr(projection_cache, "refresh_thread_projection", fake_refresh) - projection_cache._refresh_tasks.clear() - projection_cache._refresh_dirty.clear() - - projection_cache.schedule_projection_refresh(THREAD) - runner = projection_cache._refresh_tasks[THREAD] - await first_started.wait() - - # Multiple triggers during the first pass coalesce into one ordered rerun; - # no second task is allowed to race and finish out of generation order. - projection_cache.schedule_projection_refresh(THREAD) - projection_cache.schedule_projection_refresh(THREAD) - assert projection_cache._refresh_tasks[THREAD] is runner - assert projection_cache._refresh_dirty == {THREAD} - - release_first.set() - await second_started.wait() - await runner - await asyncio.sleep(0) # let the done callback release the strong ref - - assert calls == [1, 2] - assert max_active == 1 - assert THREAD not in projection_cache._refresh_tasks - assert THREAD not in projection_cache._refresh_dirty - - -async def test_live_task_streams_variants(fake_cache): - key = f"subagent:stream:{THREAD}:tsk1" - # No stream cannot prove completion (it may have TTL'd or failed to spill), - # so stay conservative and skip caching. - assert await projection_cache.live_task_streams(THREAD, {"tsk1"}) == {"tsk1"} - # Last entry is a wire string → producer still writing. - fake_cache.client.stream_tails[key] = "id: 7\nevent: message_chunk\ndata: {}\n\n" - assert await projection_cache.live_task_streams(THREAD, {"tsk1"}) == {"tsk1"} - # Finalized sentinel → terminal. - fake_cache.client.stream_tails[key] = _END_SENTINEL - assert await projection_cache.live_task_streams(THREAD, {"tsk1"}) == set() - # Every referenced task needs an explicit sentinel — only the unproven - # one reports live. - assert await projection_cache.live_task_streams(THREAD, {"tsk1", "missing"}) == { - "missing" - } - # No tasks → trivially terminal; Redis failures/disconnects → conservative live. - assert await projection_cache.live_task_streams(THREAD, set()) == set() - fake_cache.client.xrevrange = AsyncMock(side_effect=RuntimeError("redis down")) - assert await projection_cache.live_task_streams(THREAD, {"tsk1"}) == {"tsk1"} - fake_cache.client = None - assert await projection_cache.live_task_streams(THREAD, {"tsk1"}) == {"tsk1"} - - -async def test_live_subagent_turn_not_cached_until_finalized(monkeypatch, fake_cache): - # Tail mode: the dispatch turn persists while the subagent still writes. - # Its transcript is projected from task-ns state that keeps growing AFTER - # this turn's tail is fixed — caching it would freeze a partial transcript - # forever. The turn must skip the store until the task stream finalizes. - task_artifact = {"task_id": "tsk1", "action": "init", "description": "d", "prompt": "p"} - turn_msgs = [ - HumanMessage(content="hello", id="h-0"), - AIMessage(content="", id="ai-1", - tool_calls=[{"name": "Task", "args": {}, "id": "tc-t"}]), - ToolMessage(content="dispatched", tool_call_id="tc-t", name="Task", - id="tm-1", additional_kwargs={"task_artifact": task_artifact}), - ] - history = ThreadHistory( - thread_id=THREAD, turns=[_turn(0, turn_msgs, "tail-0", 0)] - ) - reader = _mock_reader( - monkeypatch, anchors=[_anchor(0, "tail-0", 0)], tip="tail-0", history=history - ) - reader.aget_task_history = AsyncMock( - return_value=TaskHistory( - messages=[AIMessage(content="partial", id="sub-1")] - ) - ) - stream_key = f"subagent:stream:{THREAD}:tsk1" - fake_cache.client.stream_tails[stream_key] = "id: 3\nevent: message_chunk\ndata: {}\n\n" - - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) - assert any(i["data"].get("agent") == "task:tsk1" for i in items) - assert fake_cache.store == {} # live stream → nothing frozen - - # Subagent finishes (sentinel lands); the transcript is final now. - fake_cache.client.stream_tails[stream_key] = _END_SENTINEL - reader.aget_task_history = AsyncMock( - return_value=TaskHistory( - messages=[ - AIMessage(content="partial", id="sub-1"), - AIMessage(content="final", id="sub-2"), - ] - ) - ) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) - entry = fake_cache.store[_key("tail-0", _response(0))] - assert sum(1 for i in entry if i["data"].get("agent") == "task:tsk1") == 2 - - -async def test_windowed_fast_path_keeps_absolute_pairing(monkeypatch, fake_cache): - entry = [{"event": "user_message", "data": {"turn_index": 2}}] - fake_cache.store[_key("tail-2", _response(2))] = entry - reader = _mock_reader( - monkeypatch, - anchors=[_anchor(0, "tail-0", 0), _anchor(1, "tail-1", 1), _anchor(2, "tail-2", 2)], - tip="tail-2", - ) - - items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1), _query(2)], - {i: _response(i) for i in range(3)}, - last_n_turns=1, - ) - - assert items == entry - reader.aget_recent_history.assert_not_called() diff --git a/tests/unit/server/services/history/test_projection_sources.py b/tests/unit/server/services/history/test_projection_sources.py new file mode 100644 index 000000000..d5e1d9959 --- /dev/null +++ b/tests/unit/server/services/history/test_projection_sources.py @@ -0,0 +1,228 @@ +"""The projection and slice digests cover the code they key. + +Stored lines carry no expiry, so a module that shapes them but sits outside +``projection_sources()`` would let a deploy keep serving what the old code +projected; a stored slice is trusted while its key matches, so code that +decides its bytes outside ``slices.slice_sources()`` would leave the slices +cut by the old rule. Every import of ours a digested module makes is +digested too, or excluded here with the reason its code cannot change what +the stored rows hold. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +from src.server.services.history import slices +from src.server.services.history.replay import keys + +ROOT = keys._SRC.parent +ANY = frozenset({"*"}) + +# Module -> (names a digested module may import from it, why its code cannot +# change stored lines). ANY admits every name; a narrower set flags the next +# name imported, which may well shape lines. +EXCLUDED: dict[str, tuple[frozenset[str], str]] = { + "src/server/services/history/reader.py": ( + ANY, + "Finds the checkpoints slices are cut between, whose ids a stored " + "row is checked against.", + ), + "src/server/services/history/slices.py": ( + ANY, + "Cuts and encodes the slices lines are projected from, digested into " + "the slice key, which the lines key carries.", + ), + "src/server/database/conversation/turn_slices.py": ( + ANY, + "Stores and reads the slices and lines themselves.", + ), + "src/server/database/pool.py": (ANY, "Connections."), + "src/server/database/runs/subagent_runs.py": ( + ANY, + "The run ledger, read as data: a run's status decides whether lines " + "are stored; the statuses themselves are in contracts/status.py.", + ), + "src/server/database/conversation/replay_rows.py": ( + frozenset({"ThreadRows"}), + "The rows a projection reads, whose versions are in the lines key.", + ), + "src/server/database/replay_facts.py": ( + frozenset({"LEGACY_KEY"}), + "Names the key legacy facts sit under: those resolve the row's stored " + "events, so where they sit cannot change what lines hold.", + ), + "src/server/utils/pg_sanitize.py": ( + frozenset({"SafeJson", "strip_pg_nul_str"}), + "Binds provenance rows on write; the projection reads them back.", + ), + "src/server/services/history/task_streams.py": ( + ANY, + "Whether a task's stream sealed, which decides whether lines are " + "stored, never what they hold.", + ), + "src/config/settings.py": ( + frozenset({"get_redis_ttl_workflow_events"}), + "How long a sealed stream outlives its task, same decision.", + ), + "src/server/services/subagent_liveness.py": ( + ANY, + "A task card's live status, stamped at read time (``finish_lines``).", + ), + "src/utils/storage/__init__.py": ( + frozenset({"get_bytes"}), + "Offloaded widget data, inlined at read time.", + ), + "src/server/services/history/replay/errors.py": ( + ANY, + "Error types: why a read gives up, never what a line holds.", + ), + "src/server/services/runs/sse_producer.py": ( + frozenset({"resolve_token_threshold"}), + "Its value, not its code, is in the lines key.", + ), + "src/llms/pricing_utils.py": ( + frozenset({"calculate_total_cost"}), + "token_counter's cost tracking; the projection calls only " + "extract_token_usage, which reads neither.", + ), + "src/llms/llm.py": ( + frozenset({"LLM"}), + "token_counter's cost tracking; the projection calls only " + "extract_token_usage, which reads neither.", + ), +} + +# The same, for the slice key's sources. +SLICE_EXCLUDED: dict[str, tuple[frozenset[str], str]] = { + "src/server/services/history/reader.py": ( + frozenset({"TurnAnchor", "CheckpointHistoryReader"}), + "Type annotations; a turn slice carries its anchor, and the cut " + "reads only whether it is a resume, a fact of the walk.", + ), + "src/server/database/conversation/turn_slices.py": ( + frozenset({"StoredSlice"}), + "A type annotation: the stored row a slice is read back from.", + ), +} + + +def _module_file(module: str) -> Path | None: + # ptc_agent is a top-level package that lives under src/. + parts = module.split(".") + base = ROOT.joinpath(*parts) if parts[0] == "src" else ROOT.joinpath("src", *parts) + for candidate in (base.with_suffix(".py"), base / "__init__.py"): + if candidate.exists(): + return candidate + return None + + +def _package_of(path: Path, level: int) -> str: + parts = list(path.relative_to(ROOT).with_suffix("").parts) + package = parts[:-1] + return ".".join(package[: len(package) - (level - 1)]) + + +def _imports(path: Path, roots: tuple[str, ...]) -> list[tuple[Path, str]]: + """``(module file, name)`` for every import in ``path`` from a package + under ``roots``; the name is ``*`` when the import is of the module + itself.""" + out: list[tuple[Path, str]] = [] + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name.split(".")[0] in roots: + found = _module_file(alias.name) + assert found, f"{path}: unresolved import {alias.name}" + out.append((found, "*")) + elif isinstance(node, ast.ImportFrom): + module = node.module or "" + if node.level: + package = _package_of(path, node.level) + module = f"{package}.{module}" if module else package + if module.split(".")[0] not in roots: + continue + for alias in node.names: + submodule = _module_file(f"{module}.{alias.name}") + if submodule is not None: + out.append((submodule, "*")) + continue + found = _module_file(module) + assert found, f"{path}: unresolved import {module}" + out.append((found, alias.name)) + return out + + +def _unaccounted( + digested: list[Path], + excluded: dict[str, tuple[frozenset[str], str]], + roots: tuple[str, ...], +) -> list[str]: + out: list[str] = [] + for path in sorted(digested): + for target, name in _imports(path, roots): + if target in digested: + continue + key = str(target.relative_to(ROOT)) + names, _why = excluded.get(key, (frozenset(), "")) + if names is not ANY and name not in names: + out.append(f"{path.relative_to(ROOT)} imports {key}:{name}") + return out + + +def _imported(digested: list[Path], roots: tuple[str, ...]) -> set[str]: + return { + str(target.relative_to(ROOT)) + for path in digested + for target, _name in _imports(path, roots) + } + + +def test_every_src_import_of_the_projection_is_digested_or_excluded(): + unaccounted = _unaccounted(keys.projection_sources(), EXCLUDED, ("src",)) + assert not unaccounted, ( + "Add each module to keys.projection_sources(), or to EXCLUDED with " + "why its code cannot change stored lines:\n" + "\n".join(unaccounted) + ) + + +def test_every_exclusion_is_still_imported(): + """An exclusion nothing imports any more would quietly admit the module + when something imports it again for a reason that does shape lines.""" + imported = _imported(keys.projection_sources(), ("src",)) + assert sorted(set(EXCLUDED) - imported) == [] + + +# The slice key follows ptc_agent imports too: the cut calls the transcript +# classifier, which no lines version would catch. +_SLICE_ROOTS = ("src", "ptc_agent") + + +def test_every_import_of_the_slice_cut_is_digested_or_excluded(): + unaccounted = _unaccounted(slices.slice_sources(), SLICE_EXCLUDED, _SLICE_ROOTS) + assert not unaccounted, ( + "Add each module to slices.slice_sources(), or to SLICE_EXCLUDED " + "with why its code cannot change a slice's bytes:\n" + + "\n".join(unaccounted) + ) + + +def test_a_dropped_slice_dependency_is_caught(): + """The cut finds a span's input by the transcript classifier, which the + key once left out: a change there moved what slices are cut without + moving the key.""" + without = [p for p in slices.slice_sources() if p.name != "classify.py"] + + assert any( + line.startswith( + "src/server/services/history/slices.py imports " + "src/ptc_agent/agent/transcript/classify.py" + ) + for line in _unaccounted(without, SLICE_EXCLUDED, _SLICE_ROOTS) + ) + + +def test_every_slice_exclusion_is_still_imported(): + imported = _imported(slices.slice_sources(), _SLICE_ROOTS) + assert sorted(set(SLICE_EXCLUDED) - imported) == [] diff --git a/tests/unit/server/services/history/test_reader.py b/tests/unit/server/services/history/test_reader.py index f190e5064..4adb225bc 100644 --- a/tests/unit/server/services/history/test_reader.py +++ b/tests/unit/server/services/history/test_reader.py @@ -14,66 +14,33 @@ from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph from langgraph.graph.message import REMOVE_ALL_MESSAGES -from langgraph.graph.ui import ui_message_reducer from langgraph.types import Command, interrupt from ptc_agent.agent.middleware.compaction.types import CompactionState from ptc_agent.agent.state import DeltaAgentState -from src.server.services.history.reader import ( - CheckpointHistoryReader, - TaskHistory, - _ReaderState, -) +from ptc_agent.agent.main_state import MainAgentState +from ptc_agent.agent.transcript.classify import is_run_boundary_message +from src.server.services.history.reader import CheckpointHistoryReader, TaskHistory from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound +from tests.unit.server.services.history.reader_graphs import ( + THREAD, + _cfg, + _echo_graph, + _history, + _run_turns, +) pytestmark = pytest.mark.asyncio -THREAD = "thread-1" - # Background subagents are invoked from inside a parent tool context, whose # config carries the pregel task id — that is what makes langgraph honor the # explicit checkpoint_ns instead of resetting it to root. Tests must mirror it. CONFIG_KEY_TASK_ID = "__pregel_task_id" -def _echo_graph(checkpointer): - """One turn = reply 'echo: ' with a stable per-turn id.""" - - def agent(state): - humans = [m for m in state["messages"] if isinstance(m, HumanMessage)] - last = humans[-1].content if humans else "?" - return { - "messages": [AIMessage(content=f"echo: {last}", id=f"ai-{len(state['messages'])}")] - } - - return ( - StateGraph(DeltaAgentState) - .add_node("agent", agent) - .add_edge(START, "agent") - .compile(checkpointer=checkpointer) - ) - - -def _cfg(thread_id=THREAD, *, turn_index=None, run_id=None, checkpoint_id=None): - cfg = {"configurable": {"thread_id": thread_id}} - if checkpoint_id: - cfg["configurable"]["checkpoint_id"] = checkpoint_id - metadata = {} - if run_id is not None: - metadata["run_id"] = run_id - if turn_index is not None: - metadata["turn_index"] = turn_index - if metadata: - cfg["metadata"] = metadata - return cfg - - -async def _run_turns(graph, n, start=0): - for i in range(start, start + n): - await graph.ainvoke( - {"messages": [HumanMessage(content=f"q{i}", id=f"h-{i}")]}, - _cfg(turn_index=i, run_id=f"run-{i}"), - ) +def _user(turn): + """The turn's own input: the first run boundary among what it added.""" + return next((m for m in turn.messages if is_run_boundary_message(m)), None) async def test_turn_slicing_and_metadata(): @@ -82,19 +49,19 @@ async def test_turn_slicing_and_metadata(): await _run_turns(graph, 3) reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) assert len(history.turns) == 3 for i, turn in enumerate(history.turns): - assert turn.turn_ordinal == i - assert turn.user_message is not None - assert turn.user_message.content == f"q{i}" + assert turn.anchor.turn_ordinal == i + assert _user(turn) is not None + assert _user(turn).content == f"q{i}" # The slice starts at the turn's own input (the input checkpoint's # state predates it) and ends before the next turn's input. contents = [m.content for m in turn.messages] assert contents == [f"q{i}", f"echo: q{i}"] - assert turn.run_id == f"run-{i}" - assert turn.turn_index == i + assert turn.anchor.run_id == f"run-{i}" + assert turn.anchor.turn_index == i assert history.interrupts == [] @@ -105,12 +72,12 @@ async def test_edit_branch_follows_requested_tip(): original_tip = (await graph.aget_state(_cfg())).config["configurable"]["checkpoint_id"] reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD, original_tip) + history = await _history(reader, tip=original_tip) # Production edit forks from the checkpoint BEFORE the turn's input # (checkpoint_handler's edit_checkpoint_id = the input checkpoint's parent), # so the stale input boundary is off the new branch. - turn2_input = history.turns[2].input_checkpoint_id + turn2_input = history.turns[2].anchor.input_checkpoint_id input_tuple = await saver.aget_tuple( {"configurable": {"thread_id": THREAD, "checkpoint_id": turn2_input}} ) @@ -123,12 +90,12 @@ async def test_edit_branch_follows_requested_tip(): new_tip = (await graph.aget_state(_cfg())).config["configurable"]["checkpoint_id"] assert new_tip != original_tip - branched = await reader.aget_thread_history(THREAD, new_tip) - assert [t.user_message.content for t in branched.turns] == ["q0", "q1", "q2-edited"] + branched = await _history(reader, tip=new_tip) + assert [_user(t).content for t in branched.turns] == ["q0", "q1", "q2-edited"] assert branched.turns[2].messages[-1].content == "echo: q2-edited" - original = await reader.aget_thread_history(THREAD, original_tip) - assert [t.user_message.content for t in original.turns] == ["q0", "q1", "q2"] + original = await _history(reader, tip=original_tip) + assert [_user(t).content for t in original.turns] == ["q0", "q1", "q2"] async def test_missing_requested_tip_fails_instead_of_reading_newest(): @@ -138,7 +105,7 @@ async def test_missing_requested_tip_fails_instead_of_reading_newest(): reader = CheckpointHistoryReader(saver) with pytest.raises(CheckpointBranchTipNotFound, match="missing-tip"): - await reader.aget_thread_history(THREAD, "missing-tip") + await _history(reader, tip="missing-tip") async def test_regenerate_branch_reuses_input_boundary(): @@ -147,16 +114,16 @@ async def test_regenerate_branch_reuses_input_boundary(): await _run_turns(graph, 2) reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) # Production regenerate re-runs FROM the input checkpoint (input=None # replays its pending input writes) — same boundary, new branch below it. await graph.ainvoke( - None, _cfg(checkpoint_id=history.turns[1].input_checkpoint_id) + None, _cfg(checkpoint_id=history.turns[1].anchor.input_checkpoint_id) ) - regenerated = await reader.aget_thread_history(THREAD) + regenerated = await _history(reader) assert len(regenerated.turns) == 2 - assert regenerated.turns[1].user_message.content == "q1" + assert _user(regenerated.turns[1]).content == "q1" assert [m.content for m in regenerated.turns[1].messages] == ["q1", "echo: q1"] @@ -185,7 +152,7 @@ def agent(state): await graph.ainvoke({"messages": [HumanMessage(content="q1", id="h-1")]}, _cfg()) reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) assert len(history.turns) == 2 assert [m.id for m in history.turns[0].messages] == ["h-0", "ai-first"] @@ -194,7 +161,7 @@ def agent(state): # there is no HumanMessage left to attribute (replay sources user text from # DB query rows, not from here). assert [m.id for m in history.turns[1].messages] == ["sum-1", "ai-fresh"] - assert history.turns[1].user_message is None + assert _user(history.turns[1]) is None async def test_hitl_resume_is_its_own_turn_boundary(): @@ -216,7 +183,7 @@ def agent(state): ) reader = CheckpointHistoryReader(saver) - pending = await reader.aget_thread_history(THREAD) + pending = await _history(reader) # A pending interrupt (__interrupt__ writes, no __resume__) is not a boundary. assert len(pending.turns) == 1 assert len(pending.interrupts) == 1 @@ -224,28 +191,28 @@ def agent(state): "action_requests": [{"description": "proceed?"}] } # Pending, not answered — it must not double as an ending interrupt. - assert pending.turns[0].ending_interrupts == [] + assert pending.turns[0].anchor.ending_interrupts == [] await graph.ainvoke(Command(resume="yes"), _cfg()) - resumed = await reader.aget_thread_history(THREAD) + resumed = await _history(reader) # The resume is a boundary of its own — it persists a resume_feedback # query row, so boundaries must stay 1:1 with persisted turns. assert len(resumed.turns) == 2 assert [m.content for m in resumed.turns[0].messages] == ["q0"] assert [m.content for m in resumed.turns[1].messages] == ["resumed: yes"] - assert resumed.turns[1].user_message is None + assert _user(resumed.turns[1]) is None # The resume checkpoint's metadata belongs to the interrupted run — the # resume turn must not inherit its run_id/turn_index. - assert resumed.turns[1].run_id is None - assert resumed.turns[1].turn_index is None + assert resumed.turns[1].anchor.run_id is None + assert resumed.turns[1].anchor.turn_index is None assert resumed.interrupts == [] # Once answered, the interrupt attributes to the turn that raised it — # read from the resume boundary's __interrupt__ writes. - assert [i["value"] for i in resumed.turns[0].ending_interrupts] == [ + assert [i["value"] for i in resumed.turns[0].anchor.ending_interrupts] == [ {"action_requests": [{"description": "proceed?"}]} ] - assert resumed.turns[0].ending_interrupts[0]["id"] is not None - assert resumed.turns[1].ending_interrupts == [] + assert resumed.turns[0].anchor.ending_interrupts[0]["id"] is not None + assert resumed.turns[1].anchor.ending_interrupts == [] async def test_task_namespace_transcript(): @@ -270,7 +237,7 @@ async def test_task_namespace_transcript(): } sub_graph = ( - StateGraph(_ReaderState) + StateGraph(MainAgentState) .add_node( "agent", lambda s: { @@ -308,7 +275,7 @@ async def test_task_namespace_transcript(): assert task_history.new_ui_records == [task_ui] # The subagent namespace does not leak turn boundaries into the main thread. - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) assert len(history.turns) == 1 @@ -321,7 +288,7 @@ async def test_project_task_transcript_serves_unprojected_lanes(monkeypatch): graph = _echo_graph(saver) await _run_turns(graph, 1) sub_graph = ( - StateGraph(_ReaderState) + StateGraph(MainAgentState) .add_node( "agent", lambda s: { @@ -345,8 +312,8 @@ async def test_project_task_transcript_serves_unprojected_lanes(monkeypatch): monkeypatch.setattr( CheckpointHistoryReader, "get_instance", classmethod(lambda cls: reader) ) - from src.server.services.history.replay import task_lane - from src.server.services.history.replay.task_lane import ( + from src.server.services.history.replay import run_lane as task_lane + from src.server.services.history.replay.run_lane import ( project_task_transcript, ) @@ -383,7 +350,7 @@ async def test_project_task_transcript_gates_on_a_settled_task(monkeypatch): graph = _echo_graph(saver) await _run_turns(graph, 1) sub_graph = ( - StateGraph(_ReaderState) + StateGraph(MainAgentState) .add_node( "agent", lambda s: {"messages": [AIMessage(content="partial", id="sub-ai-1")]}, @@ -405,7 +372,7 @@ async def test_project_task_transcript_gates_on_a_settled_task(monkeypatch): monkeypatch.setattr( CheckpointHistoryReader, "get_instance", classmethod(lambda cls: reader) ) - from src.server.services.history.replay import task_lane + from src.server.services.history.replay import run_lane as task_lane monkeypatch.setattr( task_lane, @@ -423,268 +390,6 @@ async def test_project_task_transcript_gates_on_a_settled_task(monkeypatch): assert await task_lane.project_task_transcript(THREAD, "wf1") == [] -async def test_append_ui_record_no_new_boundary(): - saver = InMemorySaver() - graph = _echo_graph(saver) - await _run_turns(graph, 2) - - reader = CheckpointHistoryReader(saver) - await reader.append_ui_record(THREAD, "image_capture", {"path_to_url": {"a.png": "https://x/a"}}) - await reader.append_ui_record(THREAD, "image_capture", {"path_to_url": {"b.png": "https://x/b"}}) - - history = await reader.aget_thread_history(THREAD) - assert len(history.turns) == 2 # updates are not input boundaries - assert [r["name"] for r in history.ui] == ["image_capture", "image_capture"] - assert history.ui[0]["props"] == {"path_to_url": {"a.png": "https://x/a"}} - assert history.ui[1]["props"] == {"path_to_url": {"b.png": "https://x/b"}} - assert all(r["id"] for r in history.ui) - - # A ui append after the last turn must not shift the turn's message slice. - assert [m.content for m in history.turns[1].messages] == ["q1", "echo: q1"] - - -async def test_append_ui_record_skipped_on_interrupted_tip(): - # An interrupted turn's tip carries a pending __interrupt__ write. Appending - # a ui record there via aupdate_state attributes the write to the reader - # graph's node and clears the pending interrupt, silently breaking HITL - # resume (the image-capture Hook B fallback fires at interrupt time, so this - # is reachable when a subagent emits a sandbox image on an interrupted turn). - # The append must skip while the tip is interrupted. - saver = InMemorySaver() - - def agent(state): - humans = [m for m in state["messages"] if isinstance(m, HumanMessage)] - if len(humans) == 1: - answer = interrupt({"action_requests": [{"description": "go?"}]}) - return {"messages": [AIMessage(content=f"resumed: {answer}", id="ai-r")]} - return {"messages": [AIMessage(content="ok", id=f"ai-{len(state['messages'])}")]} - - graph = ( - StateGraph(DeltaAgentState) - .add_node("agent", agent) - .add_edge(START, "agent") - .compile(checkpointer=saver) - ) - await graph.ainvoke({"messages": [HumanMessage(content="q0", id="h-0")]}, _cfg()) - - reader = CheckpointHistoryReader(saver) - assert (await graph.aget_state(_cfg())).next == ("agent",) - - await reader.append_ui_record( - THREAD, "image_capture", {"path_to_url": {"a.png": "https://x/a"}} - ) - - # Interrupt survives the append: the pending task is intact and resume runs. - assert (await graph.aget_state(_cfg())).next == ("agent",) - result = await graph.ainvoke(Command(resume="yes"), _cfg()) - assert any( - isinstance(m, AIMessage) and m.content == "resumed: yes" - for m in result["messages"] - ) - # The skipped record did not land. - history = await reader.aget_thread_history(THREAD) - assert history.ui == [] - - -async def test_append_ui_record_skipped_when_interrupt_check_fails(monkeypatch): - reader = CheckpointHistoryReader(InMemorySaver()) - interrupt_check = AsyncMock(side_effect=RuntimeError("checkpoint unavailable")) - update = AsyncMock() - monkeypatch.setattr(reader._checkpointer, "aget_tuple", interrupt_check) - monkeypatch.setattr(reader._updater, "aupdate_state", update) - - await reader.append_ui_record( - THREAD, "image_capture", {"path_to_url": {"a.png": "https://x/a"}} - ) - - interrupt_check.assert_awaited_once() - update.assert_not_awaited() - - -async def test_append_ui_record_advances_recorded_branch_tip(monkeypatch): - # The image-capture hook appends after turn end, creating a checkpoint - # beyond the recorded branch tip where the replay walk cannot see it. The - # append must CAS the recorded tip from the checkpoint it built on onto - # the new one. (The workflow terminal snapshot does NOT come through here - # — it writes into task:{id} via persist_task_ui_record, below.) - saver = InMemorySaver() - graph = _echo_graph(saver) - await _run_turns(graph, 1) - reader = CheckpointHistoryReader(saver) - - advance = AsyncMock(return_value=True) - monkeypatch.setattr( - "src.server.database.conversation.threads_write." - "advance_thread_checkpoint_id", - advance, - ) - root_cfg = {"configurable": {"thread_id": THREAD}} - tip_before = await saver.aget_tuple(root_cfg) - - await reader.append_ui_record( - THREAD, "image_capture", {"path_to_url": {"chart.png": "https://x/chart.png"}} - ) - - tip_after = await saver.aget_tuple(root_cfg) - assert advance.await_args.args[0] == THREAD - kwargs = advance.await_args.kwargs - assert ( - kwargs["from_checkpoint_id"] - == tip_before.config["configurable"]["checkpoint_id"] - ) - assert ( - kwargs["to_checkpoint_id"] - == tip_after.config["configurable"]["checkpoint_id"] - ) - assert kwargs["to_checkpoint_id"] != kwargs["from_checkpoint_id"] - - -async def test_append_ui_record_anchors_to_the_tip_it_read(monkeypatch): - # A turn starting on another worker between the tip read and the append - # must not re-parent the record onto that turn's uncommitted checkpoint: - # the CAS guard is the tip that was read, so an unanchored write would - # still pass it and publish partial state as the thread's commit pointer. - saver = InMemorySaver() - graph = _echo_graph(saver) - await _run_turns(graph, 1) - reader = CheckpointHistoryReader(saver) - root_cfg = {"configurable": {"thread_id": THREAD}} - stale_tip = await saver.aget_tuple(root_cfg) - stale_id = stale_tip.config["configurable"]["checkpoint_id"] - - # The concurrent turn lands. It has not finalized, so the recorded pointer - # still names the turn-1 tip — which is exactly why the CAS would pass. - await _run_turns(graph, 1, start=1) - moved = await saver.aget_tuple(root_cfg) - moved_id = moved.config["configurable"]["checkpoint_id"] - assert moved_id != stale_id - - class _StaleTip: - """Models the read/write window: the reader holds the older tip.""" - - async def aget_tuple(self, config): - return stale_tip - - monkeypatch.setattr(reader, "_checkpointer", _StaleTip()) - advance = AsyncMock(return_value=True) - monkeypatch.setattr( - "src.server.database.conversation.threads_write." - "advance_thread_checkpoint_id", - advance, - ) - - await reader.append_ui_record(THREAD, "image_capture", {"path_to_url": {}}) - - kwargs = advance.await_args.kwargs - assert kwargs["from_checkpoint_id"] == stale_id - written = await saver.aget_tuple( - { - "configurable": { - "thread_id": THREAD, - "checkpoint_id": kwargs["to_checkpoint_id"], - } - } - ) - # The guard and the write agree on one parent — the whole point. - assert written.parent_config["configurable"]["checkpoint_id"] == stale_id - - -async def test_task_ui_snapshot_readable_through_reader(): - # Cross-layer contract: the workflow driver persists its terminal ui - # snapshot middleware-side (through the run's own checkpointer, so the - # writer-guard fence applies), and the server reader materializes it. The - # snapshot may land while the launching turn is still executing — a - # root-ns append there forks a dead branch — so the writer must (a) - # surface via aget_task_history, (b) upsert by record id, and (c) leave - # the root chain untouched. - from ptc_agent.agent.middleware.background_subagent.workflow.ui_snapshot import ( - persist_task_ui_record, - ) - - saver = InMemorySaver() - graph = _echo_graph(saver) - await _run_turns(graph, 1) - reader = CheckpointHistoryReader(saver) - root_cfg = {"configurable": {"thread_id": THREAD}} - tip_before = await saver.aget_tuple(root_cfg) - - await persist_task_ui_record( - saver, - THREAD, - "wf1", - "workflow_run", - {"task_id": "wf1", "frames": [{"phase": "run_started"}]}, - record_id="workflow-run-r1", - ) - await persist_task_ui_record( - saver, - THREAD, - "wf1", - "workflow_run", - { - "task_id": "wf1", - "frames": [{"phase": "run_started"}, {"phase": "run_completed"}], - }, - record_id="workflow-run-r1", - ) - - task_history = await reader.aget_task_history(THREAD, "wf1") - records = [ - r for r in task_history.new_ui_records if r["id"] == "workflow-run-r1" - ] - assert len(records) == 1 # same-id rewrite upserts, never duplicates - assert [f["phase"] for f in records[0]["props"]["frames"]] == [ - "run_started", - "run_completed", - ] - - # Root chain untouched: same tip, no root ui records, no new boundary. - tip_after = await saver.aget_tuple(root_cfg) - assert ( - tip_after.config["configurable"]["checkpoint_id"] - == tip_before.config["configurable"]["checkpoint_id"] - ) - history = await reader.aget_thread_history(THREAD) - assert history.ui == [] - assert len(history.turns) == 1 - - -async def test_new_ui_records_attributed_to_their_turn(): - # push_ui_message inside a node (the resilience middleware's fallback - # path) checkpoints the record on the ui channel; the reader id-diffs the - # channel per turn so replay can project the notice into the right turn. - from langgraph.graph.ui import push_ui_message - - saver = InMemorySaver() - - def agent(state): - humans = [m for m in state["messages"] if isinstance(m, HumanMessage)] - if len(humans) == 2: # second turn only - push_ui_message( - name="model_fallback", - props={"from_model": "primary", "to_model": "backup"}, - id="ui-fb-1", - ) - return { - "messages": [AIMessage(content="ok", id=f"ai-{len(state['messages'])}")] - } - - graph = ( - StateGraph(DeltaAgentState) - .add_node("agent", agent) - .add_edge(START, "agent") - .compile(checkpointer=saver) - ) - await _run_turns(graph, 2) - - reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) - assert [len(t.new_ui_records) for t in history.turns] == [0, 1] - record = history.turns[1].new_ui_records[0] - assert record["name"] == "model_fallback" - assert record["props"] == {"from_model": "primary", "to_model": "backup"} - - async def test_tail_checkpoint_ids_and_anchors(): # Each turn's tail = its own last checkpoint (the tip the persist path # records) — the projection-cache key. Anchors derive the same tails from @@ -695,10 +400,10 @@ async def test_tail_checkpoint_ids_and_anchors(): await _run_turns(graph, 3) reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) anchors, tip_id = await reader.aget_turn_anchors(THREAD) - tails = [t.tail_checkpoint_id for t in history.turns] + tails = [t.anchor.tail_checkpoint_id for t in history.turns] assert all(tails) and len(set(tails)) == 3 assert [a.tail_checkpoint_id for a in anchors] == tails assert anchors[-1].tail_checkpoint_id == tip_id @@ -706,8 +411,8 @@ async def test_tail_checkpoint_ids_and_anchors(): # A non-last turn's tail is the next input boundary's parent — strictly # between the two boundaries (checkpoint ids are time-ordered). for i in range(2): - assert history.turns[i].input_checkpoint_id < tails[i] - assert tails[i] < history.turns[i + 1].input_checkpoint_id + assert history.turns[i].anchor.input_checkpoint_id < tails[i] + assert tails[i] < history.turns[i + 1].anchor.input_checkpoint_id # The tail is what the persist path records: run one more turn and the # previous tip becomes... unchanged history for turns 0-2. @@ -742,10 +447,10 @@ def agent(state): await graph.ainvoke(Command(resume="yes"), _cfg()) reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) assert len(history.turns) == 2 - assert history.turns[1].input_checkpoint_id == interrupt_tip - assert history.turns[0].tail_checkpoint_id == interrupt_tip + assert history.turns[1].anchor.input_checkpoint_id == interrupt_tip + assert history.turns[0].anchor.tail_checkpoint_id == interrupt_tip anchors, tip_id = await reader.aget_turn_anchors(THREAD) assert anchors[0].tail_checkpoint_id == interrupt_tip assert anchors[1].tail_checkpoint_id == tip_id @@ -800,11 +505,11 @@ async def test_resume_of_cancelled_calls_is_its_own_turn_boundary(): Command(resume={interrupt_id: "yes"}), _cfg(turn_index=1, run_id="run-1") ) - history = await CheckpointHistoryReader(saver).aget_thread_history(THREAD) + history = await _history(CheckpointHistoryReader(saver)) assert len(history.turns) == 2 - assert history.turns[1].input_checkpoint_id == interrupt_tip + assert history.turns[1].anchor.input_checkpoint_id == interrupt_tip assert [m.content for m in history.turns[1].messages] == ["Not run", "reply"] - assert [i["id"] for i in history.turns[0].ending_interrupts] == [interrupt_id] + assert [i["id"] for i in history.turns[0].anchor.ending_interrupts] == [interrupt_id] async def test_resume_by_an_agent_that_does_not_ask_is_its_own_turn_boundary(): @@ -817,9 +522,9 @@ async def test_resume_by_an_agent_that_does_not_ask_is_its_own_turn_boundary(): Command(resume={interrupt_id: "yes"}), _cfg(turn_index=1, run_id="run-1") ) - history = await CheckpointHistoryReader(saver).aget_thread_history(THREAD) + history = await _history(CheckpointHistoryReader(saver)) assert len(history.turns) == 2 - assert history.turns[1].input_checkpoint_id == interrupt_tip + assert history.turns[1].anchor.input_checkpoint_id == interrupt_tip assert [m.content for m in history.turns[1].messages] == ["refused", "reply"] @@ -837,7 +542,7 @@ async def test_interrupted_checkpoint_continued_by_its_own_run_is_no_boundary(): ) await graph.ainvoke(None, _cfg(run_id="run-0")) - history = await CheckpointHistoryReader(saver).aget_thread_history(THREAD) + history = await _history(CheckpointHistoryReader(saver)) assert len(history.turns) == 1 @@ -851,26 +556,15 @@ async def test_new_message_over_an_unanswered_interrupt_adds_one_boundary(): _cfg(turn_index=1, run_id="run-1"), ) - history = await CheckpointHistoryReader(saver).aget_thread_history(THREAD) - assert [t.turn_index for t in history.turns] == [0, 1] + history = await _history(CheckpointHistoryReader(saver)) + assert [t.anchor.turn_index for t in history.turns] == [0, 1] async def test_empty_thread(): reader = CheckpointHistoryReader(InMemorySaver()) - history = await reader.aget_thread_history("no-such-thread") - assert history.turns == [] - assert await reader.acount_input_boundaries("no-such-thread") == 0 + assert await reader.aget_turn_anchors("no-such-thread") == ([], None) -async def test_ui_reducer_upsert_and_unknown_remove_guard(): - a1 = {"type": "ui", "id": "u1", "name": "n", "props": {"v": 1}, "metadata": {}} - a2 = {"type": "ui", "id": "u1", "name": "n", "props": {"v": 2}, "metadata": {}} - merged = ui_message_reducer([a1], [a2]) - assert merged == [a2] # same id upserts, no duplicate - - with pytest.raises(ValueError): - ui_message_reducer([a1], [{"type": "remove-ui", "id": "missing"}]) - async def test_summarization_event_attributed_to_its_turn(): # Auto-compaction stores its summary in the _summarization_event state key @@ -907,7 +601,7 @@ def agent(state): await _run_turns(graph, 3) reader = CheckpointHistoryReader(saver) - history = await reader.aget_thread_history(THREAD) + history = await _history(reader) assert [t.new_summarization_event is not None for t in history.turns] == [ False, @@ -920,3 +614,121 @@ def agent(state): # turn that offloaded them, zero once the set stops growing. assert [t.newly_offloaded_args for t in history.turns] == [0, 2, 0] assert [t.newly_offloaded_reads for t in history.turns] == [0, 0, 0] + + +async def test_turn_slices_from_anchors_match_the_full_walk(): + """Reading chosen turns from their anchors yields exactly those turns as + a read of every turn builds them, compaction and answered interrupts + included.""" + saver = InMemorySaver() + state = {"n": 0} + + def agent(st): + state["n"] += 1 + if state["n"] == 2: + answer = interrupt({"action_requests": [{"description": "ok?"}]}) + return {"messages": [AIMessage(content=f"after {answer}", id="ai-hitl")]} + if state["n"] == 4: + return { + "messages": [ + RemoveMessage(id=REMOVE_ALL_MESSAGES), + AIMessage(content="kept", id="ai-kept"), + ] + } + return {"messages": [AIMessage(content=f"r{state['n']}", id=f"ai-{state['n']}")]} + + graph = ( + StateGraph(DeltaAgentState) + .add_node("agent", agent) + .add_edge(START, "agent") + .compile(checkpointer=saver) + ) + await graph.ainvoke({"messages": [HumanMessage(content="q0", id="h-0")]}, _cfg(turn_index=0)) + await graph.ainvoke({"messages": [HumanMessage(content="q1", id="h-1")]}, _cfg(turn_index=1)) + await graph.ainvoke(Command(resume="yes"), _cfg()) + await graph.ainvoke({"messages": [HumanMessage(content="q3", id="h-3")]}, _cfg(turn_index=3)) + + reader = CheckpointHistoryReader(saver) + anchors, _tip = await reader.aget_turn_anchors(THREAD) + assert len(anchors) == 4 + + def shape(t): + return ( + t.anchor.turn_ordinal, + t.anchor.input_checkpoint_id, + t.anchor.end_checkpoint_id, + t.anchor.tail_checkpoint_id, + t.anchor.turn_index, + [m.id for m in t.messages], + [i["value"] for i in t.anchor.ending_interrupts], + t.new_summarization_event, + t.newly_offloaded_args, + ) + + every = await reader.aget_turn_slices(THREAD, anchors) + # The resume reruns the asking node, and the last turn's REMOVE_ALL + # takes its own input with everything before it. + assert [[m.id for m in t.messages] for t in every] == [ + ["h-0", "ai-1"], + ["h-1"], + ["ai-3"], + ["ai-kept"], + ] + assert [i["value"] for i in every[1].anchor.ending_interrupts] == [ + {"action_requests": [{"description": "ok?"}]} + ] + picked = await reader.aget_turn_slices(THREAD, [anchors[1], anchors[3]]) + assert [shape(t) for t in picked] == [shape(every[1]), shape(every[3])] + + +async def test_task_runs_slice_by_their_own_boundaries(): + """Each run is the diff between its own boundary and the next run's, so a + later run that compacts the namespace does not erase an earlier run, and + each run carries only the signals that landed during it.""" + saver = InMemorySaver() + calls = {"n": 0} + + def agent(st): + calls["n"] += 1 + if calls["n"] == 1: + return {"messages": [AIMessage(content="first answer", id="sub-ai-1")]} + return { + "messages": [ + RemoveMessage(id=REMOVE_ALL_MESSAGES), + AIMessage(content="second answer", id="sub-ai-2"), + ], + "_offloaded_tool_call_ids": {"tc-9"}, + } + + sub_graph = ( + StateGraph(MainAgentState) + .add_node("agent", agent) + .add_edge(START, "agent") + .compile(checkpointer=saver) + ) + for n, run_id in ((1, "run-a"), (2, "run-b")): + await sub_graph.ainvoke( + {"messages": [HumanMessage(content=f"prompt {n}", id=f"sub-h-{n}")]}, + { + "configurable": { + "thread_id": THREAD, + "checkpoint_ns": "task:tsk1", + CONFIG_KEY_TASK_ID: "parent-task", + }, + "metadata": {"task_run_id": run_id}, + }, + ) + + reader = CheckpointHistoryReader(saver) + runs = await reader.aget_task_runs(THREAD, "tsk1") + assert [(r.ordinal, r.task_run_id) for r in runs] == [(0, "run-a"), (1, "run-b")] + assert runs[0].end_checkpoint_id == runs[1].input_checkpoint_id + + slices = await reader.aget_run_slices(THREAD, "tsk1", runs) + assert [m.id for m in slices[0].messages] == ["sub-h-1", "sub-ai-1"] + assert [m.id for m in slices[1].messages] == ["sub-ai-2"] + assert slices[0].newly_offloaded_args == 0 + assert slices[1].newly_offloaded_args == 1 + # The namespace tip alone has lost the first run entirely. + tip = await reader.aget_task_history(THREAD, "tsk1") + assert "sub-ai-1" not in [m.id for m in tip.messages] diff --git a/tests/unit/server/services/history/test_reader_ui.py b/tests/unit/server/services/history/test_reader_ui.py new file mode 100644 index 000000000..9203f2808 --- /dev/null +++ b/tests/unit/server/services/history/test_reader_ui.py @@ -0,0 +1,315 @@ +"""The reader's ui channel: records appended after a turn, a task's terminal +snapshot, and records a node pushes mid-turn.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import START, StateGraph +from langgraph.graph.ui import ui_message_reducer +from langgraph.types import Command, interrupt + +from ptc_agent.agent.state import DeltaAgentState +from src.server.services.history.reader import CheckpointHistoryReader + +from tests.unit.server.services.history.reader_graphs import ( + THREAD, + _cfg, + _echo_graph, + _history, + _run_turns, +) + +pytestmark = pytest.mark.asyncio + + +async def test_append_ui_record_no_new_boundary(): + saver = InMemorySaver() + graph = _echo_graph(saver) + await _run_turns(graph, 2) + + reader = CheckpointHistoryReader(saver) + await reader.append_ui_record(THREAD, "image_capture", {"path_to_url": {"a.png": "https://x/a"}}) + await reader.append_ui_record(THREAD, "image_capture", {"path_to_url": {"b.png": "https://x/b"}}) + + history = await _history(reader) + assert len(history.turns) == 2 # updates are not input boundaries + assert [r["name"] for r in history.ui] == ["image_capture", "image_capture"] + assert history.ui[0]["props"] == {"path_to_url": {"a.png": "https://x/a"}} + assert history.ui[1]["props"] == {"path_to_url": {"b.png": "https://x/b"}} + assert all(r["id"] for r in history.ui) + + # A ui append after the last turn must not shift the turn's message slice. + assert [m.content for m in history.turns[1].messages] == ["q1", "echo: q1"] + + +async def test_append_ui_record_skipped_on_interrupted_tip(): + # An interrupted turn's tip carries a pending __interrupt__ write. Appending + # a ui record there via aupdate_state attributes the write to the reader + # graph's node and clears the pending interrupt, silently breaking HITL + # resume (the image-capture Hook B fallback fires at interrupt time, so this + # is reachable when a subagent emits a sandbox image on an interrupted turn). + # The append must skip while the tip is interrupted. + saver = InMemorySaver() + + def agent(state): + humans = [m for m in state["messages"] if isinstance(m, HumanMessage)] + if len(humans) == 1: + answer = interrupt({"action_requests": [{"description": "go?"}]}) + return {"messages": [AIMessage(content=f"resumed: {answer}", id="ai-r")]} + return {"messages": [AIMessage(content="ok", id=f"ai-{len(state['messages'])}")]} + + graph = ( + StateGraph(DeltaAgentState) + .add_node("agent", agent) + .add_edge(START, "agent") + .compile(checkpointer=saver) + ) + await graph.ainvoke({"messages": [HumanMessage(content="q0", id="h-0")]}, _cfg()) + + reader = CheckpointHistoryReader(saver) + assert (await graph.aget_state(_cfg())).next == ("agent",) + + await reader.append_ui_record( + THREAD, "image_capture", {"path_to_url": {"a.png": "https://x/a"}} + ) + + # Interrupt survives the append: the pending task is intact and resume runs. + assert (await graph.aget_state(_cfg())).next == ("agent",) + result = await graph.ainvoke(Command(resume="yes"), _cfg()) + assert any( + isinstance(m, AIMessage) and m.content == "resumed: yes" + for m in result["messages"] + ) + # The skipped record did not land. + history = await _history(reader) + assert history.ui == [] + + +async def test_append_ui_record_skipped_when_interrupt_check_fails(monkeypatch): + reader = CheckpointHistoryReader(InMemorySaver()) + interrupt_check = AsyncMock(side_effect=RuntimeError("checkpoint unavailable")) + update = AsyncMock() + monkeypatch.setattr(reader._checkpointer, "aget_tuple", interrupt_check) + monkeypatch.setattr(reader._updater, "aupdate_state", update) + + await reader.append_ui_record( + THREAD, "image_capture", {"path_to_url": {"a.png": "https://x/a"}} + ) + + interrupt_check.assert_awaited_once() + update.assert_not_awaited() + + +async def test_append_ui_record_advances_recorded_branch_tip(monkeypatch): + # The image-capture hook appends after turn end, creating a checkpoint + # beyond the recorded branch tip where the replay walk cannot see it. The + # append must CAS the recorded tip from the checkpoint it built on onto + # the new one. (The workflow terminal snapshot does NOT come through here + # — it writes into task:{id} via persist_task_ui_record, below.) + saver = InMemorySaver() + graph = _echo_graph(saver) + await _run_turns(graph, 1) + reader = CheckpointHistoryReader(saver) + + advance = AsyncMock(return_value=True) + monkeypatch.setattr( + "src.server.database.conversation.threads_write." + "advance_thread_checkpoint_id", + advance, + ) + root_cfg = {"configurable": {"thread_id": THREAD}} + tip_before = await saver.aget_tuple(root_cfg) + + await reader.append_ui_record( + THREAD, "image_capture", {"path_to_url": {"chart.png": "https://x/chart.png"}} + ) + + tip_after = await saver.aget_tuple(root_cfg) + assert advance.await_args.args[0] == THREAD + kwargs = advance.await_args.kwargs + assert ( + kwargs["from_checkpoint_id"] + == tip_before.config["configurable"]["checkpoint_id"] + ) + assert ( + kwargs["to_checkpoint_id"] + == tip_after.config["configurable"]["checkpoint_id"] + ) + assert kwargs["to_checkpoint_id"] != kwargs["from_checkpoint_id"] + + +async def test_append_ui_record_skipped_on_a_thread_with_no_checkpoint(monkeypatch): + # With nothing to anchor to, the write would land on whatever checkpoint + # a starting turn writes first, and a CAS from no pointer would publish it. + saver = InMemorySaver() + advance = AsyncMock(return_value=True) + monkeypatch.setattr( + "src.server.database.conversation.threads_write." + "advance_thread_checkpoint_id", + advance, + ) + + await CheckpointHistoryReader(saver).append_ui_record(THREAD, "image_capture", {}) + + assert await saver.aget_tuple({"configurable": {"thread_id": THREAD}}) is None + advance.assert_not_awaited() + + +async def test_append_ui_record_anchors_to_the_tip_it_read(monkeypatch): + # A turn starting on another worker between the tip read and the append + # must not re-parent the record onto that turn's uncommitted checkpoint: + # the CAS guard is the tip that was read, so an unanchored write would + # still pass it and publish partial state as the thread's commit pointer. + saver = InMemorySaver() + graph = _echo_graph(saver) + await _run_turns(graph, 1) + reader = CheckpointHistoryReader(saver) + root_cfg = {"configurable": {"thread_id": THREAD}} + stale_tip = await saver.aget_tuple(root_cfg) + stale_id = stale_tip.config["configurable"]["checkpoint_id"] + + # The concurrent turn lands. It has not finalized, so the recorded pointer + # still names the turn-1 tip — which is exactly why the CAS would pass. + await _run_turns(graph, 1, start=1) + moved = await saver.aget_tuple(root_cfg) + moved_id = moved.config["configurable"]["checkpoint_id"] + assert moved_id != stale_id + + class _StaleTip: + """Models the read/write window: the reader holds the older tip.""" + + async def aget_tuple(self, config): + return stale_tip + + monkeypatch.setattr(reader, "_checkpointer", _StaleTip()) + advance = AsyncMock(return_value=True) + monkeypatch.setattr( + "src.server.database.conversation.threads_write." + "advance_thread_checkpoint_id", + advance, + ) + + await reader.append_ui_record(THREAD, "image_capture", {"path_to_url": {}}) + + kwargs = advance.await_args.kwargs + assert kwargs["from_checkpoint_id"] == stale_id + written = await saver.aget_tuple( + { + "configurable": { + "thread_id": THREAD, + "checkpoint_id": kwargs["to_checkpoint_id"], + } + } + ) + # The guard and the write agree on one parent — the whole point. + assert written.parent_config["configurable"]["checkpoint_id"] == stale_id + + +async def test_task_ui_snapshot_readable_through_reader(): + # Cross-layer contract: the workflow driver persists its terminal ui + # snapshot middleware-side (through the run's own checkpointer, so the + # writer-guard fence applies), and the server reader materializes it. The + # snapshot may land while the launching turn is still executing — a + # root-ns append there forks a dead branch — so the writer must (a) + # surface via aget_task_history, (b) upsert by record id, and (c) leave + # the root chain untouched. + from ptc_agent.agent.middleware.background_subagent.workflow.ui_snapshot import ( + persist_task_ui_record, + ) + + saver = InMemorySaver() + graph = _echo_graph(saver) + await _run_turns(graph, 1) + reader = CheckpointHistoryReader(saver) + root_cfg = {"configurable": {"thread_id": THREAD}} + tip_before = await saver.aget_tuple(root_cfg) + + await persist_task_ui_record( + saver, + THREAD, + "wf1", + "workflow_run", + {"task_id": "wf1", "frames": [{"phase": "run_started"}]}, + record_id="workflow-run-r1", + ) + await persist_task_ui_record( + saver, + THREAD, + "wf1", + "workflow_run", + { + "task_id": "wf1", + "frames": [{"phase": "run_started"}, {"phase": "run_completed"}], + }, + record_id="workflow-run-r1", + ) + + task_history = await reader.aget_task_history(THREAD, "wf1") + records = [ + r for r in task_history.new_ui_records if r["id"] == "workflow-run-r1" + ] + assert len(records) == 1 # same-id rewrite upserts, never duplicates + assert [f["phase"] for f in records[0]["props"]["frames"]] == [ + "run_started", + "run_completed", + ] + + # Root chain untouched: same tip, no root ui records, no new boundary. + tip_after = await saver.aget_tuple(root_cfg) + assert ( + tip_after.config["configurable"]["checkpoint_id"] + == tip_before.config["configurable"]["checkpoint_id"] + ) + history = await _history(reader) + assert history.ui == [] + assert len(history.turns) == 1 + + +async def test_new_ui_records_attributed_to_their_turn(): + # push_ui_message inside a node (the resilience middleware's fallback + # path) checkpoints the record on the ui channel; the reader id-diffs the + # channel per turn so replay can project the notice into the right turn. + from langgraph.graph.ui import push_ui_message + + saver = InMemorySaver() + + def agent(state): + humans = [m for m in state["messages"] if isinstance(m, HumanMessage)] + if len(humans) == 2: # second turn only + push_ui_message( + name="model_fallback", + props={"from_model": "primary", "to_model": "backup"}, + id="ui-fb-1", + ) + return { + "messages": [AIMessage(content="ok", id=f"ai-{len(state['messages'])}")] + } + + graph = ( + StateGraph(DeltaAgentState) + .add_node("agent", agent) + .add_edge(START, "agent") + .compile(checkpointer=saver) + ) + await _run_turns(graph, 2) + + reader = CheckpointHistoryReader(saver) + history = await _history(reader) + assert [len(t.new_ui_records) for t in history.turns] == [0, 1] + record = history.turns[1].new_ui_records[0] + assert record["name"] == "model_fallback" + assert record["props"] == {"from_model": "primary", "to_model": "backup"} + + +async def test_ui_reducer_upsert_and_unknown_remove_guard(): + a1 = {"type": "ui", "id": "u1", "name": "n", "props": {"v": 1}, "metadata": {}} + a2 = {"type": "ui", "id": "u1", "name": "n", "props": {"v": 2}, "metadata": {}} + merged = ui_message_reducer([a1], [a2]) + assert merged == [a2] # same id upserts, no duplicate + + with pytest.raises(ValueError): + ui_message_reducer([a1], [{"type": "remove-ui", "id": "missing"}]) diff --git a/tests/unit/server/services/history/test_reasoning_durations.py b/tests/unit/server/services/history/test_reasoning_durations.py index 8b45611c5..11ef41f8a 100644 --- a/tests/unit/server/services/history/test_reasoning_durations.py +++ b/tests/unit/server/services/history/test_reasoning_durations.py @@ -6,7 +6,8 @@ per block, while the projector joins those blocks into one reasoning row. """ -from src.server.services.history.replay.stored_merge import _merge_stored_payloads +from src.server.services.history.replay import legacy +from src.server.services.history.replay.stored_merge import derive_merge def _chunk(message_id: str, content: str, content_type: str, **extra): @@ -28,7 +29,7 @@ def _signal(message_id: str, content: str, **extra): def _carry(turn_items, stored): # The merge writes onto the projected items in place. - _merge_stored_payloads(turn_items, stored) + legacy.apply(turn_items, derive_merge(turn_items, stored)) return turn_items diff --git a/tests/unit/server/services/history/test_replay.py b/tests/unit/server/services/history/test_replay.py index a939db7f7..3dbe3ed55 100644 --- a/tests/unit/server/services/history/test_replay.py +++ b/tests/unit/server/services/history/test_replay.py @@ -2,24 +2,29 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from ptc_agent.agent.middleware.large_result_eviction import TOO_LARGE_TOOL_MSG -from src.server.services.history.reader import TaskHistory, ThreadHistory +from src.server.services.history.reader import TaskHistory from src.server.services.history.replay import ( CheckpointReplayUnavailable, build_checkpoint_replay_items, + build_replay_page, build_sse_replay_items, ) from tests.unit.server.services.history.replay_builders import ( + thread_rows, THREAD, + ThreadHistory, _cache_probe, _mock_reader, _query, + _reset_slice_store, _response, + _slice_store, _turn, ) @@ -40,9 +45,7 @@ async def test_legacy_backfilled_steering_falls_back(monkeypatch): _mock_reader(monkeypatch, history) with pytest.raises(CheckpointReplayUnavailable, match="no checkpoint boundary"): await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, qtype="steering")], - {0: _response(0), 1: _response(1)}, + thread_rows([_query(0), _query(1, qtype="steering")], {0: _response(0), 1: _response(1)}), ) @@ -50,35 +53,29 @@ async def test_missing_committed_tip_is_replay_unavailable(monkeypatch): from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound reader = _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD)) - reader.aget_thread_history = AsyncMock( + reader.thread_history = AsyncMock( side_effect=CheckpointBranchTipNotFound(THREAD, "missing-tip") ) with pytest.raises(CheckpointReplayUnavailable, match="missing-tip"): await build_checkpoint_replay_items( - THREAD, - [_query(0)], - {0: _response(0)}, - branch_tip_checkpoint_id="missing-tip", + thread_rows([_query(0)], {0: _response(0)}), + "missing-tip", ) -async def test_missing_committed_tip_is_unavailable_on_cache_path(monkeypatch): - from src.server.services.history import projection_cache +async def test_missing_committed_tip_in_the_anchor_walk_is_unavailable(monkeypatch): from src.server.utils.checkpoint_helpers import CheckpointBranchTipNotFound reader = _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD)) reader.aget_turn_anchors = AsyncMock( side_effect=CheckpointBranchTipNotFound(THREAD, "missing-tip") ) - monkeypatch.setattr(projection_cache, "cache_active", lambda: True) with pytest.raises(CheckpointReplayUnavailable, match="missing-tip"): await build_checkpoint_replay_items( - THREAD, - [_query(0)], - {0: _response(0)}, - branch_tip_checkpoint_id="missing-tip", + thread_rows([_query(0)], {0: _response(0)}), + "missing-tip", ) @@ -104,7 +101,9 @@ async def test_steered_turn_projects_delivered_event(monkeypatch): AIMessage(content="done, with bonds", id="ai-2"), ] _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)])) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) events = [i["event"] for i in items] assert events == ["user_message", "message_chunk", "steering_delivered", "message_chunk"] steer = items[2]["data"] @@ -145,7 +144,7 @@ async def test_stored_events_preferred_over_projected_signals(monkeypatch): ] _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)])) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) steering = [i for i in items if i["event"] == "steering_delivered"] assert len(steering) == 1 @@ -165,9 +164,11 @@ async def test_projected_token_usage_gets_threshold(monkeypatch): ] _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)])) monkeypatch.setattr( - "src.server.services.history.replay.resolve_token_threshold", lambda: 99000 + "src.server.services.history.replay.turn.resolve_token_threshold", lambda: 99000 + ) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), ) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) cw = next(i for i in items if i["event"] == "context_window") assert cw["data"]["action"] == "token_usage" assert cw["data"]["input_tokens"] == 50 @@ -192,7 +193,7 @@ async def test_summarization_event_reemitted_at_turn_head(monkeypatch): turns = [_turn(0, [AIMessage(content="a0", id="ai-0")]), turn] _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=turns)) items = await build_checkpoint_replay_items( - THREAD, [_query(0), _query(1)], {0: _response(0), 1: _response(1)} + thread_rows([_query(0), _query(1)], {0: _response(0), 1: _response(1)}), ) turn1 = [i for i in items if i["data"].get("turn_index") == 1] assert [i["event"] for i in turn1] == [ @@ -226,7 +227,7 @@ async def test_completed_turn_without_boundary_unavailable(monkeypatch): _mock_reader(monkeypatch, history) with pytest.raises(CheckpointReplayUnavailable, match="no checkpoint boundary"): await build_checkpoint_replay_items( - THREAD, [_query(0), _query(1)], {1: _response(1)} + thread_rows([_query(0), _query(1)], {1: _response(1)}), ) @@ -240,7 +241,7 @@ async def test_inflight_active_turn_replays_as_stub(monkeypatch): responses = {i: _response(i) for i in range(3)} responses[3] = _response(3, status="streaming") items = await build_checkpoint_replay_items( - THREAD, [_query(i) for i in range(4)], responses + thread_rows([_query(i) for i in range(4)], responses), ) users = [i["data"]["turn_index"] for i in items if i["event"] == "user_message"] assert users == [0, 1, 2, 3] @@ -258,7 +259,7 @@ async def test_user_message_run_id_stamped_only_when_terminal(monkeypatch): responses = {i: _response(i) for i in range(3)} responses[3] = _response(3, status="streaming") items = await build_checkpoint_replay_items( - THREAD, [_query(i) for i in range(4)], responses + thread_rows([_query(i) for i in range(4)], responses), ) stamps = { i["data"]["turn_index"]: i["data"].get("run_id") @@ -272,20 +273,13 @@ async def test_inflight_windowed_keeps_absolute_pairing(monkeypatch): # Regression for the live-proven mislabel: ?limit=N during a streaming # turn must not staple the previous turn's answer under the active turn's # index. Window covers boundaries 1-2 of a thread whose turn 3 is live. - turns = [_turn(i, [AIMessage(content=f"a{i}", id=f"ai-{i}")]) for i in (1, 2)] - reader = MagicMock() - reader.aget_recent_history = AsyncMock( - return_value=ThreadHistory(thread_id=THREAD, turns=turns) - ) - reader.aget_task_history = AsyncMock(return_value=TaskHistory()) - monkeypatch.setattr( - "src.server.services.history.replay.CheckpointHistoryReader.get_instance", - lambda: reader, - ) + turns = [_turn(i, [AIMessage(content=f"a{i}", id=f"ai-{i}")]) for i in range(3)] + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=turns)) responses = {i: _response(i) for i in range(3)} responses[3] = _response(3, status="streaming") items = await build_checkpoint_replay_items( - THREAD, [_query(i) for i in range(4)], responses, last_n_turns=2 + thread_rows([_query(i) for i in range(4)], responses), + last_n_turns=2, ) chunks = [i["data"] for i in items if i["event"] == "message_chunk"] assert [(c["turn_index"], c["content"]) for c in chunks] == [ @@ -311,7 +305,7 @@ async def test_stamped_turn_index_overrides_ordinal(monkeypatch): 2: _response(2), } items = await build_checkpoint_replay_items( - THREAD, [_query(i) for i in range(3)], responses + thread_rows([_query(i) for i in range(3)], responses), ) chunks = [i["data"] for i in items if i["event"] == "message_chunk"] assert [(c["turn_index"], c["content"]) for c in chunks] == [ @@ -322,19 +316,44 @@ async def test_stamped_turn_index_overrides_ordinal(monkeypatch): assert [i["event"] for i in turn1_items] == ["user_message"] +async def test_a_resume_after_a_dead_turn_pairs_to_the_next_turn(monkeypatch): + # Turn 1 never checkpointed, and turn 3 resumes turn 2's interrupt: the + # resume carries no stamp of its own and names the turn after turn 2, + # where its position (2) would collide with turn 2. + turns = [ + _turn(0, [AIMessage(content="a0", id="ai-0")], turn_index=0), + _turn(1, [AIMessage(content="a2", id="ai-2")], turn_index=2), + _turn(2, [AIMessage(content="a3", id="ai-3")]), + ] + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=turns)) + responses = {i: _response(i) for i in range(4)} + responses[1] = _response(1, status="error") + items = await build_checkpoint_replay_items( + thread_rows([_query(i) for i in range(4)], responses), + ) + chunks = [i["data"] for i in items if i["event"] == "message_chunk"] + assert [(c["turn_index"], c["content"]) for c in chunks] == [ + (0, "a0"), + (2, "a2"), + (3, "a3"), + ] + users = [i["data"]["turn_index"] for i in items if i["event"] == "user_message"] + assert users == [0, 1, 2, 3] + + async def test_stamped_turn_index_unknown_unavailable(monkeypatch): history = ThreadHistory( thread_id=THREAD, turns=[_turn(0, [], turn_index=5)] ) _mock_reader(monkeypatch, history) with pytest.raises(CheckpointReplayUnavailable, match="no persisted turn"): - await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + await build_checkpoint_replay_items(thread_rows([_query(0)], {0: _response(0)})) async def test_no_turns_unavailable(monkeypatch): _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD)) with pytest.raises(CheckpointReplayUnavailable, match="no checkpoint turns"): - await build_checkpoint_replay_items(THREAD, [_query(0)], {}) + await build_checkpoint_replay_items(thread_rows([_query(0)], {})) async def test_basic_turn_projection_and_enrichment(monkeypatch): @@ -343,7 +362,7 @@ async def test_basic_turn_projection_and_enrichment(monkeypatch): ) _mock_reader(monkeypatch, history) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) assert [i["event"] for i in items] == ["user_message", "message_chunk"] assert items[0]["data"]["content"] == "hello" @@ -387,7 +406,7 @@ async def test_stored_widget_replaces_projected(monkeypatch): _mock_reader(monkeypatch, history) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, [stored_widget, stored_cw])} + thread_rows([_query(0)], {0: _response(0, [stored_widget, stored_cw])}), ) widgets = [ i for i in items @@ -429,7 +448,7 @@ async def test_passthrough_events_keep_mid_turn_position(monkeypatch): _mock_reader(monkeypatch, history) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) events = [i["event"] for i in items] # context_window lands between the tool result and the final text — its @@ -469,7 +488,9 @@ async def test_evicted_tool_result_restored_from_stored(monkeypatch): history = ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)]) _mock_reader(monkeypatch, history) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0, stored)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0, stored)}), + ) result = next(i for i in items if i["event"] == "tool_call_result") assert result["data"]["content"] == full # restored, not the pointer @@ -493,7 +514,9 @@ async def test_evicted_pointer_in_both_streams_is_noop(monkeypatch): history = ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)]) _mock_reader(monkeypatch, history) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0, stored)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0, stored)}), + ) result = next(i for i in items if i["event"] == "tool_call_result") assert result["data"]["content"] == _EVICTION_POINTER # unchanged @@ -527,7 +550,7 @@ async def test_stored_interrupt_and_error_pass_through_in_position(monkeypatch): _mock_reader(monkeypatch, history) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) events = [i["event"] for i in items] # The resolved interrupt keeps its original slot: after the tool call, @@ -561,7 +584,9 @@ async def test_image_map_applied_from_ui_records(monkeypatch): ], ) _mock_reader(monkeypatch, history) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) chunk = next(i for i in items if i["event"] == "message_chunk") assert chunk["data"]["content"] == "![chart](https://cdn/x/chart.png)" @@ -598,16 +623,21 @@ async def test_image_maps_are_scoped_to_their_turn(monkeypatch): _mock_reader(monkeypatch, history) items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="next")], - {0: _response(0), 1: _response(1)}, + thread_rows([_query(0), _query(1, content="next")], {0: _response(0), 1: _response(1)}), ) chunks = [i["data"]["content"] for i in items if i["event"] == "message_chunk"] assert chunks == ["![old](https://cdn/old.png)", "![new](https://cdn/new.png)"] -async def test_unresolved_images_fall_back_to_stored(monkeypatch): +# Two images whose captured URLs share a basename: no basename match tells +# them apart, so the turn's main lane replays from its stored copy (the +# wholesale branch). +_AMBIGUOUS_IMAGES = "![a](work/a/chart.png) ![b](work/b/chart.png)" +_AMBIGUOUS_STORED = "![a](https://cdn/1/chart.png) ![b](https://cdn/2/chart.png)" + + +async def test_unresolved_images_resolve_to_their_stored_urls(monkeypatch): history = ThreadHistory( thread_id=THREAD, turns=[_turn(0, [AIMessage(content="![chart](work/chart.png)", id="ai-1")])], @@ -617,17 +647,41 @@ async def test_unresolved_images_fall_back_to_stored(monkeypatch): { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": "![chart](https://cdn/x/chart.png)", "content_type": "text", "role": "assistant", }, } ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) chunk = next(i for i in items if i["event"] == "message_chunk") - assert chunk["data"]["content"] == "![chart](https://cdn/rewritten.png)" + assert chunk["data"]["content"] == "![chart](https://cdn/x/chart.png)" + assert chunk["data"]["id"] == "ai-1" # the projected message + + +async def test_ambiguous_images_replay_the_stored_main_lane(monkeypatch): + history = ThreadHistory( + thread_id=THREAD, + turns=[_turn(0, [AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1")])], + ) + _mock_reader(monkeypatch, history) + stored = [ + { + "event": "message_chunk", + "data": { + "content": _AMBIGUOUS_STORED, + "content_type": "text", + "role": "assistant", + }, + } + ] + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0, stored)}), + ) + chunk = next(i for i in items if i["event"] == "message_chunk") + assert chunk["data"]["content"] == _AMBIGUOUS_STORED # The wholesale replay copies each stored event's nested data; enrichment # must not stamp turn/response context back into the pristine source row. assert "turn_index" not in stored[0]["data"] @@ -650,7 +704,7 @@ async def test_wholesale_fallback_preserves_task_lane_and_watermark(monkeypatch) "task_run_id": "run-1", } turn_msgs = [ - AIMessage(content="![chart](work/chart.png)", id="ai-1"), + AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1"), AIMessage( content="", id="ai-2", @@ -669,9 +723,9 @@ async def test_wholesale_fallback_preserves_task_lane_and_watermark(monkeypatch) ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)]), task_messages=[AIMessage(content="task answer", id="sub-ai-1")], ) - reader.aget_task_run_stamps = AsyncMock(return_value=["run-1"]) + reader.task_run_stamps = AsyncMock(return_value=["run-1"]) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -686,7 +740,7 @@ async def test_wholesale_fallback_preserves_task_lane_and_watermark(monkeypatch) { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -707,14 +761,14 @@ async def test_wholesale_fallback_preserves_task_lane_and_watermark(monkeypatch) ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) main_chunks = [ i["data"]["content"] for i in items if i["event"] == "message_chunk" and i["data"].get("agent") != "task:tsk1" ] - assert "![chart](https://cdn/rewritten.png)" in main_chunks + assert _AMBIGUOUS_STORED in main_chunks # The projected task transcript survived the substitution. assert [ i["data"]["content"] @@ -746,7 +800,7 @@ async def test_wholesale_fallback_dedups_legacy_interleaved_archive(monkeypatch) "prompt": "p", } turn_msgs = [ - AIMessage(content="![chart](work/chart.png)", id="ai-1"), + AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1"), AIMessage( content="", id="ai-2", @@ -769,7 +823,7 @@ async def test_wholesale_fallback_dedups_legacy_interleaved_archive(monkeypatch) { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -785,7 +839,7 @@ async def test_wholesale_fallback_dedups_legacy_interleaved_archive(monkeypatch) ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) assert [ i["data"]["content"] @@ -807,7 +861,7 @@ async def test_wholesale_fallback_tool_only_archive_renders_once(monkeypatch): "prompt": "p", } turn_msgs = [ - AIMessage(content="![chart](work/chart.png)", id="ai-1"), + AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1"), AIMessage( content="", id="ai-2", @@ -840,7 +894,7 @@ async def test_wholesale_fallback_tool_only_archive_renders_once(monkeypatch): { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -873,7 +927,7 @@ async def test_wholesale_fallback_tool_only_archive_renders_once(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) task_rows = [ i for i in items if (i["data"] or {}).get("agent") == "task:tsk1" @@ -916,7 +970,7 @@ async def test_wholesale_fallback_stop_snapshot_does_not_evict_projection( "task_run_id": "run-1", } turn_msgs = [ - AIMessage(content="![chart](work/chart.png)", id="ai-1"), + AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1"), AIMessage( content="", id="ai-2", @@ -938,9 +992,9 @@ async def test_wholesale_fallback_stop_snapshot_does_not_evict_projection( AIMessage(content="partial answer plus the checkpointed tail", id="sub-ai-1"), ], ) - reader.aget_task_run_stamps = AsyncMock(return_value=["run-1"]) + reader.task_run_stamps = AsyncMock(return_value=["run-1"]) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -955,7 +1009,7 @@ async def test_wholesale_fallback_stop_snapshot_does_not_evict_projection( { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -1000,7 +1054,7 @@ async def test_wholesale_fallback_stop_snapshot_does_not_evict_projection( ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) # The full checkpointed transcript renders, once; the partial snapshot # rows (opener, truncated chunk, synthetic close) are gone. @@ -1029,8 +1083,8 @@ async def test_wholesale_fallback_stop_snapshot_does_not_evict_projection( def _wholesale_task_turn(): - """A turn launching tsk1 whose main lane has an unresolved sandbox image - (activates the wholesale branch).""" + """A turn launching tsk1 whose main lane shows images no basename match + tells apart (activates the wholesale branch).""" task_artifact = { "task_id": "tsk1", "action": "init", @@ -1038,7 +1092,7 @@ def _wholesale_task_turn(): "prompt": "p", } return [ - AIMessage(content="![chart](work/chart.png)", id="ai-1"), + AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1"), AIMessage( content="", id="ai-2", @@ -1081,7 +1135,7 @@ async def test_wholesale_fallback_restores_evicted_task_results(monkeypatch): { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -1099,7 +1153,7 @@ async def test_wholesale_fallback_restores_evicted_task_results(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) results = [ i @@ -1139,7 +1193,7 @@ async def test_wholesale_fallback_task_signals_render_once_in_position( { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -1176,7 +1230,7 @@ async def test_wholesale_fallback_task_signals_render_once_in_position( ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) task_rows = [ i for i in items if (i["data"] or {}).get("agent") == "task:tsk1" @@ -1202,9 +1256,9 @@ def _ledgered_run(monkeypatch, reader, status): from datetime import datetime, timezone started = datetime(2026, 1, 3, tzinfo=timezone.utc) - reader.aget_task_run_stamps = AsyncMock(return_value=["run-1"]) + reader.task_run_stamps = AsyncMock(return_value=["run-1"]) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -1257,11 +1311,11 @@ def _ledgered_task_turn(): ] -async def test_wholesale_fallback_errored_run_partial_text_survives(monkeypatch): +async def test_wholesale_fallback_errored_run_partial_text_is_dropped(monkeypatch): """A model call that raises mid-stream leaves partial text in the capture - that the checkpoint never committed (only the opener checkpointed). For - the errored run's lane the merge resurrects those trailing stored chunks - after the projected opener — text the user saw live must survive reload.""" + that the checkpoint never committed (only the opener checkpointed). The + model never sees that text again, so replay leaves it out and keeps the + run's error.""" reader = _mock_reader( monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, _ledgered_task_turn())]), @@ -1272,7 +1326,7 @@ async def test_wholesale_fallback_errored_run_partial_text_survives(monkeypatch) { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -1281,24 +1335,21 @@ async def test_wholesale_fallback_errored_run_partial_text_survives(monkeypatch) ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) task_rows = [ i for i in items if (i["data"] or {}).get("agent") == "task:tsk1" ] assert [(i["event"], i["data"].get("content")) for i in task_rows] == [ ("user_message", "p"), - ("message_chunk", "The answer is"), - ("message_chunk", " incomplete"), ("error", None), ] # The single opener is the projected one. assert task_rows[0]["data"].get("id") == "sub-h-1" -async def test_normal_path_errored_run_partial_text_survives(monkeypatch): - """Same resurrection on the normal (no unresolved image) path — the fix - lives in the shared merge, not the wholesale branch.""" +async def test_normal_path_errored_run_partial_text_is_dropped(monkeypatch): + """Same on the normal (no unresolved image) path.""" msgs = _ledgered_task_turn() msgs[0] = AIMessage(content="plain main text", id="ai-1") reader = _mock_reader( @@ -1310,21 +1361,21 @@ async def test_normal_path_errored_run_partial_text_survives(monkeypatch): stored = list(_ERRORED_RUN_STORED_TASK_ROWS) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) chunks = [ i["data"]["content"] for i in items if i["event"] == "message_chunk" and i["data"].get("agent") == "task:tsk1" ] - assert chunks == ["The answer is", " incomplete"] + assert chunks == [] async def test_completed_run_phantom_partials_stay_dropped(monkeypatch): """A completed run can leave phantom partials in the capture (a model attempt that failed mid-stream before an in-run retry re-streamed the - full text). Resurrection is off for completed runs — replaying the - phantom beside the checkpointed message would double-render.""" + full text). Replaying the phantom beside the checkpointed message would + double-render.""" msgs = _ledgered_task_turn() msgs[0] = AIMessage(content="plain main text", id="ai-1") reader = _mock_reader( @@ -1360,7 +1411,7 @@ async def test_completed_run_phantom_partials_stay_dropped(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) chunks = [ i["data"]["content"] @@ -1370,11 +1421,11 @@ async def test_completed_run_phantom_partials_stay_dropped(monkeypatch): assert chunks == ["the full final text"] -async def test_errored_run_committed_copy_is_not_resurrected(monkeypatch): +async def test_errored_run_renders_committed_output_once(monkeypatch): """A phantom partial in an errored run shifts stored lane ordinals: the - committed message's stored copy lands beyond the projected count and - looks trailing. Content matching marks it as the checkpointed message's - duplicate — only genuinely capture-only output resurrects, once.""" + committed message's stored copy lands beyond the projected count. + Content matching pairs it with the checkpointed message, so the + committed text renders once and the uncommitted tail not at all.""" msgs = _ledgered_task_turn() msgs[0] = AIMessage(content="plain main text", id="ai-1") reader = _mock_reader( @@ -1420,57 +1471,26 @@ async def test_errored_run_committed_copy_is_not_resurrected(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) chunks = [ i["data"]["content"] for i in items if i["event"] == "message_chunk" and i["data"].get("agent") == "task:tsk1" ] - # The committed copy renders once (projected); the capture-only tail - # survives; the phantom stays consumed. - assert chunks == ["the full final text", "and then it died"] - - -async def test_lossy_lane_awaiting_collector_stays_uncacheable(monkeypatch): - """The collector races the refresh-at-finalize and never invalidates the - projection cache: an errored run's turn built before the captured rows - reach Postgres must stay uncacheable (rebuild per read), else the - opener-only build freezes the loss for the cache TTL. Once the lane has - stored rows the atomic collector write has landed and the turn caches.""" - cached = _cache_probe(monkeypatch) - msgs = _ledgered_task_turn() - msgs[0] = AIMessage(content="plain main text", id="ai-1") - turn = _turn(0, msgs) - turn.tail_checkpoint_id = "tail-0" - reader = _mock_reader( - monkeypatch, - ThreadHistory(thread_id=THREAD, turns=[turn]), - task_messages=[HumanMessage(content="p", id="sub-h-1")], - ) - _ledgered_run(monkeypatch, reader, "error") - - # Archive not yet written: no task-lane stored rows -> uncacheable. - await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) - assert cached == [] - - # Collector landed: lane rows present -> cacheable again. - stored = list(_ERRORED_RUN_STORED_TASK_ROWS) - await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} - ) - assert cached == ["tail-0"] + assert chunks == ["the full final text"] -async def test_evicted_pointer_awaiting_collector_stays_uncacheable(monkeypatch): - """Same race, completed run: the checkpoint holds only the eviction - pointer and the fuller stored result exists only in the collector's - pending archive — caching the pointer build would freeze it.""" - cached = _cache_probe(monkeypatch) +async def test_archive_landing_reprojects_a_stored_turn(monkeypatch): + """The checkpoint holds only the eviction pointer and the fuller result + reaches the response row later, with the collector's archive. The turn's + lines store as soon as its run settles; the archive rewrites the row, so + the key moves and the next read projects the fuller copy.""" + stored_tails = _cache_probe(monkeypatch) msgs = _ledgered_task_turn() msgs[0] = AIMessage(content="plain main text", id="ai-1") turn = _turn(0, msgs) - turn.tail_checkpoint_id = "tail-0" + turn.anchor.tail_checkpoint_id = "tail-0" reader = _mock_reader( monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn]), @@ -1491,113 +1511,46 @@ async def test_evicted_pointer_awaiting_collector_stays_uncacheable(monkeypatch) ) _ledgered_run(monkeypatch, reader, "completed") - await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) - assert cached == [] - - -async def test_task_custom_artifact_does_not_clear_archive_gate(monkeypatch): - """Task custom artifacts (todo/ui) are written to the root archive LIVE, - before any collector runs — an artifact row is lane presence, not - archive evidence. Only transcript-class rows (which the atomic archive - writers alone produce) clear the awaiting-archive debt.""" - cached = _cache_probe(monkeypatch) - msgs = _ledgered_task_turn() - msgs[0] = AIMessage(content="plain main text", id="ai-1") - turn = _turn(0, msgs) - turn.tail_checkpoint_id = "tail-0" - reader = _mock_reader( - monkeypatch, - ThreadHistory(thread_id=THREAD, turns=[turn]), - task_messages=[HumanMessage(content="p", id="sub-h-1")], - ) - _ledgered_run(monkeypatch, reader, "error") - artifact_only = [ - { - "event": "artifact", - "data": { - "artifact_type": "todo_list", - "artifact_id": "todo-1", - "agent": "task:tsk1", - "payload": {"todos": []}, - }, - } - ] - - await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, list(artifact_only))} - ) - assert cached == [] + def task_results(items): + return [ + i["data"]["content"] + for i in items + if i["event"] == "tool_call_result" + and i["data"].get("agent") == "task:tsk1" + ] - # Transcript rows present -> the archive landed -> cacheable. - await build_checkpoint_replay_items( - THREAD, - [_query(0)], - {0: _response(0, artifact_only + list(_ERRORED_RUN_STORED_TASK_ROWS))}, + before = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), ) - assert cached == ["tail-0"] + assert task_results(before) == [_EVICTION_POINTER] + assert stored_tails == ["tail-0"] - -async def test_repeated_text_capture_only_message_resurrects(monkeypatch): - """Identical content is not identity: a NEW capture-only message whose - text repeats an earlier checkpointed message (agent loops emit repeated - status texts) must still resurrect. Alignment consumes each projected - message once, in order — only the first stored occurrence is the copy.""" - msgs = _ledgered_task_turn() - msgs[0] = AIMessage(content="plain main text", id="ai-1") - reader = _mock_reader( - monkeypatch, - ThreadHistory(thread_id=THREAD, turns=[_turn(0, msgs)]), - task_messages=[ - HumanMessage(content="p", id="sub-h-1"), - AIMessage(content="Still working", id="sub-ai-1"), + archived = _response( + 0, + [ + { + "event": "tool_call_result", + "data": { + "agent": "task:tsk1", + "role": "assistant", + "tool_call_id": "tc-1", + "content": "the full captured result", + "content_type": "text", + }, + } ], ) - _ledgered_run(monkeypatch, reader, "error") - stored = [ - { - "event": "message_chunk", - "data": { - "agent": "task:tsk1", - "id": "lc-copy", - "content": "Still working", - "content_type": "text", - "role": "assistant", - }, - }, - { - "event": "message_chunk", - "data": { - "agent": "task:tsk1", - "id": "lc-new", - "content": "Still working", - "content_type": "text", - "role": "assistant", - }, - }, - ] - - items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} - ) - task_chunks = [ - (i["data"].get("id"), i["data"]["content"]) - for i in items - if i["event"] == "message_chunk" and i["data"].get("agent") == "task:tsk1" - ] - # The projected message renders once; the distinct repeated message - # resurrects under its own id. - assert task_chunks == [ - ("sub-ai-1", "Still working"), - ("lc-new", "Still working"), - ] + after = await build_checkpoint_replay_items(thread_rows([_query(0)], {0: archived})) + assert task_results(after) == ["the full captured result"] + assert stored_tails == ["tail-0", "tail-0"] -async def test_rewritten_image_copy_is_not_resurrected(monkeypatch): +async def test_rewritten_image_copy_renders_once(monkeypatch): """The checkpoint copy of a message carries the durable image URL while the archive copy keeps the sandbox path split across token fragments (the row-level archive rewrite can't see fragmented markdown). Matching normalizes image targets, so the shifted committed copy is still - recognized and only genuinely lost output resurrects.""" + recognized and does not render a second time.""" msgs = _ledgered_task_turn() msgs[0] = AIMessage(content="plain main text", id="ai-1") reader = _mock_reader( @@ -1656,17 +1609,14 @@ async def test_rewritten_image_copy_is_not_resurrected(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) chunks = [ i["data"]["content"] for i in items if i["event"] == "message_chunk" and i["data"].get("agent") == "task:tsk1" ] - assert chunks == [ - "Here: ![chart](https://cdn/chart.png) done", - "and then it died", - ] + assert chunks == ["Here: ![chart](https://cdn/chart.png) done"] async def test_phantom_matching_later_text_does_not_displace_copies(monkeypatch): @@ -1719,7 +1669,7 @@ async def test_phantom_matching_later_text_does_not_displace_copies(monkeypatch) ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) task_chunks = [ (i["data"].get("id"), i["data"]["content"]) @@ -1781,7 +1731,7 @@ async def test_distinct_image_targets_do_not_collide(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) task_chunks = [ (i["data"].get("id"), i["data"]["content"]) @@ -1813,7 +1763,7 @@ async def test_wholesale_fallback_task_custom_artifact_is_not_ownership(monkeypa "task_run_id": "run-1", } turn_msgs = [ - AIMessage(content="![chart](work/chart.png)", id="ai-1"), + AIMessage(content=_AMBIGUOUS_IMAGES, id="ai-1"), AIMessage( content="", id="ai-2", @@ -1832,9 +1782,9 @@ async def test_wholesale_fallback_task_custom_artifact_is_not_ownership(monkeypa ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)]), task_messages=[AIMessage(content="task answer", id="sub-ai-1")], ) - reader.aget_task_run_stamps = AsyncMock(return_value=["run-1"]) + reader.task_run_stamps = AsyncMock(return_value=["run-1"]) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -1849,7 +1799,7 @@ async def test_wholesale_fallback_task_custom_artifact_is_not_ownership(monkeypa { "event": "message_chunk", "data": { - "content": "![chart](https://cdn/rewritten.png)", + "content": _AMBIGUOUS_STORED, "content_type": "text", "role": "assistant", }, @@ -1879,7 +1829,7 @@ async def test_wholesale_fallback_task_custom_artifact_is_not_ownership(monkeypa ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) # The projected task transcript survived despite the stored custom artifact. assert [ @@ -1956,9 +1906,7 @@ async def test_subagent_transcript_projected_once_with_image_map(monkeypatch): ], ) items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="next")], - {0: _response(0), 1: _response(1)}, + thread_rows([_query(0), _query(1, content="next")], {0: _response(0), 1: _response(1)}), ) reader.aget_task_history.assert_awaited_once_with(THREAD, "tsk1") sub_chunks = [ @@ -2030,7 +1978,7 @@ async def test_subagent_private_state_and_ui_are_checkpoint_projected(monkeypatc ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) task_items = [i for i in items if i["data"].get("agent") == "task:tsk1"] @@ -2082,7 +2030,7 @@ async def test_claimed_run_stamps_projection_watermark(monkeypatch): task_messages=[AIMessage(content="task answer", id="sub-ai-1")], ) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -2095,7 +2043,7 @@ async def test_claimed_run_stamps_projection_watermark(monkeypatch): ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) artifacts = [ i @@ -2149,9 +2097,9 @@ async def test_in_progress_run_is_excluded_from_the_watermark(monkeypatch): ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)]), task_messages=[AIMessage(content="live text", id="sub-ai-1")], ) - reader.aget_task_run_stamps = AsyncMock(return_value=["run-1"]) + reader.task_run_stamps = AsyncMock(return_value=["run-1"]) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -2164,7 +2112,7 @@ async def test_in_progress_run_is_excluded_from_the_watermark(monkeypatch): ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) artifacts = [ i @@ -2216,9 +2164,7 @@ async def test_subagent_checkpoint_read_failure_makes_replay_unavailable(monkeyp CheckpointReplayUnavailable, match="subagent checkpoint state unavailable for task:tsk1", ): - await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} - ) + await build_checkpoint_replay_items(thread_rows([_query(0)], {0: _response(0)})) async def test_terminal_interrupts_appended(monkeypatch): @@ -2230,13 +2176,40 @@ async def test_terminal_interrupts_appended(monkeypatch): ], ) _mock_reader(monkeypatch, history) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) assert items[-1]["event"] == "interrupt" assert items[-1]["data"]["interrupt_id"] == "int-1" assert items[-1]["data"]["action_requests"] == [{"description": "go?"}] assert items[-1]["data"]["finish_reason"] == "interrupt" +async def test_the_newest_page_raises_its_turns_error_before_its_interrupts( + monkeypatch, +): + """The tip's interrupts are read beside the turns: when both reads fail + the turns' error is the page's, and an interrupt read failing alone still + fails the page.""" + history = ThreadHistory( + thread_id=THREAD, turns=[_turn(0, [AIMessage(content="hi", id="ai-1")])] + ) + reader = _mock_reader(monkeypatch, history) + reader.aget_tip_interrupts = AsyncMock(side_effect=RuntimeError("interrupts")) + turn_slices = reader.aget_turn_slices + reader.aget_turn_slices = AsyncMock( + side_effect=CheckpointReplayUnavailable("turn slices") + ) + rows = thread_rows([_query(0)], {0: _response(0)}) + + with pytest.raises(CheckpointReplayUnavailable, match="turn slices"): + await build_replay_page(rows) + + reader.aget_turn_slices = turn_slices + with pytest.raises(RuntimeError, match="interrupts"): + await build_replay_page(rows) + + async def test_system_query_tagged_and_synthetic_human_dropped(monkeypatch): # Report-back (flash dispatches PTC, the PTC completion is injected as a new # flash turn): the turn's boundary input is a -wrapped HumanMessage @@ -2253,11 +2226,11 @@ async def test_system_query_tagged_and_synthetic_human_dropped(monkeypatch): HumanMessage(content=system_text, id="h-sys"), AIMessage(content="Here is your summary.", id="ai-1"), ] - history = ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs, user=system_text)]) + history = ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)]) _mock_reader(monkeypatch, history) items = await build_checkpoint_replay_items( - THREAD, [_query(0, content=system_text, qtype="system")], {0: _response(0)} + thread_rows([_query(0, content=system_text, qtype="system")], {0: _response(0)}), ) users = [i for i in items if i["event"] == "user_message"] assert len(users) == 1 @@ -2270,28 +2243,20 @@ async def test_system_query_tagged_and_synthetic_human_dropped(monkeypatch): async def test_last_n_turns_windows_to_recent(monkeypatch): # Windowed replay materializes only the last N turns and pairs them with the - # matching tail of query rows (aget_recent_history, not aget_thread_history). + # matching tail of query rows: the light anchor walk spans the thread, the + # state reads only the window. turns = [_turn(i, [AIMessage(content=f"a{i}", id=f"ai-{i}")]) for i in range(4)] - reader = MagicMock() - reader.aget_thread_history = AsyncMock( - return_value=ThreadHistory(thread_id=THREAD, turns=turns) - ) - reader.aget_recent_history = AsyncMock( - return_value=ThreadHistory(thread_id=THREAD, turns=turns[-2:]) - ) - reader.aget_task_history = AsyncMock(return_value=TaskHistory()) - monkeypatch.setattr( - "src.server.services.history.replay.CheckpointHistoryReader.get_instance", - lambda: reader, - ) + reader = _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=turns)) queries = [_query(i) for i in range(4)] items = await build_checkpoint_replay_items( - THREAD, queries, {i: _response(i) for i in range(4)}, last_n_turns=2 + thread_rows(queries, {i: _response(i) for i in range(4)}), + last_n_turns=2, ) users = [i for i in items if i["event"] == "user_message"] assert [u["data"]["turn_index"] for u in users] == [2, 3] # last two turns only - reader.aget_recent_history.assert_awaited_once() - reader.aget_thread_history.assert_not_awaited() + reader.aget_turn_slices.assert_awaited_once() + read = reader.aget_turn_slices.await_args.args[1] + assert [a.input_checkpoint_id for a in read] == ["cp-in-2", "cp-in-3"] async def test_widget_data_ref_resolved_from_storage(monkeypatch): @@ -2320,7 +2285,9 @@ async def test_widget_data_ref_resolved_from_storage(monkeypatch): "src.server.services.history.replay.widgets.get_bytes", lambda key: b'{"file.csv": "a,b"}' if key == "widgets/t/abc.json" else None, ) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) widgets = [ i for i in items if i["event"] == "artifact" and i["data"]["artifact_type"] == "html_widget" @@ -2354,7 +2321,9 @@ def _boom(key): raise AssertionError("storage must not be read for inline data") monkeypatch.setattr("src.server.services.history.replay.widgets.get_bytes", _boom) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) widgets = [ i for i in items if i["event"] == "artifact" and i["data"]["artifact_type"] == "html_widget" @@ -2385,7 +2354,9 @@ async def test_widget_data_ref_unreadable_left_in_place(monkeypatch): monkeypatch.setattr( "src.server.services.history.replay.widgets.get_bytes", lambda key: None ) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) widgets = [ i for i in items if i["event"] == "artifact" and i["data"]["artifact_type"] == "html_widget" @@ -2431,12 +2402,14 @@ async def test_answered_interrupt_projected_at_turn_end(monkeypatch): them in favor of the stored copies.""" turn_msgs = [AIMessage(content="asking", id="a-0")] turn = _turn(0, turn_msgs) - turn.ending_interrupts = [ + turn.anchor.ending_interrupts = [ {"id": "int-1", "value": {"action_requests": [{"description": "pick one"}]}} ] _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn])) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) assert [i["event"] for i in items] == ["user_message", "message_chunk", "interrupt"] card = items[2]["data"] assert card["interrupt_id"] == "int-1" @@ -2449,7 +2422,7 @@ async def test_answered_interrupt_projected_at_turn_end(monkeypatch): {"event": "interrupt", "data": {"interrupt_id": "int-1", "action_requests": [{"description": "pick one"}], "extra": "stored"}}, ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) cards = [i for i in items if i["event"] == "interrupt"] assert len(cards) == 1 @@ -2474,7 +2447,7 @@ async def test_credit_usage_synthesized_from_usage_row(monkeypatch): } items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)}, usages=[usage_row] + thread_rows([_query(0)], {0: _response(0)}, usages=[usage_row]), ) assert items[-1]["event"] == "credit_usage" data = items[-1]["data"] @@ -2491,7 +2464,7 @@ async def test_credit_usage_synthesized_from_usage_row(monkeypatch): # A turn with stored events keeps the stored credit_usage instead. stored = [{"event": "credit_usage", "data": {"total_credits": 9.99}}] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)}, usages=[usage_row] + thread_rows([_query(0)], {0: _response(0, stored)}, usages=[usage_row]), ) credits = [i for i in items if i["event"] == "credit_usage"] assert len(credits) == 1 @@ -2499,7 +2472,7 @@ async def test_credit_usage_synthesized_from_usage_row(monkeypatch): # Errored runs never emitted the event live — no synthesis. items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, status="error")}, usages=[usage_row] + thread_rows([_query(0)], {0: _response(0, status="error")}, usages=[usage_row]), ) assert not [i for i in items if i["event"] == "credit_usage"] @@ -2546,10 +2519,7 @@ async def test_credit_usage_ignores_later_subagent_usage_rows(monkeypatch): } items = await build_checkpoint_replay_items( - THREAD, - [_query(0)], - {0: _response(0)}, - usages=[main, task], + thread_rows([_query(0)], {0: _response(0)}, usages=[main, task]), ) credit = next(i for i in items if i["event"] == "credit_usage") @@ -2600,7 +2570,7 @@ async def test_provenance_synthesized_from_rows_anchored(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)}, provenance=rows + thread_rows([_query(0)], {0: _response(0)}, provenance=rows), ) events = [i["event"] for i in items] result_idx = events.index("tool_call_result") @@ -2616,7 +2586,7 @@ async def test_provenance_synthesized_from_rows_anchored(monkeypatch): # Stored events win during the transition. stored = [{"event": "provenance", "data": {"record_id": "live-1", "source_type": "web_page"}}] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)}, provenance=rows + thread_rows([_query(0)], {0: _response(0, stored)}, provenance=rows), ) prov = [i for i in items if i["event"] == "provenance"] assert [p["data"]["record_id"] for p in prov] == ["live-1"] @@ -2632,7 +2602,7 @@ async def test_terminal_error_synthesized_on_both_paths(monkeypatch): turn_msgs = [AIMessage(content="partial", id="a-0")] _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs)])) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: response}) + items = await build_checkpoint_replay_items(thread_rows([_query(0)], {0: response})) assert items[-1]["event"] == "error" data = items[-1]["data"] assert data["error"] == "boom exploded" @@ -2651,9 +2621,7 @@ async def test_terminal_error_synthesized_on_both_paths(monkeypatch): ThreadHistory(thread_id=THREAD, turns=[_turn(0, turn_msgs, turn_index=0)]), ) items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="errored")], - {0: _response(0), 1: response | {"conversation_response_id": "resp-1"}}, + thread_rows([_query(0), _query(1, content="errored")], {0: _response(0), 1: response | {"conversation_response_id": "resp-1"}}), ) assert [i["event"] for i in items] == [ "user_message", @@ -2683,9 +2651,7 @@ async def test_terminal_error_replay_sanitizes_legacy_rows(monkeypatch): ), ) - items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: response} - ) + items = await build_checkpoint_replay_items(thread_rows([_query(0)], {0: response})) error = next(i for i in items if i["event"] == "error") assert "hunter2" not in error["data"]["error"] @@ -2717,7 +2683,9 @@ async def test_model_fallback_projected_from_ui_records(monkeypatch): ] turn = _turn(0, [AIMessage(content="ok", id="ai-1")], new_ui_records=records) _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn])) - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) assert [i["event"] for i in items] == ["user_message", "model_fallback", "message_chunk"] data = items[1]["data"] assert data["agent"] == "main" @@ -2749,7 +2717,7 @@ async def test_model_fallback_stored_events_preferred(monkeypatch): turn = _turn(0, [AIMessage(content="ok", id="ai-1")], new_ui_records=records) _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn])) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, sse_events=stored)} + thread_rows([_query(0)], {0: _response(0, sse_events=stored)}), ) fallbacks = [i for i in items if i["event"] == "model_fallback"] assert len(fallbacks) == 1 @@ -2805,9 +2773,7 @@ def launch(ordinal, action, prompt): ], ) items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="r2"), _query(2, content="r3")], - {0: _response(0), 1: _response(1), 2: _response(2)}, + thread_rows([_query(0), _query(1, content="r2"), _query(2, content="r3")], {0: _response(0), 1: _response(1), 2: _response(2)}), ) openers = [ i["data"] @@ -2895,9 +2861,7 @@ def launch(ordinal, action, prompt): events = [ i["event"] for i in await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="r2")], - {0: _response(0), 1: _response(1)}, + thread_rows([_query(0), _query(1, content="r2")], {0: _response(0), 1: _response(1)}), ) if i["data"].get("agent") == "task:tsk1" ] @@ -2915,7 +2879,7 @@ async def test_live_task_final_launch_defers_to_stream(monkeypatch): # seq 1, and both together render the instruction bubble twice. Earlier # settled runs still attribute; the live remainder is not salvaged as # trailing either. - from src.server.services.history.replay import task_lane as task_lane_module + from src.server.services.history.replay import run_lane as task_lane_module async def fake_details(thread_id, task_ids): return {tid: {"status": "running", "error": None} for tid in task_ids} @@ -2963,9 +2927,7 @@ def launch(ordinal, action, prompt): ], ) items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1, content="r2")], - {0: _response(0), 1: _response(1)}, + thread_rows([_query(0), _query(1, content="r2")], {0: _response(0), 1: _response(1)}), ) task_items = [ i["data"].get("content") @@ -2976,6 +2938,7 @@ def launch(ordinal, action, prompt): # Init tail (a freshly spawned live task): the only launch is the final # one, so replay projects nothing from the namespace. + _reset_slice_store(monkeypatch) _mock_reader( monkeypatch, ThreadHistory( @@ -2987,41 +2950,33 @@ def launch(ordinal, action, prompt): ], ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) assert not [i for i in items if i["data"].get("agent") == "task:tsk1"] -async def test_trailing_salvage_keeps_its_turn_uncacheable(monkeypatch): +async def test_run_without_a_committed_launch_is_not_replayed(monkeypatch): # An orphan run beyond the last committed launch (its launching turn never - # persisted) is salvaged under that launch's stamps. The salvage-stamped - # turn must stay OUT of the projection cache: the all-cache-hit fast path - # never runs the task lane, so a cached entry would replay the turn - # without the salvage on the next refresh. Fixed regression. + # committed the launch) is not in the model's context, so replay leaves + # it out and the turns cache as usual. from src.server.services.history import replay as replay_module - from src.server.services.history.replay import task_lane as task_lane_module + from src.server.services.history.replay import run_lane as task_lane_module async def fake_details(thread_id, task_ids): return {} # settled — no live writer owns the trailing segment async def fake_live(thread_id, task_ids): - return set() - - stored: list[str | None] = [] - deleted: list[str] = [] - - async def fake_store(thread_id, tail_checkpoint_id, fingerprint, items): - stored.append(tail_checkpoint_id) - - async def fake_delete(thread_id, turn_keys): - deleted.extend(tail for tail, _fp in turn_keys) + return dict.fromkeys(task_ids, "sealed") monkeypatch.setattr(task_lane_module, "resolve_task_details", fake_details) monkeypatch.setattr( - replay_module.projection_cache, "live_task_streams", fake_live + replay_module.run_lane.task_streams, "task_stream_states", fake_live ) - monkeypatch.setattr(replay_module.projection_cache, "store_turn", fake_store) - monkeypatch.setattr(replay_module.projection_cache, "delete_turns", fake_delete) + # A pre-ledger task: the ledger reads fine and holds no run for it. + monkeypatch.setattr( + task_lane_module.sr_db, "list_runs_for_thread", AsyncMock(return_value=[]) + ) + stored = _slice_store(monkeypatch).lined def launch(ordinal, action, prompt): artifact = { @@ -3046,9 +3001,9 @@ def launch(ordinal, action, prompt): ] turn0 = _turn(0, launch(0, "init", "run one")) - turn0.tail_checkpoint_id = "tail-0" + turn0.anchor.tail_checkpoint_id = "tail-0" turn1 = _turn(1, [AIMessage(content="plain turn", id="ai-plain")]) - turn1.tail_checkpoint_id = "tail-1" + turn1.anchor.tail_checkpoint_id = "tail-1" _mock_reader( monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn0, turn1]), @@ -3061,19 +3016,16 @@ def launch(ordinal, action, prompt): ) items = await build_checkpoint_replay_items( - THREAD, [_query(0), _query(1)], {0: _response(0), 1: _response(1)} + thread_rows([_query(0), _query(1)], {0: _response(0), 1: _response(1)}), ) - salvaged = [ - i["data"] + task_chunks = [ + i["data"]["content"] for i in items - if i["event"] == "message_chunk" - and i["data"].get("agent") == "task:tsk1" - and i["data"].get("content") == "orphan done" + if i["event"] == "message_chunk" and i["data"].get("agent") == "task:tsk1" ] - assert salvaged and salvaged[0]["turn_index"] == 0 # last launch's stamps - assert stored == ["tail-1"] # the salvage-stamped turn is never stored - assert deleted == ["tail-0"] # and any pre-orphan entry is evicted + assert task_chunks == ["one done"] + assert stored == ["tail-0", "tail-1"] async def test_workflow_launch_replays_ns_ui_snapshot(monkeypatch): @@ -3131,12 +3083,12 @@ async def test_workflow_launch_replays_ns_ui_snapshot(monkeypatch): ), ) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock(return_value=[]), ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) lifecycle = [i for i in items if i["event"] == "workflow_lifecycle"] @@ -3199,7 +3151,7 @@ async def test_workflow_launch_terminal_status_follows_ledger(monkeypatch): ), ) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock( return_value=[ { @@ -3213,7 +3165,7 @@ async def test_workflow_launch_terminal_status_follows_ledger(monkeypatch): ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) terminal = [ @@ -3277,12 +3229,12 @@ async def test_workflow_turn_uncacheable_when_ledger_read_fails(monkeypatch): ), ) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock(side_effect=RuntimeError("ledger down")), ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0)} + thread_rows([_query(0)], {0: _response(0)}), ) # Frames still replay from the snapshot; the build just must not cache. @@ -3359,12 +3311,12 @@ async def test_ledger_read_failure_leaves_merge_semantics_alone(monkeypatch): ), ) monkeypatch.setattr( - "src.server.services.history.replay.task_lane.sr_db.list_runs_for_thread", + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", AsyncMock(side_effect=RuntimeError("ledger down")), ) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored)} + thread_rows([_query(0)], {0: _response(0, stored)}), ) assert not [ @@ -3377,7 +3329,7 @@ async def test_stream_probe_precedes_task_namespace_reads(monkeypatch): monotonic, so "sealed before the read" proves the read saw final state. A post-read probe would let a seal landing in between freeze pre-terminal state for the cache TTL. Pin the ordering (fixed regression).""" - from src.server.services.history.replay import task_lane as task_lane_module + from src.server.services.history.replay import run_lane as task_lane_module order: list[str] = [] task_artifact = { @@ -3415,14 +3367,14 @@ async def read_spy(*args, **kwargs): async def probe_spy(thread_id, task_ids): order.append("probe") - assert task_ids == {"tsk1"} - return set() + assert set(task_ids) == {"tsk1"} + return dict.fromkeys(task_ids, "sealed") monkeypatch.setattr( - task_lane_module.projection_cache, "live_task_streams", probe_spy + task_lane_module.task_streams, "task_stream_states", probe_spy ) - await build_checkpoint_replay_items(THREAD, [_query(0)], {0: _response(0)}) + await build_checkpoint_replay_items(thread_rows([_query(0)], {0: _response(0)})) assert order[0] == "probe" assert "read" in order diff --git a/tests/unit/server/services/history/test_replay_facts.py b/tests/unit/server/services/history/test_replay_facts.py new file mode 100644 index 000000000..2e1d345a9 --- /dev/null +++ b/tests/unit/server/services/history/test_replay_facts.py @@ -0,0 +1,503 @@ +"""Replay reads what a run's checkpoint does not keep from the run's facts. + +Reasoning durations, captured subagent images and returned steering are +measured or learned around the graph, so they reach replay through the run +rows' ``replay_facts`` rather than the stored events. A row settled before +those existed keeps replaying from its stored events, and a turn holding both +must replay the same either way. Run facts land after a run settles, so a +turn's stored lines are keyed on them. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from src.server.services.history.replay import ( + CheckpointReplayUnavailable, + build_checkpoint_replay_items, + build_replay_page, + read_replay_page, +) +from src.server.services.history.replay import cold +from src.server.services.history.replay import facts as replay_facts +from src.server.services.history.replay import legacy +from src.server.services.history.replay.stored_merge import derive_merge +from tests.unit.server.services.history.replay_builders import ( + thread_rows, + THREAD, + ThreadHistory, + _cache_probe, + _mock_reader, + _query, + _response, + _slice_store, + _turn, +) + + +def _merged(projected, stored, **kwargs): + return legacy.apply(projected, derive_merge(projected, stored, **kwargs)) + + +def _signal(message_id, content, agent="main", **extra): + return { + "event": "message_chunk", + "data": { + "thread_id": THREAD, + "agent": agent, + "id": message_id, + "role": "assistant", + "content": content, + "content_type": "reasoning_signal", + **extra, + }, + } + + +def _text(message_id, content, agent="main"): + return { + "event": "message_chunk", + "data": { + "agent": agent, + "id": message_id, + "role": "assistant", + "content": content, + "content_type": "text", + }, + } + + +def _returned(input_id, agent="task:tsk1", content="also AMD", reason="run_ended"): + return { + "event": "steering_returned", + "data": { + "agent": agent, + "content": content, + "input_id": input_id, + "reason": reason, + "thread_id": THREAD, + }, + } + + +# ------------------------------------------------------------ reasoning + + +def test_durations_sum_onto_each_main_message_close(): + """The projector joins a message's thinking blocks into one row, so its + one close stands for every block the stream timed.""" + items = [ + _signal("m1", "start"), + _signal("m1", "complete"), + _signal("m2", "start"), + _signal("m2", "complete"), + _signal("t1", "complete", agent="task:k"), + ] + replay_facts.apply_reasoning_durations( + items, {"m1": [3000, 12000, 8000], "m2": [900], "t1": [5]} + ) + assert items[1]["data"]["elapsed_ms"] == 23000 + assert items[3]["data"]["elapsed_ms"] == 900 + assert "elapsed_ms" not in items[0]["data"] + # Task lanes were never timed live; a stray key does not invent one. + assert "elapsed_ms" not in items[4]["data"] + + +def test_durations_ignore_untimed_and_malformed_spans(): + items = [_signal("m1", "complete"), _signal("m2", "complete")] + replay_facts.apply_reasoning_durations(items, {"m1": [True, "9"], "m2": "900"}) + assert "elapsed_ms" not in items[0]["data"] + assert "elapsed_ms" not in items[1]["data"] + + +def test_known_durations_are_not_carried_again_from_stored_events(): + """With the row's facts applied, the stored copy of the same closes is + not summed on top.""" + projected = [_signal("ai-1", "start"), _signal("ai-1", "complete")] + replay_facts.apply_reasoning_durations(projected, {"ai-1": [1200, 800]}) + stored = [ + _signal("lc-1", "start"), + _signal("lc-1", "complete", elapsed_ms=1200), + _signal("lc-1", "start"), + _signal("lc-1", "complete", elapsed_ms=800), + ] + merged = _merged(projected, stored, main_durations_known=True) + assert merged[1]["data"]["elapsed_ms"] == 2000 + + +def test_stored_durations_still_carry_for_a_row_without_facts(): + projected = [_signal("ai-1", "start"), _signal("ai-1", "complete")] + stored = [_signal("lc-1", "start"), _signal("lc-1", "complete", elapsed_ms=4200)] + merged = _merged(projected, stored) + assert merged[1]["data"]["elapsed_ms"] == 4200 + + +# ------------------------------------------------------------ run facts + + +def _run_items(): + return [ + _text("a1", "first", agent="task:tsk1"), + _text("a2", "second", agent="task:tsk1"), + ] + + +def _placed(items): + return [ + (i["event"], i["data"].get("input_id") or i["data"].get("content")) + for i in items + ] + + +def test_returns_at_run_end_follow_everything_the_run_wrote(): + out = replay_facts.apply_run_facts( + _run_items(), + { + "steering_returned": [ + {"content": "x", "input_id": "in-1", "reason": "run_ended"} + ] + }, + agent="task:tsk1", + thread_id=THREAD, + message_ids=["h1", "a1", "a2"], + ) + assert _placed(out) == [ + ("message_chunk", "first"), + ("message_chunk", "second"), + ("steering_returned", "in-1"), + ] + assert out[-1]["data"] == { + "agent": "task:tsk1", + "content": "x", + "input_id": "in-1", + "reason": "run_ended", + "thread_id": THREAD, + } + + +def test_returns_before_a_model_call_follow_the_message_the_run_had_reached(): + returned = [ + {"input_id": "r-head", "reason": "run_mismatch", "after": None}, + {"input_id": "r-a1", "reason": "run_mismatch", "after": "a1"}, + # The opener projects no item of its own: the return leads the run. + {"input_id": "r-h1", "reason": "run_mismatch", "after": "h1"}, + # A message this run's transcript does not hold: end of the run. + {"input_id": "r-gone", "reason": "run_mismatch", "after": "elsewhere"}, + {"input_id": "r-a1", "reason": "run_mismatch", "after": "a1"}, # repeated write + ] + out = replay_facts.apply_run_facts( + _run_items(), + {"steering_returned": returned}, + agent="task:tsk1", + thread_id=THREAD, + message_ids=["h1", "a1", "a2"], + ) + assert _placed(out) == [ + ("steering_returned", "r-head"), + ("steering_returned", "r-h1"), + ("message_chunk", "first"), + ("steering_returned", "r-a1"), + ("message_chunk", "second"), + ("steering_returned", "r-gone"), + ] + + +def test_run_images_resolve_the_runs_text(): + items = [ + _text("a1", "![c](work/sub.png) and ![d](work/other.png)", agent="task:tsk1") + ] + replay_facts.apply_run_facts( + items, + {"images": {"work/sub.png": "https://cdn/sub.png"}}, + agent="task:tsk1", + thread_id=THREAD, + message_ids=["a1"], + ) + assert items[0]["data"]["content"] == ( + "![c](https://cdn/sub.png) and ![d](work/other.png)" + ) + + +def test_a_placed_return_is_not_merged_again_from_stored_events(): + projected = [_text("a1", "first", agent="task:tsk1"), _returned("in-1")] + stored = [ + _text("lc-a1", "first", agent="task:tsk1"), + _returned("in-1"), + _returned("in-2"), + ] + merged = _merged(projected, stored) + # A run settled before facts keeps its stored returns; one with facts + # shows each of its returns once. + assert sorted( + i["data"]["input_id"] for i in merged if i["event"] == "steering_returned" + ) == ["in-1", "in-2"] + + +# ------------------------------------------------------------ through replay + + +def _task_turn(): + task_artifact = { + "task_id": "tsk1", + "action": "init", + "description": "d", + "prompt": "p", + "task_run_id": "run-1", + } + return [ + AIMessage( + content="", + id="ai-1", + tool_calls=[{"name": "Task", "args": {}, "id": "tc-t"}], + ), + ToolMessage( + content="dispatched", + tool_call_id="tc-t", + name="Task", + id="tm-1", + additional_kwargs={"task_artifact": task_artifact}, + ), + ] + + +def _ledgered(monkeypatch, reader): + from datetime import datetime, timezone + + reader.task_run_stamps = AsyncMock(return_value=["run-1"]) + monkeypatch.setattr( + "src.server.services.history.replay.run_lane.sr_db.list_runs_for_thread", + AsyncMock( + return_value=[ + { + "task_run_id": "run-1", + "status": "completed", + "started_at": datetime(2026, 1, 3, tzinfo=timezone.utc), + } + ] + ), + ) + + +def _task_reader(monkeypatch): + turn = _turn(0, _task_turn()) + turn.anchor.tail_checkpoint_id = "tail-0" + reader = _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=[turn]), + task_messages=[ + HumanMessage(content="p", id="sub-h-1"), + AIMessage(content="see ![c](work/sub.png)", id="sub-ai-1"), + ], + ) + _ledgered(monkeypatch, reader) + return reader + + +_RUN_FACTS = { + "turn_index": 0, + "task_run_id": "run-1", + "replay_facts": { + "v": 1, + "images": {"work/sub.png": "https://cdn/sub.png"}, + "steering_returned": [ + {"content": "also AMD", "input_id": "in-1", "reason": "run_ended"} + ], + }, +} + + +def _task_lane(items): + return [ + (i["event"], i["data"].get("content")) + for i in items + if str(i["data"].get("agent", "")).startswith("task:") + and i["event"] in ("message_chunk", "steering_returned") + ] + + +@pytest.mark.asyncio +async def test_run_facts_reach_the_projected_run(monkeypatch): + _cache_probe(monkeypatch) + _task_reader(monkeypatch) + stored_copy = _response(0, [_returned("in-1")]) + + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: stored_copy}, run_facts=[_RUN_FACTS]), + ) + + assert _task_lane(items) == [ + ("message_chunk", "see ![c](https://cdn/sub.png)"), + ("steering_returned", "also AMD"), + ] + + +@pytest.mark.asyncio +async def test_run_facts_landing_later_reproject_a_stored_turn(monkeypatch): + """A run's facts are written after it settles, by when its turn's lines + may already be stored.""" + stored_tails = _cache_probe(monkeypatch) + _task_reader(monkeypatch) + + before = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}), + ) + assert _task_lane(before) == [("message_chunk", "see ![c](work/sub.png)")] + assert stored_tails == ["tail-0"] + + after = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0)}, run_facts=[_RUN_FACTS]), + ) + assert _task_lane(after) == [ + ("message_chunk", "see ![c](https://cdn/sub.png)"), + ("steering_returned", "also AMD"), + ] + assert stored_tails == ["tail-0", "tail-0"] + + +@pytest.mark.asyncio +async def test_a_turn_with_facts_and_stored_events_replays_the_same(monkeypatch): + """The dual-write guarantee: the facts reproduce the stored durations.""" + reasoning = AIMessage( + content=[ + {"type": "thinking", "thinking": "weighing it"}, + {"type": "text", "text": "the answer"}, + ], + id="msg_1", + ) + stored = [ + _signal("lc-1", "start"), + _signal("lc-1", "complete", elapsed_ms=1500), + _signal("lc-1", "start"), + _signal("lc-1", "complete", elapsed_ms=500), + _text("lc-1", "the answer"), + ] + + def closes(items): + return [ + i["data"].get("elapsed_ms") + for i in items + if i["event"] == "message_chunk" + and i["data"].get("content_type") == "reasoning_signal" + and i["data"].get("content") == "complete" + ] + + async def replay_with(response): + _mock_reader( + monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, [reasoning])]) + ) + return await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: response}), + ) + + legacy = await replay_with(_response(0, stored)) + with_facts = _response(0, stored) + with_facts["replay_facts"] = {"v": 1, "reasoning_ms": {"msg_1": [1500, 500]}} + dual = await replay_with(with_facts) + facts_only = _response(0, []) + facts_only["replay_facts"] = with_facts["replay_facts"] + alone = await replay_with(facts_only) + + assert closes(legacy) == closes(dual) == closes(alone) == [2000] + assert legacy == dual + + +# ------------------------------------------------------------ per-turn fallback + + +def _two_turns(monkeypatch): + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=[ + _turn(0, [AIMessage(content="zero", id="ai-0")]), + _turn(1, [AIMessage(content="one", id="ai-1")]), + ], + ), + ) + real = cold.project_turn + + def flaky(inputs, turn_index, turn, response, lane, derive=None): + if turn_index == 1: + raise ValueError("a turn this build cannot project") + return real(inputs, turn_index, turn, response, lane, derive) + + monkeypatch.setattr(cold, "project_turn", flaky) + return ( + [_query(0), _query(1, content="next")], + {0: _response(0), 1: _response(1, [_text("lc-1", "one, as stored")])}, + ) + + +def _texts(lines): + return [ + (line.item()["data"].get("turn_index"), line.item()["data"]["content"]) + for line in lines + if line.event == "message_chunk" + ] + + +@pytest.mark.asyncio +async def test_a_turn_that_cannot_project_replays_alone_from_storage(monkeypatch): + queries, responses = _two_turns(monkeypatch) + page = await build_replay_page(thread_rows(queries, responses), turn_fallback=True) + assert _texts(page.lines) == [(0, "zero"), (1, "one, as stored")] + # The stored-events copy is never stored as the turn's lines: the next + # read projects it again. + assert "resp-1" not in _slice_store(monkeypatch).turns or not ( + _slice_store(monkeypatch).turns["resp-1"].get("lines") + ) + + +@pytest.mark.asyncio +async def test_without_turn_fallback_the_failure_reaches_the_caller(monkeypatch): + queries, responses = _two_turns(monkeypatch) + with pytest.raises(ValueError): + await build_replay_page(thread_rows(queries, responses)) + + +# ------------------------------------------------------------ read_replay_page + + +@pytest.mark.asyncio +async def test_auto_falls_back_to_storage_when_turns_cannot_pair(monkeypatch): + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[])) + page, source = await read_replay_page( + thread_rows([_query(0)], {0: _response(0, [_text("lc-0", "stored")])}), + ) + assert source == "sse" + assert _texts(page.lines) == [(0, "stored")] + + +@pytest.mark.asyncio +async def test_checkpoint_source_never_falls_back(monkeypatch): + _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[])) + with pytest.raises(CheckpointReplayUnavailable): + await read_replay_page( + thread_rows([_query(0)], {0: _response(0)}), + source="checkpoint", + ) + + +@pytest.mark.asyncio +async def test_a_thread_without_a_commit_pointer_replays_from_storage(monkeypatch): + reader = _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[])) + page, source = await read_replay_page( + thread_rows( + [_query(0), _query(1), _query(2)], + { + 0: _response(0, [_text("lc-0", "zero")]), + 1: _response(1, [_text("lc-1", "one")]), + 2: _response(2, [_text("lc-2", "two")]), + }, + tip=None, + ), + limit=2, + ) + assert source == "sse" + reader.thread_history.assert_not_awaited() + assert _texts(page.lines) == [(1, "one"), (2, "two")] + assert (page.first_turn_index, page.has_more) == (1, True) diff --git a/tests/unit/server/services/history/test_replay_refresh.py b/tests/unit/server/services/history/test_replay_refresh.py new file mode 100644 index 000000000..4d47e2ebf --- /dev/null +++ b/tests/unit/server/services/history/test_replay_refresh.py @@ -0,0 +1,86 @@ +"""A caller that waits on a refresh never pays for a whole-thread pass: the +turns that need one go to the background runner.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, call + +import pytest + +from src.server.services.history import replay +from src.server.services.history.replay import refresh +from tests.unit.server.services.history.replay_builders import thread_rows + +pytestmark = pytest.mark.asyncio + +_ROWS = thread_rows([], {3: {"turn_index": 3, "conversation_response_id": "r-3"}}) + + +@pytest.fixture +def wired(monkeypatch): + monkeypatch.setattr( + "src.server.database.conversation.get_replay_thread_data", + AsyncMock(return_value=_ROWS), + ) + project = AsyncMock(return_value={}) + monkeypatch.setattr(replay, "project_turns", project) + schedule = MagicMock() + monkeypatch.setattr(refresh, "schedule_replay_refresh", schedule) + return project, schedule + + +async def test_turns_that_need_the_whole_thread_are_left_to_the_runner(wired): + project, schedule = wired + project.side_effect = replay.ClaimsPassNeeded("tsk1") + + await refresh.refresh_thread("t-1", {"r-3"}, claims_pass=False) + + assert project.await_args.kwargs["claims_pass"] is False + assert project.await_args.kwargs["turn_indexes"] == [3] + schedule.assert_called_once_with("t-1", "r-3") + + +async def test_a_thread_the_checkpoints_cannot_cover_is_not_retried(wired): + project, schedule = wired + project.side_effect = replay.CheckpointReplayUnavailable("unpaired") + + await refresh.refresh_thread("t-1", {"r-3"}, claims_pass=False) + + schedule.assert_not_called() + + +async def test_the_runner_takes_the_pass_itself(wired): + project, _schedule = wired + + await refresh.refresh_thread("t-1", {"r-3"}) + + assert project.await_args.kwargs["claims_pass"] is True + + +async def test_a_newest_pass_that_needs_the_whole_thread_is_rescheduled(wired): + project, schedule = wired + project.side_effect = replay.ClaimsPassNeeded("tsk1") + + await refresh.refresh_thread("t-1", {"r-3"}, newest=True, claims_pass=False) + + assert project.await_args.kwargs["last_n_turns"] == 2 + assert schedule.call_args_list == [call("t-1"), call("t-1", "r-3")] + + +async def test_triggers_landing_before_the_runner_starts_fold_into_one_pass( + monkeypatch, +): + monkeypatch.setattr(refresh, "_runners", {}) + monkeypatch.setattr(refresh, "_pending", {}) + passes = [] + + async def refresh_thread(thread_id, response_ids, *, newest=False): + passes.append((thread_id, set(response_ids), newest)) + + monkeypatch.setattr(refresh, "refresh_thread", refresh_thread) + + refresh.schedule_replay_refresh("t-1") + refresh.schedule_replay_refresh("t-1", "r-2") + await refresh._runners["t-1"] + + assert passes == [("t-1", {"r-2"}, True)] diff --git a/tests/unit/server/services/history/test_replay_stopped.py b/tests/unit/server/services/history/test_replay_stopped.py index d31eb6083..726d8a742 100644 --- a/tests/unit/server/services/history/test_replay_stopped.py +++ b/tests/unit/server/services/history/test_replay_stopped.py @@ -1,21 +1,26 @@ -"""A stopped or failed turn replays the partial answer the user watched stream.""" +"""A stopped or failed turn replays what it committed, then its stop close or error. + +Output that streamed from a step the checkpoint never committed is not in +the model's context, so replay leaves it out.""" from __future__ import annotations import pytest from langchain_core.messages import AIMessage, ToolMessage -from src.server.services.history.reader import ThreadHistory from src.server.services.history.replay import ( build_checkpoint_replay_items, build_sse_replay_items, ) from tests.unit.server.services.history.replay_builders import ( + thread_rows, THREAD, + ThreadHistory, _cache_probe, _mock_reader, _query, _response, + _slice_store, _turn, ) @@ -112,30 +117,38 @@ def _replayed(items, turn_index=0): return rows -async def test_stopped_flash_turn_replays_its_partial_answer(monkeypatch): +def _turn_close(turn_index): + """The close for a stopped turn with no committed message to name.""" + return { + "thread_id": THREAD, + "role": "assistant", + "finish_reason": "stopped", + "turn_index": turn_index, + "response_id": f"resp-{turn_index}", + } + + +async def test_stopped_flash_turn_drops_its_uncommitted_partial(monkeypatch): """A stop inside the only model call commits nothing past the input, so - the checkpoint slice is the HumanMessage alone. The partial the user - watched stream, and the close that marks it stopped, replay from the - finalize's archive in their live order. The archived close is the turn's - only one: replay adds none beside it.""" + the checkpoint slice is the HumanMessage alone. The partial and its + archived close name a message the checkpoint never held: the turn + replays its question and a close on the turn.""" _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, [])])) stored = [*_partial_answer(), _stop_close("lc-1")] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _user_stopped(0, stored)} + thread_rows([_query(0)], {0: _user_stopped(0, stored)}), ) - assert items[0]["event"] == "user_message" - assert [i["data"] for i in items[1:]] == [ - {**row["data"], "thread_id": THREAD, "turn_index": 0, "response_id": "resp-0"} - for row in stored - ] + assert [i["event"] for i in items] == ["user_message", "message_chunk"] + assert items[1]["data"] == _turn_close(0) -async def test_stopped_ptc_turn_replays_steps_once_and_the_partial(monkeypatch): +async def test_stopped_ptc_turn_replays_committed_steps_once(monkeypatch): """Steps that finished before the stop are committed and project from the checkpoint; their stored copies only anchor. The in-flight answer exists - only in the archive and replays once, after the last step.""" + only in the archive and is dropped; the close lands on the last + committed message.""" turn_msgs = [ AIMessage( content="Pulling the filings.", @@ -162,7 +175,7 @@ async def test_stopped_ptc_turn_replays_steps_once_and_the_partial(monkeypatch): ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _user_stopped(0, stored)} + thread_rows([_query(0)], {0: _user_stopped(0, stored)}), ) assert _replayed(items) == [ @@ -172,16 +185,14 @@ async def test_stopped_ptc_turn_replays_steps_once_and_the_partial(monkeypatch): ("ai-2", "Now the margins."), ("tool_calls", "tc-2"), ("tool_call_result", "tc-2"), - ("lc-3", "The revenue"), - ("lc-3", " grew"), - ("lc-3", "stopped"), + ("ai-2", "stopped"), ] async def test_stopped_turn_replays_a_parallel_round_once(monkeypatch): """The checkpoint holds a parallel round as one tool_calls row, the live archive as one row per call. Every call's row anchors on the projected - round, so none of them replays a second time as lost output.""" + round, so none of them replays a second time.""" turn_msgs = [ AIMessage( content="Pulling both filings.", @@ -202,7 +213,7 @@ async def test_stopped_turn_replays_a_parallel_round_once(monkeypatch): stored = [text, calls_1, calls_2, result_1, result_2, _main_chunk("lc-2", "text", "Revenue")] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored, status="cancelled")} + thread_rows([_query(0)], {0: _response(0, stored, status="cancelled")}), ) assert [ @@ -215,13 +226,13 @@ async def test_stopped_turn_replays_a_parallel_round_once(monkeypatch): ("tool_calls", "tc-1"), ("tool_call_result", "tc-1"), ("tool_call_result", "tc-2"), - ("lc-2", "Revenue"), ] -async def test_stopped_tool_round_replays_the_result_that_streamed(monkeypatch): +async def test_stopped_tool_round_drops_the_uncommitted_result(monkeypatch): """A stop while parallel tools run commits the call but none of the - results; one that finished before the stop streamed live and replays.""" + results. One that finished before the stop streamed live, but the model + never sees it, so replay leaves it out.""" turn_msgs = [ AIMessage( content="Checking both.", @@ -239,17 +250,14 @@ async def test_stopped_tool_round_replays_the_result_that_streamed(monkeypatch): stored[1]["data"]["tool_calls"].append({"name": "bash", "args": {}, "id": "tc-2"}) items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _user_stopped(0, stored)} + thread_rows([_query(0)], {0: _user_stopped(0, stored)}), ) assert _replayed(items) == [ ("ai-1", "Checking both."), ("tool_calls", "tc-1"), - ("tool_call_result", "tc-1"), ("ai-1", "stopped"), ] - result = next(i for i in items if i["event"] == "tool_call_result") - assert result["data"]["content"] == "first done" def _stopped_during_search(stored_close): @@ -296,9 +304,7 @@ async def test_stop_between_messages_closes_the_committed_message( _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_search_turn()])) items = await build_checkpoint_replay_items( - THREAD, - [_query(0)], - {0: _user_stopped(0, _stopped_during_search(stored_close))}, + thread_rows([_query(0)], {0: _user_stopped(0, _stopped_during_search(stored_close))}), ) assert _replayed(items) == [ @@ -330,7 +336,7 @@ async def test_cancel_the_user_did_not_ask_for_gets_no_stop_close( response = _response(0, _stopped_during_search(False), status="cancelled") items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: response | {"metadata": metadata}} + thread_rows([_query(0)], {0: response | {"metadata": metadata}}), ) assert _replayed(items) == [ @@ -364,7 +370,7 @@ async def test_settled_main_lane_keeps_its_phantom_partial_dropped( ): """A turn that ended on a committed boundary replays the checkpoint alone. Its archive can still hold a partial from a model attempt an in-run retry - replaced; resurrecting it would double-render the answer.""" + replaced; replaying it would double-render the answer.""" _mock_reader( monkeypatch, ThreadHistory( @@ -378,35 +384,70 @@ async def test_settled_main_lane_keeps_its_phantom_partial_dropped( ] items = await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _response(0, stored, status=status)} + thread_rows([_query(0)], {0: _response(0, stored, status=status)}), ) assert _replayed(items) == [("ai-1", "The answer is 42.")] -async def test_errored_turn_replays_its_partial_before_the_error(monkeypatch): - """A model call that raises commits nothing, but the error finalize - archives what streamed in the same write that records the failure. The - partial replays, then the terminal error. No stop close exists here.""" +async def test_a_dropped_partial_leaves_later_stored_rows_in_place(monkeypatch): + """Rows that streamed after a partial no checkpoint kept still replay + where they streamed, not at the end of the turn.""" + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=[ + _turn( + 0, + [ + AIMessage( + content="Searching", + id="ai-1", + tool_calls=[{"name": "bash", "args": {}, "id": "call-1"}], + ), + ToolMessage(content="out", tool_call_id="call-1"), + ], + ) + ], + ), + ) + step = _main_step("lc-1", "Searching", "call-1", "out") + stored = [ + *step[:2], + _main_chunk("lc-phantom", "text", "Sear"), + {"event": "context_window", "data": {"agent": _MAIN_AGENT, "action": "noop"}}, + step[2], + ] + + items = await build_checkpoint_replay_items( + thread_rows([_query(0)], {0: _response(0, stored)}), + ) + + assert _replayed(items) == [ + ("ai-1", "Searching"), + ("tool_calls", "call-1"), + ("context_window", None), + ("tool_call_result", "call-1"), + ] + + +async def test_errored_turn_replays_the_error_without_its_partial(monkeypatch): + """A model call that raises commits nothing. What streamed before it + raised is dropped; the terminal error replays. No stop close exists + here.""" _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[_turn(0, [])])) response = _response(0, _partial_answer(), status="error") response["errors"] = ["provider exploded"] - items = await build_checkpoint_replay_items(THREAD, [_query(0)], {0: response}) + items = await build_checkpoint_replay_items(thread_rows([_query(0)], {0: response})) - assert _replayed(items) == [ - ("lc-1", "start"), - ("lc-1", "Reading the filing"), - ("lc-1", "complete"), - ("lc-1", "The answer"), - ("lc-1", " is"), - ("error", "provider exploded"), - ] + assert _replayed(items) == [("error", "provider exploded")] -async def test_stopped_turn_without_a_boundary_replays_its_partial(monkeypatch): +async def test_stopped_turn_without_a_boundary_replays_as_a_stub(monkeypatch): """A stopped turn whose finalize could not advance the commit pointer has - no slice and replays as a stub. What it streamed still replays.""" + no slice and replays as a stub: its question and a close on the turn.""" _mock_reader( monkeypatch, ThreadHistory( @@ -417,17 +458,12 @@ async def test_stopped_turn_without_a_boundary_replays_its_partial(monkeypatch): stored = [*_partial_answer(), _stop_close("lc-1")] items = await build_checkpoint_replay_items( - THREAD, - [_query(0), _query(1)], - {0: _response(0), 1: _user_stopped(1, stored)}, + thread_rows([_query(0), _query(1)], {0: _response(0), 1: _user_stopped(1, stored)}), ) turn1 = [i for i in items if i["data"].get("turn_index") == 1] - assert turn1[0]["event"] == "user_message" - assert [i["data"] for i in turn1[1:]] == [ - {**row["data"], "thread_id": THREAD, "turn_index": 1, "response_id": "resp-1"} - for row in stored - ] + assert [i["event"] for i in turn1] == ["user_message", "message_chunk"] + assert turn1[1]["data"] == _turn_close(1) async def test_stop_before_the_first_assistant_event_closes_the_turn(monkeypatch): @@ -444,45 +480,30 @@ async def test_stop_before_the_first_assistant_event_closes_the_turn(monkeypatch queries = [_query(0), _query(1)] responses = {0: _response(0), 1: _user_stopped(1)} - checkpoint = await build_checkpoint_replay_items(THREAD, queries, responses) + checkpoint = await build_checkpoint_replay_items(thread_rows(queries, responses)) stored = build_sse_replay_items(THREAD, queries, responses) for items in (checkpoint, stored): turn1 = [i for i in items if i["data"].get("turn_index") == 1] assert [i["event"] for i in turn1] == ["user_message", "message_chunk"] - assert turn1[1]["data"] == { - "thread_id": THREAD, - "role": "assistant", - "finish_reason": "stopped", - "turn_index": 1, - "response_id": "resp-1", - } + assert turn1[1]["data"] == _turn_close(1) -async def test_stopped_main_lane_caches_with_its_partial(monkeypatch): - """A stopped turn's partial lands in the same finalize write as its - status, so the turn owes no archive and caches, partial included.""" - _cache_probe(monkeypatch) - segments: list[list[dict]] = [] - - async def capture(thread_id, tail_checkpoint_id, fingerprint, items): - segments.append(items) - - from src.server.services.history import replay as replay_module - - monkeypatch.setattr(replay_module.projection_cache, "store_turn", capture) +async def test_stopped_turn_caches_with_its_close(monkeypatch): + """A stopped turn owes no later write, so it caches as it replays.""" + stored_tails = _cache_probe(monkeypatch) turn = _turn(0, []) - turn.tail_checkpoint_id = "tail-0" + turn.anchor.tail_checkpoint_id = "tail-0" _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn])) stored = [*_partial_answer(), _stop_close("lc-1")] await build_checkpoint_replay_items( - THREAD, [_query(0)], {0: _user_stopped(0, stored)} + thread_rows([_query(0)], {0: _user_stopped(0, stored)}), ) - assert len(segments) == 1 - assert _replayed(segments[0])[-3:] == [ - ("lc-1", "The answer"), - ("lc-1", " is"), - ("lc-1", "stopped"), - ] + from src.server.services.history.replay import lines + + assert stored_tails == ["tail-0"] + row = _slice_store(monkeypatch).turns["resp-0"] + segment = [line.item() for line in lines.decode(row["lines"])] + assert _replayed(segment) == [(None, "stopped")] diff --git a/tests/unit/server/services/history/test_replay_store.py b/tests/unit/server/services/history/test_replay_store.py new file mode 100644 index 000000000..ec0717e63 --- /dev/null +++ b/tests/unit/server/services/history/test_replay_store.py @@ -0,0 +1,679 @@ +"""What replay stores, and when a stored row is trusted. + +A stored run slice is trusted as final forever, stored lines as the turn's +replay while their key matches, and stored claims as each legacy launch's +run: each may only be written from evidence that outlives the read. +""" + +from __future__ import annotations + +import dataclasses +from unittest.mock import AsyncMock, patch + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from src.server.services.history import replay, slices, task_streams +from src.server.services.history.slices import SpanDelta +from src.server.services.history.replay import ( + build_checkpoint_replay_items, + build_replay_page, + run_lane, +) +from src.server.services.history.replay import cold +from src.server.services.history.replay import lines as replay_lines +from src.server.services.history.replay.turn import Inputs +from tests.unit.server.services.history.replay_builders import ( + thread_rows, + THREAD, + ThreadHistory, + _mock_reader, + _query, + _response, + _slice_store, + _turn, +) + +pytestmark = pytest.mark.asyncio + + +def _launch(ordinal, prompt, *, action="init", run_id=None, task_id="tsk1"): + artifact = { + "task_id": task_id, + "action": action, + "description": "d", + "prompt": prompt, + } + if run_id: + artifact["task_run_id"] = run_id + return [ + AIMessage( + content="", + id=f"ai-{ordinal}", + tool_calls=[{"name": "Task", "args": {}, "id": f"tc-{ordinal}"}], + ), + ToolMessage( + content="dispatched", + tool_call_id=f"tc-{ordinal}", + name="Task", + id=f"tm-{ordinal}", + additional_kwargs={"task_artifact": artifact}, + ), + ] + + +def _plain(ordinal): + return [AIMessage(content=f"answer {ordinal}", id=f"ai-{ordinal}")] + + +def _turns(messages_by_turn): + out = [] + for ordinal, messages in enumerate(messages_by_turn): + turn = _turn(ordinal, messages) + turn.anchor.tail_checkpoint_id = f"tail-{ordinal}" + out.append(turn) + return out + + +def _rows(n): + return [_query(i) for i in range(n)], {i: _response(i) for i in range(n)} + + +def _seal_evidence( + monkeypatch, *, probe=None, stream=task_streams.STREAM_SEALED, ledger=None +): + """Liveness probe, task stream states and ledger rows for the lane.""" + + async def streams(thread_id, task_ids): + return dict.fromkeys(task_ids, stream) + + monkeypatch.setattr( + run_lane, "resolve_task_details", probe or AsyncMock(return_value={}) + ) + monkeypatch.setattr(task_streams, "task_stream_states", streams) + monkeypatch.setattr( + run_lane.sr_db, + "list_runs_for_thread", + ledger if ledger is not None else AsyncMock(return_value=[]), + ) + + +def _task_text(items): + return [ + i["data"]["content"] + for i in items + if i["event"] == "message_chunk" + and i["data"].get("agent") == "task:tsk1" + and i["data"].get("role") == "assistant" + ] + + +async def test_a_run_final_only_by_a_failed_probe_is_never_stored(monkeypatch): + """The liveness probe fails open, so a legacy run still writing reads as + finished. Its slice serves that read only: stored, it would be trusted + as final ever after, and the turn's lines sealed with the run cut short. + """ + store = _slice_store(monkeypatch) + _seal_evidence( + monkeypatch, + probe=AsyncMock(side_effect=RuntimeError("probe down")), + stream=task_streams.STREAM_OPEN, + ) + reader = _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=_turns([_launch(0, "go")])), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="partial", id="sub-ai-1"), + ], + ) + queries, responses = _rows(1) + + first = await build_checkpoint_replay_items(thread_rows(queries, responses)) + assert _task_text(first) == ["partial"] + assert store.runs == {} + assert store.lined == [] + + # The run wrote on; the next read cuts it again rather than trusting a + # copy from the first. + reader.aget_task_history.return_value.messages.append( + AIMessage(content="more", id="sub-ai-2") + ) + second = await build_checkpoint_replay_items(thread_rows(queries, responses)) + assert _task_text(second) == ["partial", "more"] + assert store.runs == {} + assert store.lined == [] + + # The stream's end sentinel settles it without the probe. + _seal_evidence(monkeypatch, probe=AsyncMock(side_effect=RuntimeError("probe down"))) + third = await build_checkpoint_replay_items(thread_rows(queries, responses)) + assert _task_text(third) == ["partial", "more"] + assert [r["tail_checkpoint_id"] for r in store.runs.values()] == ["tsk1:tail-0"] + assert store.lined == ["tail-0"] + + +async def test_nothing_resolved_without_the_ledger_is_stored(monkeypatch): + """With the ledger unreadable every run's end rests on the probe: no run + slice, no claims and no lines of a launching turn are stored. A turn + that launched nothing never asked the ledger and stores as usual.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch, ledger=AsyncMock(side_effect=RuntimeError("down"))) + _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=_turns([_launch(0, "go"), _plain(1)])), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + ], + ) + queries, responses = _rows(2) + + items = await build_checkpoint_replay_items(thread_rows(queries, responses)) + + assert _task_text(items) == ["done"] + assert store.runs == {} + assert store.turns["resp-0"]["lines"] is None + assert store.turns["resp-0"]["run_claims"] is None + assert store.lined == ["tail-1"] + + +async def test_a_stored_run_the_walk_ends_elsewhere_is_cut_again(monkeypatch): + """A run that wrote on after its slice was stored no longer ends at the + stored tail; once the task is walked, that slice is dropped and the run + cut again.""" + store = _slice_store(monkeypatch) + _seal_evidence( + monkeypatch, + ledger=AsyncMock( + return_value=[ + {"task_run_id": "run-1", "status": "completed", "started_at": None}, + {"task_run_id": "run-2", "status": "completed", "started_at": None}, + ] + ), + ) + reader = _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=_turns( + [ + _launch(0, "first", run_id="run-1"), + _launch(1, "second", action="resume", run_id="run-2"), + ] + ), + ), + task_messages=[ + HumanMessage(content="first", id="sub-h-1"), + AIMessage(content="first done", id="sub-ai-1"), + HumanMessage(content="second", id="sub-h-2"), + AIMessage(content="second done", id="sub-ai-2"), + ], + ) + reader.task_run_stamps = AsyncMock(return_value=["run-1", "run-2"]) + codec, data = slices.encode( + reader.serde, + SpanDelta( + messages=[ + HumanMessage(content="first", id="sub-h-1"), + AIMessage(content="cut short", id="sub-ai-0"), + ] + ), + ) + store.runs[("tsk1", "tsk1:in-0")] = { + "task_id": "tsk1", + "input_checkpoint_id": "tsk1:in-0", + "task_run_id": "run-1", + "tail_checkpoint_id": "tsk1:tail-before", + "slice_key": slices.SLICE_KEY, + "slice_codec": codec, + "slice": data, + } + queries, responses = _rows(2) + + items = await build_checkpoint_replay_items(thread_rows(queries, responses)) + + assert _task_text(items) == ["first done", "second done"] + assert store.runs[("tsk1", "tsk1:in-0")]["tail_checkpoint_id"] == "tsk1:tail-0" + + +async def test_the_claims_pass_stores_the_settled_runs_outside_the_page( + monkeypatch, +): + """The pass reads every task's seal evidence, not only the evidence of + the batch that needed it, so a sealed run it claims for a turn outside + the page is stored rather than cut again by the read that pages to it.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=_turns([_launch(0, "go"), _launch(1, "go", task_id="tsk2")]), + ), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + ], + ) + queries, responses = _rows(2) + + page = await build_replay_page(thread_rows(queries, responses), limit=1) + + assert page.first_turn_index == 1 + assert sorted(store.runs) == [("tsk1", "tsk1:in-0"), ("tsk2", "tsk2:in-0")] + assert store.turns["resp-0"]["run_claims"] == ["tsk1:in-0"] + assert store.lined == ["tail-1"] + + +async def test_a_claim_on_a_run_the_walk_no_longer_holds_is_unavailable( + monkeypatch, +): + """Stored claims name the run's input checkpoint. When the walk no + longer has that run and nothing stored its slice, its transcript cannot + be read, and the read says so rather than cut a slice it cannot bound.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + turn = _turns([_launch(0, "go")])[0] + reader = _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=[turn]), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + ], + ) + anchors, _tip = await reader.aget_turn_anchors(THREAD) + codec, data = slices.encode(reader.serde, turn) + store.turns["resp-0"] = { + "response_id": "resp-0", + "input_checkpoint_id": anchors[0].input_checkpoint_id, + "tail_checkpoint_id": anchors[0].tail_checkpoint_id, + "slice_key": slices.SLICE_KEY, + "slice_codec": codec, + "slice": data, + "lines_key": None, + "lines": None, + "run_claims": ["tsk1:in-gone"], + } + queries, responses = _rows(1) + + with pytest.raises(replay.CheckpointReplayUnavailable): + await build_checkpoint_replay_items(thread_rows(queries, responses)) + reader.aget_run_slices.assert_not_awaited() + assert store.runs == {} + + +async def test_claims_count_only_beside_their_slice(monkeypatch): + """Claims were placed against the slice stored beside them; once the + branch no longer agrees with that slice, neither are they trusted.""" + store = _slice_store(monkeypatch) + turn = _turns([_launch(0, "go")])[0] + reader = _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[turn])) + anchors, _tip = await reader.aget_turn_anchors(THREAD) + queries, responses = _rows(1) + inputs = Inputs.of(thread_rows(queries, responses)) + codec, data = slices.encode(reader.serde, turn) + store.turns["resp-0"] = { + "response_id": "resp-0", + "input_checkpoint_id": anchors[0].input_checkpoint_id, + "tail_checkpoint_id": anchors[0].tail_checkpoint_id, + "slice_key": slices.SLICE_KEY, + "slice_codec": codec, + "slice": data, + "lines_key": None, + "lines": None, + "run_claims": ["tsk1:in-0"], + } + + _turns_read, claims, fresh = await cold.load_turns( + reader, inputs, [(0, anchors[0])] + ) + assert claims == {0: ["tsk1:in-0"]} + assert fresh == set() + + store.turns["resp-0"]["tail_checkpoint_id"] = "tail-moved" + _turns_read, claims, fresh = await cold.load_turns( + reader, inputs, [(0, anchors[0])] + ) + assert claims == {0: None} + assert fresh == {0} + + +async def test_undecodable_stored_rows_are_rebuilt(monkeypatch): + """A row that no longer decodes costs a rebuild of that row, never the + read: it is treated as stale and stored over.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + reader = _mock_reader( + monkeypatch, + ThreadHistory(thread_id=THREAD, turns=_turns([_launch(0, "go"), _plain(1)])), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + ], + ) + queries, responses = _rows(2) + first = await build_checkpoint_replay_items(thread_rows(queries, responses)) + + store.turns["resp-0"]["lines"] = "no tabs here" + store.turns["resp-0"]["slice"] = b"\x00not a slice" + store.turns["resp-1"]["lines"] = "no tabs here either" + (run_key,) = store.runs + store.runs[run_key]["slice"] = b"\x00not a run slice" + + second = await build_checkpoint_replay_items(thread_rows(queries, responses)) + + assert second == first + for row in store.turns.values(): + replay_lines.decode(row["lines"]) + slices.decode( + reader.serde, store.turns["resp-0"]["slice_codec"], store.turns["resp-0"]["slice"] + ) + slices.decode( + reader.serde, store.runs[run_key]["slice_codec"], store.runs[run_key]["slice"] + ) + + +async def test_a_cold_open_writes_each_turn_once(monkeypatch): + """The claims pass reads every turn; the ones it cuts from checkpoints are + projected from memory and written once, slice and lines together, a + statement per batch. A pass that starts mid-thread waits for the stores + before it, and the request returns only once every store has landed.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + n = 20 + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=_turns( + [_plain(i) for i in range(n - 2)] + [_launch(18, "go"), _plain(19)] + ), + ), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + ], + ) + queries, responses = _rows(n) + + await build_checkpoint_replay_items(thread_rows(queries, responses)) + + assert sorted(store.written) == sorted(f"resp-{i}" for i in range(n)) + assert [kind for kind, _ in store.statements] == [ + "upsert_turn_slices", + "upsert_run_slices", + "upsert_turn_slices", + ] + assert sorted(store.lined) == sorted(f"tail-{i}" for i in range(n)) + assert store.turns["resp-18"]["run_claims"] == ["tsk1:in-0"] + + +async def test_a_paged_cold_open_stores_the_pass_once(monkeypatch): + """A page's claims pass cuts turns outside the page too: those store + their slices and claims once, the page's turns their lines as well.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + reader = _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=_turns( + [ + _launch(0, "go"), + _plain(1), + _plain(2), + _launch(3, "again", action="resume"), + ] + ), + ), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + HumanMessage(content="again", id="sub-h-2"), + AIMessage(content="done again", id="sub-ai-2"), + ], + ) + queries, responses = _rows(4) + + page = await build_replay_page(thread_rows(queries, responses), limit=2) + + assert page.first_turn_index == 2 + assert sorted(store.written) == ["resp-0", "resp-1", "resp-2", "resp-3"] + assert store.turns["resp-0"]["run_claims"] == ["tsk1:in-0"] + assert store.turns["resp-3"]["run_claims"] == ["tsk1:in-1"] + assert sorted(store.lined) == ["tail-2", "tail-3"] + assert sorted(store.runs) == [("tsk1", "tsk1:in-0"), ("tsk1", "tsk1:in-1")] + + # The next page reads the pass's work back: no checkpoint is cut again. + before = len(store.statements) + cuts = reader.aget_turn_slices.await_count, reader.aget_run_slices.await_count + older = await build_replay_page( + thread_rows(queries, responses), + before_turn=2, + limit=2, + ) + assert older.first_turn_index == 0 + assert ( + reader.aget_turn_slices.await_count, + reader.aget_run_slices.await_count, + ) == cuts + assert [kind for kind, _ in store.statements[before:]] == ["update_turns"] + + +async def test_an_uncached_replay_reads_and_writes_no_stored_row(monkeypatch): + """``cache=False`` cuts every turn and run from checkpoints, though + current rows are stored, and stores nothing: what a parity run compares + is what the code projects today.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + reader = _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=_turns([_launch(0, "go"), _plain(1)]), + ), + task_messages=[ + HumanMessage(content="go", id="sub-h-1"), + AIMessage(content="done", id="sub-ai-1"), + ], + ) + rows = thread_rows(*_rows(2)) + cached = await build_checkpoint_replay_items(rows) + assert store.lined and store.runs + + for name in ("get_turn_lines", "get_turn_slices", "get_run_slices"): + monkeypatch.setattr( + cold.slices_db, name, AsyncMock(side_effect=AssertionError(name)) + ) + before = len(store.statements) + cuts = reader.aget_turn_slices.await_count, reader.aget_run_slices.await_count + + assert await build_checkpoint_replay_items(rows, cache=False) == cached + assert len(store.statements) == before + assert reader.aget_turn_slices.await_count > cuts[0] + assert reader.aget_run_slices.await_count > cuts[1] + + +async def test_a_change_to_the_cut_moves_the_slice_key(monkeypatch, tmp_path): + """The slice key digests the code that cuts a slice, so a fix there + rebuilds the slices stored before it with no manual bump.""" + copies = [] + for source in slices.slice_sources(): + copy = tmp_path / source.relative_to(slices._SRC) + copy.parent.mkdir(parents=True, exist_ok=True) + copy.write_bytes(source.read_bytes()) + copies.append(copy) + monkeypatch.setattr(slices, "_SRC", tmp_path) + monkeypatch.setattr(slices, "slice_sources", lambda: copies) + before = slices._slice_key() + assert before == slices.SLICE_KEY + + cut = copies[0] + cut.write_text( + cut.read_text().replace("has_input: bool = True", "has_input: bool = False") + ) + assert slices._slice_key() != before + assert before.startswith(f"{slices.SLICE_VERSION}.") + + +# ------------------------------------------------------------ fallback scope + + +def _three_batches(monkeypatch): + """A batch launching a run, one launching nothing, and one more launching, + every launch unledgered so only the claims pass can place it.""" + size = cold.PROJECT_BATCH + n = 2 * size + 2 + launches = {0: "go", 2 * size: "again"} + _mock_reader( + monkeypatch, + ThreadHistory( + thread_id=THREAD, + turns=_turns( + [ + _launch(i, launches[i]) if i in launches else _plain(i) + for i in range(n) + ] + ), + ), + ) + return size, n + + +def _answers(page): + return { + line.item()["data"]["content"] + for line in page.lines + if line.event == "message_chunk" + and str(line.item()["data"].get("content", "")).startswith("answer ") + } + + +async def test_a_failed_claims_pass_runs_once_per_read(monkeypatch): + """Retrying turn by turn cannot place a run the whole-thread pass could + not: the batches that need it replay from storage, the one that does not + projects, and the pass is never run again in the read.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + size, n = _three_batches(monkeypatch) + place = AsyncMock(side_effect=RuntimeError("a ledger this build cannot read")) + monkeypatch.setattr(cold, "_place_launches", place) + queries, responses = _rows(n) + + page = await build_replay_page(thread_rows(queries, responses), turn_fallback=True) + + assert place.await_count == 1 + assert _answers(page) == {f"answer {i}" for i in range(size, 2 * size)} + assert sorted(store.lined) == sorted(f"tail-{i}" for i in range(size, 2 * size)) + + +async def test_a_failed_claims_pass_is_unavailable_without_turn_fallback( + monkeypatch, +): + """A read that may not fall back reports the thread unavailable (the + endpoint's 409), not a failure.""" + _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + _size, n = _three_batches(monkeypatch) + place = AsyncMock(side_effect=RuntimeError("a ledger this build cannot read")) + monkeypatch.setattr(cold, "_place_launches", place) + queries, responses = _rows(n) + + with pytest.raises(replay.CheckpointReplayUnavailable): + await replay.read_replay_page( + thread_rows(queries, responses), + source="checkpoint", + ) + + +async def test_a_failed_claims_pass_answers_409_for_a_checkpoint_replay( + monkeypatch, threads_client +): + """The endpoint's side of the contract above: a thread the pass cannot + place is unavailable from checkpoints, not a server error.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + _size, n = _three_batches(monkeypatch) + place = AsyncMock(side_effect=RuntimeError("a ledger this build cannot read")) + monkeypatch.setattr(cold, "_place_launches", place) + queries, responses = _rows(n) + rows = dataclasses.replace( + thread_rows(queries, responses), owner_id="test-user-123" + ) + + with patch( + "src.server.app.threads.messaging.get_replay_thread_data", + new=AsyncMock(return_value=rows), + ): + resp = await threads_client.get( + f"/api/v1/threads/{THREAD}/messages/replay", + params={"source": "checkpoint"}, + ) + + assert resp.status_code == 409 + assert resp.json()["detail"].startswith("Checkpoint replay unavailable") + place.assert_awaited_once() + assert store.lined == [] + + +async def test_a_caller_that_cannot_wait_is_never_made_to_read_the_whole_thread( + monkeypatch, +): + """A turn whose launch only the claims pass can place raises instead of + running it; a turn that needs no pass still projects.""" + store = _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + _size, n = _three_batches(monkeypatch) + place = AsyncMock(side_effect=AssertionError("ran the claims pass")) + monkeypatch.setattr(cold, "_place_launches", place) + queries, responses = _rows(n) + + with pytest.raises(replay.ClaimsPassNeeded): + await replay.project_turns( + thread_rows(queries, responses), + "cp-tip", + turn_indexes=[0], + claims_pass=False, + ) + lines = await replay.project_turns( + thread_rows(queries, responses), + "cp-tip", + turn_indexes=[1], + claims_pass=False, + ) + + place.assert_not_awaited() + assert list(lines) == [1] + assert store.lined == ["tail-1"] + + +async def test_a_database_fault_ends_the_read_without_per_turn_retries( + monkeypatch, +): + from src.server.database.conversation import turn_slices as slices_db + from src.server.database.pool import AppDataPoolTimeout + + _slice_store(monkeypatch) + _seal_evidence(monkeypatch) + _size, n = _three_batches(monkeypatch) + reads = AsyncMock(side_effect=AppDataPoolTimeout("couldn't get a connection")) + monkeypatch.setattr(slices_db, "get_turn_slices", reads) + queries, responses = _rows(n) + + # Nor does it fall back to storage, which waits on the same pool. + with pytest.raises(AppDataPoolTimeout): + await replay.read_replay_page(thread_rows(queries, responses)) + assert reads.await_count == 1 + + +async def test_a_busy_checkpointer_pool_still_falls_back_to_storage(monkeypatch): + from psycopg_pool import PoolTimeout + + reader = _mock_reader(monkeypatch, ThreadHistory(thread_id=THREAD, turns=[])) + reader.thread_history.side_effect = PoolTimeout("couldn't get a connection") + queries, responses = _rows(1) + + # Storage is read through the app-data pool, which this wait never held. + page, source = await replay.read_replay_page(thread_rows(queries, responses)) + assert (source, page.first_turn_index) == ("sse", 0) diff --git a/tests/unit/server/services/history/test_slices.py b/tests/unit/server/services/history/test_slices.py new file mode 100644 index 000000000..25a3587d2 --- /dev/null +++ b/tests/unit/server/services/history/test_slices.py @@ -0,0 +1,101 @@ +"""Which messages a turn added: the diff between a turn's start and end state.""" + +from __future__ import annotations + +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from src.server.services.history.slices import _new_messages + + +def _state(*messages): + return {"messages": list(messages)} + + +def _human(content, id=None): + message = HumanMessage(content=content) + message.id = id + return message + + +def test_an_idless_message_is_new_only_once(): + """A message saved without an id loads id-less from some checkpoints and + with an id from others; either way it is the same message, not new in + every turn that follows it.""" + legacy = [_human("ask"), _human("continue")] + asked = _human("continue", id="h-new") + + assert _new_messages(_state(*legacy), _state(*legacy, asked)) == [asked] + # The same messages with ids at the end, and without them at the start. + ided = [_human("ask", id="a"), _human("continue", id="c")] + assert _new_messages(_state(*legacy), _state(*ided, asked)) == [asked] + assert _new_messages(_state(*ided), _state(*legacy, asked)) == [asked] + + +def test_a_repeat_after_compaction_is_new(): + """Compaction drops the first "continue"; the user's next one is new.""" + first = _human("continue", id="c1") + again = _human("continue", id="c2") + + assert _new_messages(_state(first), _state(again)) == [again] + + +def test_a_repeat_of_a_compacted_legacy_message_is_new(): + """Compaction removed the id-less legacy messages, and the user's next + input repeats one of them under a fresh id: it is the turn's input, not + the old message reloaded, because carried messages precede the input.""" + legacy = [_human("ask"), _human("continue")] + summary = HumanMessage( + content="summary of the work so far", + id="s", + additional_kwargs={"lc_source": "summarization"}, + ) + again = _human("continue", id="c2") + answer = AIMessage(content="carrying on", id="ai-1") + + assert _new_messages(_state(*legacy), _state(summary, again, answer)) == [ + summary, + again, + answer, + ] + + +def test_a_resume_has_no_input_to_cut_at(): + """A HITL resume adds no message of its own, so its carried legacy + messages still match by content wherever they sit.""" + legacy = [_human("ask"), AIMessage(content="which one?")] + ided = [_human("ask", id="a"), AIMessage(content="which one?", id="q")] + answer = AIMessage(content="done", id="ai-2") + + assert _new_messages( + _state(*legacy), _state(*ided, answer), has_input=False + ) == [answer] + + +def test_a_resume_that_trimmed_the_window_keeps_its_new_messages(): + """An orchestrator re-entry passes ``before_agent``, so the window can + trim inside a resume span, where no input bounds the content match. The + id-less legacy messages it dropped are not in the end state, so a new + message that repeats one of them is new; a kept one the trim stamped a + fresh id on is still the same message.""" + call = {"name": "bash", "args": {"c": "ls"}, "id": "c1"} + dropped = [ + _human("continue"), + AIMessage(content="", tool_calls=[call]), + ToolMessage(content="ok", tool_call_id="c1"), + ] + for message in dropped: + message.id = None + kept = [_human("run two"), AIMessage(content="answer two", id="a2")] + stamped = [_human("run two", id="h2-stamped"), kept[1]] + again = {"name": "bash", "args": {"c": "pwd"}, "id": "c3"} + new = [ + AIMessage(content="", id="a3", tool_calls=[again]), + ToolMessage(content="ok", tool_call_id="c3", id="t3"), + _human("continue", id="h-notice"), + AIMessage(content="final", id="a4"), + ] + + got = _new_messages( + _state(*dropped, *kept), _state(*stamped, *new), has_input=False + ) + assert [m.id for m in got] == [m.id for m in new] diff --git a/tests/unit/server/services/history/test_window.py b/tests/unit/server/services/history/test_window.py new file mode 100644 index 000000000..c8c3ede61 --- /dev/null +++ b/tests/unit/server/services/history/test_window.py @@ -0,0 +1,455 @@ +"""The server's side of the window: replay reads the same turns from a trimmed +thread, the coverage the agent asks before a trim, the trimmed runs read back +for a transcript with no stored copy to carry them, and the commit pointer a +manual compaction moves. + +The thread is ``window_harness``'s scenario, run with the trim and without +it on the real graph and checkpointer, read by the real history reader.""" + +from __future__ import annotations + +import asyncio +import json +from dataclasses import dataclass + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph import START, MessagesState, StateGraph +from langgraph.graph.message import REMOVE_ALL_MESSAGES +from langsmith import tracing_context + +from ptc_agent.agent.state import DeltaAgentState +from ptc_agent.agent.transcript import ( + EarlierTurnsMissing, + TranscriptTarget, + Window, + build_directory, +) +from ptc_agent.agent.transcript.render import split_runs +from src.server.services.history import slices +from src.server.database.conversation import turn_slices as slices_db +from src.server.services import transcripts +from src.server.services.history import replay +from src.server.services.history import window as window_module +from src.server.services.history.reader import CheckpointHistoryReader +from tests.unit.middleware.compaction import window_harness as h +from tests.unit.server.services.history.replay_builders import ( + thread_rows, + SliceStore, + _cache_probe, + _query, + _response, +) + +pytestmark = pytest.mark.asyncio + +TURNS = len(h.SCENARIO) + + +@dataclass +class Thread: + run: h.ThreadRun + saver: InMemorySaver + + @property + def tip(self) -> str: + return self.run.tails[-1] + + @property + def base(self) -> int: + return self.run.held[-1][1] + + def head(self, runs: int) -> Window: + """The thread's first ``runs`` runs, as a trim of them would leave + its head.""" + whole = h.whole_threads(self.run)[-1] + return Window().extend([m for run in split_runs(whole)[:runs] for m in run]) + + async def trimmed(self) -> Window: + """What the window left trimmed at the tip.""" + values = (await self.run.graph.aget_state(h.config())).values + return Window.of(values) + + +@pytest.fixture(scope="module") +def threads() -> tuple[Thread, Thread]: + async def run() -> tuple[Thread, Thread]: + off, on = InMemorySaver(), InMemorySaver() + return ( + Thread(await h.run_scenario(off, window=False), off), + Thread(await h.run_scenario(on, window=True), on), + ) + + return asyncio.run(run()) + + +def _serve(monkeypatch: pytest.MonkeyPatch, thread: Thread, store: SliceStore) -> None: + """The real reader over ``thread``'s checkpoints, and ``store`` for the + slice tables, the anchor-keyed reads the window makes included.""" + reader = CheckpointHistoryReader(thread.saver) + monkeypatch.setattr( + CheckpointHistoryReader, "get_instance", classmethod(lambda cls: reader) + ) + for name in ( + "get_turn_lines", + "get_turn_slices", + "has_turn_slices", + "get_slices_at", + "upsert_turn_slices", + "update_turns", + "get_run_slices", + "upsert_run_slices", + ): + monkeypatch.setattr(slices_db, name, getattr(store, name)) + + async def pointer(thread_id): + return thread.tip + + monkeypatch.setattr(window_module, "get_thread_checkpoint_id", pointer) + + +async def _page(monkeypatch, thread: Thread, store: SliceStore) -> str: + _serve(monkeypatch, thread, store) + page = await replay.build_replay_page( + thread_rows([_query(i, content=f"Question {i + 1}") for i in range(TURNS)], {i: _response(i) for i in range(TURNS)}, thread_id=h.THREAD), + ) + return h.Normalizer()(json.dumps([line.item() for line in page.lines], default=str)) + + +def _text(messages) -> str: + return h.Normalizer()(h.dump(list(messages))) + + +# ------------------------------------------------------------ replay + + +async def test_replay_reads_a_trimmed_thread_as_the_whole_one( + monkeypatch, threads +) -> None: + off, on = threads + _cache_probe(monkeypatch) + stores = SliceStore(), SliceStore() + cold = ( + await _page(monkeypatch, off, stores[0]), + await _page(monkeypatch, on, stores[1]), + ) + assert cold[1] == cold[0] + assert len(stores[1].turns) == TURNS + warm = ( + await _page(monkeypatch, off, stores[0]), + await _page(monkeypatch, on, stores[1]), + ) + assert warm == cold + + +# ------------------------------------------------------------ coverage + + +async def test_the_agent_and_the_server_name_the_same_head(threads) -> None: + _, on = threads + trimmed = await on.trimmed() + assert trimmed.runs == on.base > 0 + assert trimmed == on.head(on.base) + + +async def test_coverage_holds_once_the_slices_are_stored(monkeypatch, threads) -> None: + _, on = threads + _cache_probe(monkeypatch) + store = SliceStore() + _serve(monkeypatch, on, store) + assert await window_module.runs_held(h.THREAD, Window()) + assert not await window_module.runs_held(h.THREAD, on.head(on.base)) + await _page(monkeypatch, on, store) + assert await window_module.runs_held(h.THREAD, on.head(on.base)) + whole = on.head(TURNS) + assert await window_module.runs_held(h.THREAD, whole) + beyond = Window(TURNS + 1, whole.messages + 1, whole.digest) + assert not await window_module.runs_held(h.THREAD, beyond) + other = Window().extend([HumanMessage("Question 1, edited", id="h-1")]) + assert not await window_module.runs_held(h.THREAD, other) + + +async def test_a_slice_from_another_branch_does_not_count(monkeypatch, threads) -> None: + _, on = threads + _cache_probe(monkeypatch) + store = SliceStore() + await _page(monkeypatch, on, store) + row = store.turns["resp-1"] + row["tail_checkpoint_id"] = "another-branch" + assert await window_module.runs_held(h.THREAD, on.head(1)) + assert not await window_module.runs_held(h.THREAD, on.head(2)) + + +async def test_a_slice_replay_would_not_serve_does_not_count( + monkeypatch, threads +) -> None: + """A trim leaves the slices as the trimmed runs' only copy, so a slice + counts only while replay would read it: under the running slice key, and + decodable. The read back takes such a turn from its checkpoints.""" + _, on = threads + _cache_probe(monkeypatch) + store = SliceStore() + await _page(monkeypatch, on, store) + head = on.head(on.base) + extracted = await window_module.earlier_runs(h.THREAD, on.tip, head) + row = store.turns["resp-0"] + assert await window_module.runs_held(h.THREAD, head) + + key, row["slice_key"] = row["slice_key"], "an-older-key" + assert not await window_module.runs_held(h.THREAD, head) + row["slice_key"], row["slice"] = key, b"undecodable" + assert not await window_module.runs_held(h.THREAD, head) + read_back = await window_module.earlier_runs(h.THREAD, on.tip, head) + assert read_back is not None and _text(read_back) == _text(extracted) + + +# ------------------------------------------------------------ the trimmed runs + + +async def test_the_trimmed_runs_read_back_as_they_were(monkeypatch, threads) -> None: + """From stored slices where they are, from the checkpoints where not.""" + off, on = threads + _cache_probe(monkeypatch) + store = SliceStore() + _serve(monkeypatch, on, store) + extracted = await window_module.earlier_runs(h.THREAD, on.tip, await on.trimmed()) + await _page(monkeypatch, on, store) + stored = await window_module.earlier_runs(h.THREAD, on.tip, await on.trimmed()) + assert extracted is not None and stored is not None + assert _text(stored) == _text(extracted) + + whole_off = (await off.run.graph.aget_state(h.config())).values["messages"] + trimmed_off = [m for run in split_runs(whole_off)[: on.base] for m in run] + assert _text(stored) == _text(trimmed_off) + window = (await on.run.graph.aget_state(h.config())).values["messages"] + assert ( + build_directory([*stored, *window], window=Window()).manifest + == build_directory(h.whole_threads(on.run)[-1], window=Window()).manifest + ) + + +def _job(on: Thread, window, previous: str | None = None) -> transcripts._Job: + return transcripts._Job( + TranscriptTarget(h.THREAD), + list(window), + {}, + "fp", + on.tip, + previous, + on.head(on.base), + ) + + +async def test_a_render_with_no_copy_reads_the_trimmed_runs( + monkeypatch, threads +) -> None: + _, on = threads + _cache_probe(monkeypatch) + store = SliceStore() + await _page(monkeypatch, on, store) + window = (await on.run.graph.aget_state(h.config())).values["messages"] + rendered = await transcripts._render_job(_job(on, window), inline=True) + whole = build_directory(h.whole_threads(on.run)[-1], header={}, window=Window()) + assert rendered.copy.manifest == whole.manifest + + +async def test_a_copy_from_another_branch_is_rendered_over(monkeypatch, threads) -> None: + """A truncation keeps a live copy labelled with the fork point, which can + hold the other branch's turns under the same numbers: its entries for + the trimmed runs are not carried, and the runs are read back instead.""" + _, on = threads + _cache_probe(monkeypatch) + store = SliceStore() + await _page(monkeypatch, on, store) + whole = h.whole_threads(on.run)[-1] + other = [ + m.model_copy(update={"content": "Question 1, asked on another branch"}) + if m.id == "h-1" + else m + for m in whole + ] + copy = build_directory(other, header={}, window=Window()) + window = (await on.run.graph.aget_state(h.config())).values["messages"] + rendered = await transcripts._render_job(_job(on, window, copy.manifest), inline=True) + assert rendered.copy.manifest == build_directory(whole, header={}, window=Window()).manifest + + +async def test_a_render_that_cannot_read_them_fails(monkeypatch, threads) -> None: + _, on = threads + _serve(monkeypatch, on, SliceStore()) + + async def unreadable(thread_id, tip, head): + return None + + monkeypatch.setattr(window_module, "earlier_runs", unreadable) + window = (await on.run.graph.aget_state(h.config())).values["messages"] + with pytest.raises(EarlierTurnsMissing): + await transcripts._render_job(_job(on, window), inline=True) + + +async def test_the_read_back_must_be_the_trimmed_messages(monkeypatch, threads) -> None: + _, on = threads + _serve(monkeypatch, on, SliceStore()) + reader = CheckpointHistoryReader.get_instance() + real_slices = reader.aget_turn_slices + + async def merged(thread_id, anchors): + turns = await real_slices(thread_id, anchors) + for turn in turns[1:]: + turn.messages[:] = [m for m in turn.messages if m.type != "human"] + return turns + + monkeypatch.setattr(reader, "aget_turn_slices", merged) + assert await window_module.earlier_runs(h.THREAD, on.tip, on.head(on.base)) is None + + +# ------------------------------------------------------------ thread shapes + +SHAPED = "5e1f0000-0000-4000-8000-0000000005a9" + + +async def _shaped(monkeypatch, requests: list[tuple[list, list | None]]): + """A thread of one turn per request, ``(input, reply)``, a reply of None + being a turn killed before it ran, with every turn's slice stored as + replay stores it. Returns the messages at the tip.""" + saver = InMemorySaver() + replies = [reply for _, reply in requests] + + def answer(state): + reply = replies.pop(0) + return {"messages": reply} if reply else {} + + graph = ( + StateGraph(DeltaAgentState) + .add_node("answer", answer) + .add_edge(START, "answer") + .compile(checkpointer=saver) + ) + with tracing_context(enabled=False): + for i, (given, _) in enumerate(requests): + config = { + "configurable": {"thread_id": SHAPED}, + "metadata": {"turn_index": i, "run_id": f"run-{i}"}, + } + await graph.ainvoke({"messages": given}, config) + state = await graph.aget_state({"configurable": {"thread_id": SHAPED}}) + tip = state.config["configurable"]["checkpoint_id"] + + reader = CheckpointHistoryReader(saver) + monkeypatch.setattr(CheckpointHistoryReader, "get_instance", classmethod(lambda cls: reader)) + anchors, _ = await reader.aget_turn_anchors(SHAPED, tip) + assert len(anchors) == len(requests) + rows = [ + slices_db.StoredSlice( + a.input_checkpoint_id, + a.tail_checkpoint_id, + slices.SLICE_KEY, + *slices.encode(reader.serde, turn), + ) + for a, turn in zip(anchors, await reader.aget_turn_slices(SHAPED, anchors)) + ] + + async def held(thread_id, key): + return bool(rows) + + async def at(thread_id, key, inputs): + return [row for row in rows if row.input_checkpoint_id in inputs] + + async def pointer(thread_id): + return tip + + monkeypatch.setattr(slices_db, "has_turn_slices", held) + monkeypatch.setattr(slices_db, "get_slices_at", at) + monkeypatch.setattr(window_module, "get_thread_checkpoint_id", pointer) + return tip, list(state.values["messages"]) + + +async def test_turns_that_open_no_run_or_several(monkeypatch) -> None: + """Counted by turns, each of these misplaces the runs: a request that + opens with an assistant message (its message belongs to the run + before), one killed before it ran, and two questions in one request.""" + tip, messages = await _shaped( + monkeypatch, + [ + ([AIMessage("Welcome.", id="a-0"), HumanMessage("Q1", id="h-1")], [AIMessage("A1", id="a-1")]), + ([], None), + ([HumanMessage("Q2", id="h-2"), HumanMessage("Q3", id="h-3")], [AIMessage("A3", id="a-3")]), + ([AIMessage("Also:", id="a-3b"), HumanMessage("Q4", id="h-4")], [AIMessage("A4", id="a-4")]), + ([HumanMessage("Q5", id="h-5")], [AIMessage("A5", id="a-5")]), + ], + ) + runs = split_runs(messages) + assert [[m.id for m in run] for run in runs[:4]] == [ + ["a-0", "h-1", "a-1"], + ["h-2"], + ["h-3", "a-3", "a-3b"], + ["h-4", "a-4"], + ] + for n in range(1, 5): + trimmed = [m for run in runs[:n] for m in run] + head = Window().extend(trimmed) + assert head.runs == n + assert await window_module.runs_held(SHAPED, head), n + read_back = await window_module.earlier_runs(SHAPED, tip, head) + assert [m.id for m in read_back or ()] == [m.id for m in trimmed], n + + +async def test_a_thread_the_old_summarizer_compacted_is_not_trimmed(monkeypatch) -> None: + """Its summary replaced the messages before it in the checkpoint, while + the slices still hold them: no trim may leave the summary's turn to + them.""" + summary = HumanMessage("Summary of the conversation to date: Q1 and Q2.", id="s-1") + tip, messages = await _shaped( + monkeypatch, + [ + ([HumanMessage("Q1", id="h-1")], [AIMessage("A1", id="a-1")]), + ([HumanMessage("Q2", id="h-2")], [AIMessage("A2", id="a-2")]), + ( + [HumanMessage("Q3", id="h-3")], + [ + RemoveMessage(id=REMOVE_ALL_MESSAGES), + summary, + AIMessage("A2", id="a-2"), + HumanMessage("Q3", id="h-3"), + AIMessage("A3", id="a-3"), + ], + ), + ([HumanMessage("Q4", id="h-4")], [AIMessage("A4", id="a-4")]), + ], + ) + assert [m.id for m in messages] == ["s-1", "a-2", "h-3", "a-3", "h-4", "a-4"] + head = Window().extend(messages[:4]) + assert not await window_module.runs_held(SHAPED, head) + assert await window_module.earlier_runs(SHAPED, tip, head) is None + + +# ------------------------------------------------------------ manual writes + + +async def test_a_manual_write_moves_the_commit_pointer(monkeypatch) -> None: + from src.server.database.conversation import threads_write + from src.server.handlers.thread_maintenance import _update_graph_state + + graph = ( + StateGraph(MessagesState) + .add_node("echo", lambda state: {"messages": [AIMessage("ok")]}) + .add_edge(START, "echo") + .compile(checkpointer=InMemorySaver()) + ) + config = {"configurable": {"thread_id": "t"}} + with tracing_context(enabled=False): + await graph.ainvoke({"messages": [HumanMessage("q")]}, config) + state = await graph.aget_state(config) + moved = [] + + async def advance(thread_id, *, from_checkpoint_id, to_checkpoint_id): + moved.append((thread_id, from_checkpoint_id, to_checkpoint_id)) + return False + + monkeypatch.setattr(threads_write, "advance_thread_checkpoint_id", advance) + await _update_graph_state(graph, state, {"messages": [AIMessage("x")]}, "t", "compact") + written = await graph.aget_state(config) + built_on = state.config["configurable"]["checkpoint_id"] + assert written.parent_config["configurable"]["checkpoint_id"] == built_on + assert moved == [("t", built_on, written.config["configurable"]["checkpoint_id"])] diff --git a/tests/unit/server/services/runs/test_coordinator.py b/tests/unit/server/services/runs/test_coordinator.py index 9e1d15141..ce402a060 100644 --- a/tests/unit/server/services/runs/test_coordinator.py +++ b/tests/unit/server/services/runs/test_coordinator.py @@ -35,7 +35,7 @@ def _row(status: str, msg_type: str = "ptc") -> dict: def _coordinator() -> RunCoordinator: coordinator = RunCoordinator() - coordinator._schedule_projection_refresh = MagicMock() + coordinator._schedule_replay_refresh = MagicMock() coordinator._nudge_hook_drainer = MagicMock() coordinator._latest_checkpoint_id = AsyncMock(return_value="cp-1") return coordinator @@ -129,5 +129,5 @@ async def test_an_export_that_fails_to_schedule_only_logs(): with patch(EXPORT, side_effect=RuntimeError("no loop")): coordinator.post_finalize_tail("thread-1", _row("completed")) - coordinator._schedule_projection_refresh.assert_called_once_with("thread-1") + coordinator._schedule_replay_refresh.assert_called_once_with("thread-1") coordinator._nudge_hook_drainer.assert_called_once() diff --git a/tests/unit/server/services/runs/test_finalize_credit_pause_stamp.py b/tests/unit/server/services/runs/test_finalize_credit_pause_stamp.py new file mode 100644 index 000000000..11bef8808 --- /dev/null +++ b/tests/unit/server/services/runs/test_finalize_credit_pause_stamp.py @@ -0,0 +1,415 @@ +"""What the finalize records at the source so its readers never reopen the +turn's event archive. + +A credit pause's denial is stamped into the run row's metadata, which is what +settling the run's automation relays; rows a release before the stamp +finalized still answer from the pause's stored interrupt. The producer's record +(its reasoning durations and provenance) rides the outcome, and a record that +cannot be read costs the row its facts, never its finalize. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +from langgraph.types import Interrupt + +from src.server.contracts.status import ( + INTERRUPT_REASON_CREDIT_PAUSE, + classify_interrupts, +) +from src.server.database.runs.lifecycle import ProducerRecord +from src.server.services.automation_settlement import run_failure_message +from src.server.services.runs.finalization import ( + assemble_finalize_artifacts, + classify_outcome, +) +from src.server.services.runs.recovery import RecoveryScanner +from src.server.services.runs.sse_producer import RunSSEProducer + +_DENIAL = "You have used this month's credits." +_KEY = ("t-pause", "r-pause") + + +def _credit_pause(message=_DENIAL): + request = {"type": INTERRUPT_REASON_CREDIT_PAUSE} + if message is not None: + request["message"] = message + return request + + +def _question(message="Which filing?"): + return {"type": "ask_user_question", "message": message} + + +def _payload(*requests): + """One interrupt payload, as the graph raises it.""" + return Interrupt(value={"action_requests": list(requests)}, id="intr-1") + + +class _ScriptedGraph: + def __init__(self, script): + self._script = script + + def astream( + self, + _input_state, + config=None, + stream_mode=None, + subgraphs=None, + durability=None, + ): + async def _gen(): + for step in self._script: + yield step + + return _gen() + + +async def _streamed(*script, durable=True): + """A producer that ran the real stream loop over ``script``; a pause in + it passes the durability barrier when ``durable``.""" + handler = RunSSEProducer(thread_id=_KEY[0], run_id=_KEY[1]) + + async def _verify(_graph): + return durable + + handler._verify_interrupt_durable = _verify + async for _ in handler.stream_workflow( + _ScriptedGraph(script), input_state={}, config={} + ): + pass + return handler + + +def _paused_on(*requests): + return ((), "updates", {"__interrupt__": (_payload(*requests),)}) + + +async def _outcome(handler, kind="stream_end", *, graph=None, stop_events=None): + """The outcome the executor's finalize would write for ``handler``'s run + ending as ``kind``.""" + status, phase, interrupt_reason, pause_message, error = await classify_outcome( + kind, handler=handler, graph=graph, thread_id=_KEY[0], error=None + ) + return await assemble_finalize_artifacts( + _KEY, + metadata={"handler": handler}, + status=status, + phase=phase, + interrupt_reason=interrupt_reason, + pause_message=pause_message, + error=error, + handler=handler, + cancelled_by_user=False, + workspace_id=None, + user_id=None, + stop_events=stop_events, + ) + + +async def _persist_metadata(handler, kind="stream_end", *, graph=None): + return (await _outcome(handler, kind, graph=graph)).metadata + + +# ─── The finalize stamps the pause's words ──────────────────────────── + + +@pytest.mark.asyncio +async def test_a_credit_pause_stamps_its_denial(): + handler = await _streamed(_paused_on(_credit_pause())) + assert handler.interrupt_reason == INTERRUPT_REASON_CREDIT_PAUSE + + persist_metadata = await _persist_metadata(handler) + + assert persist_metadata["credit_pause_message"] == _DENIAL + + +@pytest.mark.asyncio +async def test_a_credit_pause_without_words_is_stamped_null(): + # Settling then knows there is nothing to relay without the archive. + handler = await _streamed(_paused_on(_credit_pause(message=None))) + + persist_metadata = await _persist_metadata(handler) + + assert persist_metadata["credit_pause_message"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("requests", "kind", "durable"), + [ + ((_question(),), "stream_end", True), + ((_credit_pause(),), "cancelled", True), + ((_credit_pause(),), "stream_end", False), + ], + ids=["paused_on_a_question", "settled_cancelled", "pause_not_durable"], +) +async def test_nothing_is_stamped_unless_a_credit_pause_settles( + requests, kind, durable +): + handler = await _streamed(_paused_on(*requests), durable=durable) + + persist_metadata = await _persist_metadata(handler, kind) + + assert "credit_pause_message" not in persist_metadata + + +class _PausedGraph: + """A graph whose checkpoint is paused on ``requests``, for the + handler-less finalize's state probe.""" + + def __init__(self, *requests): + task = SimpleNamespace(interrupts=(_payload(*requests),)) + self._snapshot = SimpleNamespace(next=("tools",), tasks=(task,)) + + async def aget_state(self, _config): + return self._snapshot + + +@pytest.mark.asyncio +async def test_a_handler_less_credit_pause_stamps_its_denial(): + # The state probe classifies the pause from the same interrupts the + # producer would have buffered, so it stamps the same words. + metadata = await _persist_metadata(None, graph=_PausedGraph(_credit_pause())) + assert metadata["credit_pause_message"] == _DENIAL + + metadata = await _persist_metadata(None, graph=_PausedGraph(_question())) + assert "credit_pause_message" not in metadata + assert "credit_pause_message" not in await _persist_metadata(None) + + +# ─── Recovery stamps a pause it settles ─────────────────────────────── + + +async def _recovered_outcome(*requests, salvaged=None): + """The outcome the scanner finalizes a dead run with, its tip checkpoint + pending ``requests`` as one interrupt and its stream salvaged as + ``salvaged``.""" + tip = SimpleNamespace( + metadata={"run_id": _KEY[1]}, + config={"configurable": {"checkpoint_id": "cp-tip"}}, + pending_writes=[("task-1", "__interrupt__", [_payload(*requests)])], + ) + saver = SimpleNamespace(aget_tuple=AsyncMock(return_value=tip)) + finalize = AsyncMock(return_value=SimpleNamespace(applied=True, run={})) + coordinator = SimpleNamespace(finalize_detached_run=finalize) + scanner = RecoveryScanner() + scanner._salvage_stream = AsyncMock(return_value=(salvaged, "empty_stream")) + with ( + patch("src.server.app.setup.checkpointer", saver), + patch( + "src.server.services.runs.coordinator.RunCoordinator.get_instance", + return_value=coordinator, + ), + ): + await scanner._recover_run({}, _KEY[1], _KEY[0]) + outcome = finalize.await_args.args[2] + assert outcome.status == "interrupted" + return outcome + + +async def _recovered_metadata(*requests): + return (await _recovered_outcome(*requests)).metadata + + +@pytest.mark.asyncio +async def test_a_recovered_credit_pause_stamps_its_denial(): + # The scanner's salvaged archive is not where settling reads it. + metadata = await _recovered_metadata(_credit_pause()) + assert metadata["credit_pause_message"] == _DENIAL + assert metadata["recovery"] == "scanner" + + +@pytest.mark.asyncio +async def test_a_recovered_pause_is_stamped_only_when_it_is_a_credit_pause(): + assert (await _recovered_metadata(_credit_pause(message=None)))[ + "credit_pause_message" + ] is None + assert "credit_pause_message" not in await _recovered_metadata(_question()) + + +@pytest.mark.asyncio +async def test_a_recovered_run_records_its_salvaged_provenance_but_no_durations(): + # Nothing timed the dead run's thinking, so its row keeps no reasoning + # facts and replay reads the durations off the salvaged events instead. + found = { + "event": "provenance", + "data": {"agent": "main", "source_type": "web", "identifier": "https://a.test"}, + } + chunk = {"event": "message_chunk", "data": {"agent": "main", "content": "x"}} + + outcome = await _recovered_outcome(_question(), salvaged=[chunk, found]) + + assert outcome.record == ProducerRecord(provenance=[found], reasoning_ms=None) + assert (await _recovered_outcome(_question())).record == ProducerRecord() + + +# ─── classify_interrupts: the pause's words ────────────────────────── + + +def _words(interrupts): + return classify_interrupts(interrupts)[1] + + +def test_the_last_credit_pause_with_words_wins(): + assert ( + _words( + [_payload(_credit_pause("First denial.")), _payload(_credit_pause(_DENIAL))] + ) + == _DENIAL + ) + # A later pause without words does not erase the earlier one's, as the + # stored-event reader never did either. + assert ( + _words( + [_payload(_credit_pause(_DENIAL)), _payload(_credit_pause(message=None))] + ) + == _DENIAL + ) + + +def test_another_questions_words_are_never_relayed(): + assert classify_interrupts( + [_payload(_credit_pause(_DENIAL)), _payload(_question("Proceed anyway?"))] + ) == (INTERRUPT_REASON_CREDIT_PAUSE, _DENIAL) + assert classify_interrupts([_payload(_question("Proceed anyway?"))]) == ( + "user_question", + None, + ) + + +@pytest.mark.parametrize( + "interrupts", + [ + [], + [Interrupt(value="not a dict")], + [ + _payload( + {"type": INTERRUPT_REASON_CREDIT_PAUSE, "message": {"text": _DENIAL}} + ) + ], + ], + ids=["none", "not_a_payload", "words_not_a_string"], +) +def test_no_credit_pause_words_is_none(interrupts): + assert _words(interrupts) is None + + +# ─── run_failure_message reads the stamp first ──────────────────────── + + +def _paused_row(metadata, *events): + return { + "status": "interrupted", + "interrupt_reason": INTERRUPT_REASON_CREDIT_PAUSE, + "conversation_response_id": _KEY[1], + "errors": None, + "metadata": metadata, + "sse_events": list(events) if events else None, + } + + +def _pause_event(message): + return {"event": "interrupt", "data": {"action_requests": [_credit_pause(message)]}} + + +def test_the_stamp_is_preferred_over_the_stored_events(): + row = _paused_row({"credit_pause_message": _DENIAL}, _pause_event("Older words.")) + assert run_failure_message(row) == _DENIAL + + +def test_the_stamp_needs_no_stored_events(): + # The settlement read leaves the archive out once the stamp is there. + assert ( + run_failure_message(_paused_row({"credit_pause_message": _DENIAL})) == _DENIAL + ) + + +@pytest.mark.parametrize( + "metadata", [{}, None, {"deepthinking": False}], ids=["empty", "none", "unstamped"] +) +def test_an_older_row_still_answers_from_its_stored_pause(metadata): + assert run_failure_message(_paused_row(metadata, _pause_event(_DENIAL))) == _DENIAL + + +def test_the_stamp_is_read_only_for_a_credit_pause(): + row = { + "status": "error", + "interrupt_reason": None, + "conversation_response_id": _KEY[1], + "errors": ["upstream refused"], + "metadata": {"credit_pause_message": _DENIAL}, + } + assert run_failure_message(row) == "upstream refused" + + +# ─── The producer's record ──────────────────────────────────────────── + + +def _provenance(agent, identifier): + return { + "event": "provenance", + "data": {"source_type": "web_search", "identifier": identifier, "agent": agent}, + } + + +@pytest.mark.asyncio +async def test_the_outcome_carries_the_record_with_the_stop_drains_provenance(): + handler = await _streamed( + ( + (), + "custom", + { + "type": "provenance", + "source_type": "web_search", + "identifier": "https://example.test/a", + "agent": None, + }, + ) + ) + main = handler.provenance_events + assert [e["data"]["agent"] for e in main] == ["main"] + drained = _provenance("task:k7", "https://example.test/b") + + outcome = await _outcome( + handler, + "cancelled", + stop_events=[{"event": "message_chunk", "data": {}}, drained, "not an event"], + ) + + assert outcome.record == ProducerRecord( + provenance=[*main, drained], reasoning_ms={} + ) + + +@pytest.mark.asyncio +async def test_a_main_lane_with_no_provenance_records_an_empty_list(): + """A producer always timed the turn, even when nothing was measured, so + its facts are written and replay never reads them off the events.""" + handler = await _streamed() + drained = _provenance("task:k7", "https://example.test/b") + + assert handler.record() == ProducerRecord(provenance=[], reasoning_ms={}) + assert handler.record([drained]).provenance == [drained] + + +@pytest.mark.asyncio +async def test_a_record_that_cannot_be_read_never_blocks_the_finalize(): + # The finalize must still reach its CAS: a raise here would leave the + # row in_progress for recovery. + handler = await _streamed() + + def _raises(_stop_events=None): + raise RuntimeError("record unavailable") + + handler.record = _raises + outcome = await _outcome(handler) + + assert outcome.status == "completed" + assert outcome.record is None + + +@pytest.mark.asyncio +async def test_a_run_without_a_producer_has_no_record(): + assert (await _outcome(None)).record is None diff --git a/tests/unit/server/services/runs/test_recovery_task_chains.py b/tests/unit/server/services/runs/test_recovery_task_chains.py index 7d6919ec4..19a96df67 100644 --- a/tests/unit/server/services/runs/test_recovery_task_chains.py +++ b/tests/unit/server/services/runs/test_recovery_task_chains.py @@ -15,6 +15,7 @@ from src.server.services.runs.recovery import RecoveryScanner SR_DB = "src.server.services.runs.recovery.sr_db" +REPAIR = "src.server.services.runs.recovery.subagent_repair" TL_DB = "src.server.services.runs.recovery.tl_db" @@ -44,7 +45,7 @@ async def test_sweep_runs_even_when_there_is_nothing_to_recover(): async def test_heal_reports_what_it_repaired(): repair = AsyncMock(return_value={"rewound": 2, "deleted": 1}) - with patch(f"{SR_DB}.repair_dangling_task_chains", new=repair): + with patch(f"{REPAIR}.repair_dangling_task_chains", new=repair): healed = await RecoveryScanner().heal_task_chains() assert healed == {"rewound": 2, "deleted": 1} @@ -57,7 +58,7 @@ async def test_heal_failure_never_aborts_the_scan(): concerns are independent and only one of them is time-critical.""" repair = AsyncMock(side_effect=RuntimeError("connection reset")) - with patch(f"{SR_DB}.repair_dangling_task_chains", new=repair): + with patch(f"{REPAIR}.repair_dangling_task_chains", new=repair): healed = await RecoveryScanner().heal_task_chains() assert healed == {"rewound": 0, "deleted": 0} @@ -69,7 +70,7 @@ async def test_scan_survives_a_failing_sweep(): repair = AsyncMock(side_effect=RuntimeError("connection reset")) with open_runs, open_task_runs, patch( - f"{SR_DB}.repair_dangling_task_chains", new=repair + f"{REPAIR}.repair_dangling_task_chains", new=repair ): assert await RecoveryScanner().scan_once(assume_dead=True) == 0 diff --git a/tests/unit/server/services/runs/test_sse_producer_reasoning_facts.py b/tests/unit/server/services/runs/test_sse_producer_reasoning_facts.py new file mode 100644 index 000000000..dced616b9 --- /dev/null +++ b/tests/unit/server/services/runs/test_sse_producer_reasoning_facts.py @@ -0,0 +1,280 @@ +"""How long the main agent thought, kept for replay. + +The checkpoint keeps a turn's reasoning but not its duration, so the producer +records each closed reasoning block's ``elapsed_ms`` and the finalize stores +them on the run's row. A model call streams under a chunk id the checkpoint +never sees, so a close is attributed only when the node's update names the +message it committed, and a call that committed nothing takes its closes with +it. Each test drives the real ``stream_workflow`` loop over a scripted graph. +""" + +import json +import time +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage +from langgraph.types import Overwrite + +from src.server.services.runs import sse_producer +from src.server.services.runs.sse_producer import RunSSEProducer + +_RUNNER = "src.server.services.thread_mutation.ThreadMutationRunner.get_instance" +_SUMMARIZE_NODE = "SummarizationMiddleware.before_model" + + +class _Clock: + def __init__(self): + self.now = 100.0 + + def monotonic(self): + return self.now + + def advance(self, seconds): + return lambda: setattr(self, "now", self.now + seconds) + + +class _ScriptedGraph: + """Yields each ``(namespace, stream_mode, payload)`` tuple of the script + and runs each callable in place, so a test can move the clock or look at + the producer between two events.""" + + def __init__(self, script): + self._script = script + + def astream( + self, + _input_state, + config=None, + stream_mode=None, + subgraphs=None, + durability=None, + ): + async def _gen(): + for step in self._script: + if callable(step): + step() + else: + yield step + + return _gen() + + +@pytest.fixture +def clock(monkeypatch): + # Swap the module's ``time`` rather than ``time.monotonic`` itself, which + # the event loop reads too. + fake = _Clock() + monkeypatch.setattr( + sse_producer, "time", SimpleNamespace(monotonic=fake.monotonic, time=time.time) + ) + return fake + + +def _producer(): + return RunSSEProducer(thread_id="t-facts", run_id="r-facts") + + +async def _run(handler, script): + return [ + event + async for event in handler.stream_workflow( + _ScriptedGraph(script), input_state={}, config={} + ) + ] + + +def _think(chunk_id, *, ns=(), node="model"): + chunk = AIMessageChunk( + content="", + id=chunk_id, + additional_kwargs={"reasoning_content": "Weighing the filing."}, + ) + return (ns, "messages", (chunk, {"langgraph_node": node})) + + +def _say(chunk_id, *, ns=(), node="model"): + chunk = AIMessageChunk(content="The answer.", id=chunk_id) + return (ns, "messages", (chunk, {"langgraph_node": node})) + + +def _update(node, messages, *, ns=()): + return (ns, "updates", {node: {"messages": messages}}) + + +def _committed(message_id): + return AIMessage(content="The answer.", id=message_id) + + +def _reasoning_closes(events): + """The ``complete`` reasoning signals the stream emitted, as payloads.""" + closes = [] + for event in events: + data = json.loads(event.split("data: ", 1)[1]) + if ( + data.get("content_type") == "reasoning_signal" + and data.get("content") == "complete" + ): + closes.append((event.split("\n")[1], data)) + return closes + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "messages", + [ + [HumanMessage(content="What changed?", id="h-1"), _committed("msg_123")], + Overwrite( + [ + _committed("msg_old"), + HumanMessage(content="And now?"), + _committed("msg_123"), + ] + ), + _committed("msg_123"), + ], + ids=["list", "overwrite", "single"], +) +async def test_a_close_lands_on_the_message_the_node_committed(clock, messages): + handler = _producer() + await _run( + handler, + [ + _think("lc_run--abc"), + clock.advance(0.25), + _say("lc_run--abc"), + _update("model", messages), + ], + ) + assert handler.record().reasoning_ms == {"msg_123": [250]} + + +@pytest.mark.asyncio +async def test_two_thinking_blocks_in_one_call_land_in_order(clock): + handler = _producer() + await _run( + handler, + [ + _think("lc_run--abc"), + clock.advance(0.25), + _say("lc_run--abc"), + _think("lc_run--abc"), + clock.advance(0.5), + _say("lc_run--abc"), + _update("model", [_committed("msg_123")]), + ], + ) + assert handler.record().reasoning_ms == {"msg_123": [250, 500]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retry", "expected"), + [ + ( + lambda clock: [ + _think("lc_run--second"), + clock.advance(0.5), + _say("lc_run--second"), + ], + {"msg_2": [500]}, + ), + (lambda clock: [_say("lc_run--second")], {}), + ], + ids=["retry_thought", "retry_did_not_think"], +) +async def test_a_replaced_attempt_takes_its_closes_with_it(clock, retry, expected): + """A retry inside the node streams under a new chunk id; only the newest + call's thinking belongs to the message the node commits.""" + handler = _producer() + await _run( + handler, + [ + _think("lc_run--first"), + clock.advance(0.25), + _say("lc_run--first"), + *retry(clock), + _update("model", [_committed("msg_2")]), + ], + ) + assert handler.record().reasoning_ms == expected + + +@pytest.mark.asyncio +async def test_another_nodes_update_leaves_the_closes_waiting(clock): + handler = _producer() + seen = [] + await _run( + handler, + [ + _think("lc_run--abc"), + clock.advance(0.25), + _say("lc_run--abc"), + _update( + "tools", [ToolMessage(content="42", tool_call_id="call-1", id="tool-1")] + ), + lambda: seen.append(handler.record().reasoning_ms), + _update("model", [_committed("msg_123")]), + ], + ) + assert seen == [{}] + assert handler.record().reasoning_ms == {"msg_123": [250]} + + +@pytest.mark.asyncio +async def test_a_subagents_thinking_records_nothing(clock): + ns = ("task:k7",) + handler = _producer() + events = await _run( + handler, + [ + _think("lc_run--sub", ns=ns), + clock.advance(0.25), + _say("lc_run--sub", ns=ns), + _update("model", [_committed("msg_sub")], ns=ns), + _say("lc_run--main"), + _update("model", [_committed("msg_main")]), + ], + ) + # The block still closed with its duration on the wire; it just is not + # the main reply's thinking, not even the reply the main agent commits + # after it. + assert [ + (line, data["agent"], data["elapsed_ms"]) + for line, data in _reasoning_closes(events) + ] == [("event: message_chunk", "task:k7", 250)] + assert handler.record().reasoning_ms == {} + + +@pytest.mark.asyncio +async def test_a_compactions_thinking_records_nothing(clock): + summarize = {"type": "context_window", "action": "summarize"} + handler = _producer() + with patch(_RUNNER, return_value=MagicMock()): + events = await _run( + handler, + [ + ((), "custom", {**summarize, "signal": "start"}), + _think("lc_run--sum", node=_SUMMARIZE_NODE), + clock.advance(0.25), + _say("lc_run--sum", node=_SUMMARIZE_NODE), + ((), "custom", {**summarize, "signal": "complete"}), + _update(_SUMMARIZE_NODE, [_committed("msg_sum")]), + ], + ) + assert [(line, data["elapsed_ms"]) for line, data in _reasoning_closes(events)] == [ + ("event: compaction_chunk", 250) + ] + assert handler.record().reasoning_ms == {} + + +@pytest.mark.asyncio +async def test_nothing_measured_is_an_empty_record(clock): + assert _producer().record().reasoning_ms == {} + + handler = _producer() + await _run( + handler, [_say("lc_run--abc"), _update("model", [_committed("msg_123")])] + ) + assert handler.record().reasoning_ms == {} diff --git a/tests/unit/server/services/runs/test_subagent_archive_run_facts.py b/tests/unit/server/services/runs/test_subagent_archive_run_facts.py new file mode 100644 index 000000000..97d00aed4 --- /dev/null +++ b/tests/unit/server/services/runs/test_subagent_archive_run_facts.py @@ -0,0 +1,130 @@ +"""A subagent archive records its lanes' captured images as run facts. + +The lanes settle after their turn did, so a checkpoint ui record appended now +would land on whatever turn is newest. Each lane's path-to-URL map goes onto +the runs the turn launched for that task instead, limited to the paths that +lane referenced, and the launching turn is re-projected once a run took it. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from src.server.services import subagent_run_coordinator +from src.server.services.runs.subagent_archive import persist_collected_events + +_A = "/home/workspace/a.png" +_B = "/home/workspace/b.png" +_URL_A = "https://cdn.example.com/a.png" +_URL_B = "https://cdn.example.com/b.png" + + +def _text(agent: str | None, content: str) -> dict: + return { + "event": "message_chunk", + "data": {"agent": agent, "content_type": "text", "content": content}, + } + + +async def _persist(events: list[dict], captured: dict[str, str]): + record = AsyncMock() + reader = MagicMock() + reader.append_ui_record = AsyncMock() + with ( + patch( + "src.server.services.persistence.image_capture.capture_images", + AsyncMock(return_value=captured), + ), + patch.object(subagent_run_coordinator, "record_lane_images", record), + patch( + "src.server.services.runs.subagent_archive.replace_agent_events", + AsyncMock(return_value=True), + ), + patch( + "src.server.services.history.reader.CheckpointHistoryReader.get_instance", + return_value=reader, + ), + ): + ok = await persist_collected_events( + events, "resp-1", "thread-1", "ws-1", sandbox=object() + ) + return ok, record, reader + + +@pytest.mark.asyncio +async def test_each_lane_records_only_the_images_it_referenced(): + events = [_text("task:k1", f"see ![a]({_A})"), _text("task:k2", f"see ![b]({_B})")] + + ok, record, reader = await _persist(events, {_A: _URL_A, _B: _URL_B}) + + assert ok is True + assert sorted(c.args for c in record.await_args_list) == [ + ("thread-1", "resp-1", "k1", {_A: _URL_A}), + ("thread-1", "resp-1", "k2", {_B: _URL_B}), + ] + reader.append_ui_record.assert_not_called() + + +@pytest.mark.asyncio +async def test_a_lane_with_nothing_captured_records_nothing(): + events = [_text("task:k1", f"see ![a]({_A})"), _text("task:k2", "no image")] + + _, record, _ = await _persist(events, {}) + + record.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_rows_outside_a_task_lane_record_nothing(): + _, record, _ = await _persist([_text(None, f"see ![a]({_A})")], {_A: _URL_A}) + + record.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("landed", "refreshed"), [(True, True), (False, False)] +) +async def test_the_launching_turn_is_reprojected_once_a_run_took_the_images( + landed, refreshed +): + append = AsyncMock(return_value=landed) + refresh = MagicMock() + with ( + patch.object(subagent_run_coordinator.facts_db, "append_lane_images", append), + patch( + "src.server.services.history.replay.refresh.schedule_replay_refresh", + refresh, + ), + ): + await subagent_run_coordinator.record_lane_images( + "thread-1", "resp-1", "k1", {_A: _URL_A} + ) + + append.assert_awaited_once_with("thread-1", "resp-1", "k1", {_A: _URL_A}) + assert refresh.called is refreshed + if refreshed: + refresh.assert_called_once_with("thread-1", "resp-1") + + +@pytest.mark.asyncio +async def test_images_that_fail_to_land_never_fail_the_archive(): + refresh = MagicMock() + with ( + patch.object( + subagent_run_coordinator.facts_db, + "append_lane_images", + AsyncMock(side_effect=RuntimeError("pool closed")), + ), + patch( + "src.server.services.history.replay.refresh.schedule_replay_refresh", + refresh, + ), + ): + await subagent_run_coordinator.record_lane_images( + "thread-1", "resp-1", "k1", {_A: _URL_A} + ) + + refresh.assert_not_called() diff --git a/tests/unit/server/services/test_automation_excerpt.py b/tests/unit/server/services/test_automation_excerpt.py index b27ad5cc2..9cf8b91f1 100644 --- a/tests/unit/server/services/test_automation_excerpt.py +++ b/tests/unit/server/services/test_automation_excerpt.py @@ -28,3 +28,47 @@ def test_a_long_answer_is_cut_at_a_word(): assert excerpt.endswith("…") assert len(excerpt) <= 321 assert not excerpt[:-1].endswith(" ") + + +async def _two_runs(monkeypatch): + from langchain_core.messages import AIMessage, HumanMessage + from langgraph.checkpoint.memory import InMemorySaver + from langgraph.graph import START, StateGraph + + from ptc_agent.agent.state import DeltaAgentState + from src.server.services.history.reader import CheckpointHistoryReader + + saver = InMemorySaver() + graph = ( + StateGraph(DeltaAgentState) + .add_node( + "agent", + lambda s: {"messages": [AIMessage(f"**Answer** {len(s['messages'])}")]}, + ) + .add_edge(START, "agent") + .compile(checkpointer=saver) + ) + for n in range(2): + await graph.ainvoke( + {"messages": [HumanMessage(f"q{n}")]}, + {"configurable": {"thread_id": "t"}, "metadata": {"run_id": f"run-{n}"}}, + ) + monkeypatch.setattr( + CheckpointHistoryReader, "get_instance", lambda: CheckpointHistoryReader(saver) + ) + + +@pytest.mark.asyncio +async def test_the_newest_run_reads_its_answer(monkeypatch): + from src.server.services.automation_excerpt import read_run_excerpt + + await _two_runs(monkeypatch) + assert await read_run_excerpt("t", "run-1") == "Answer 3" + + +@pytest.mark.asyncio +async def test_a_run_a_newer_turn_followed_leaves_none(monkeypatch): + from src.server.services.automation_excerpt import read_run_excerpt + + await _two_runs(monkeypatch) + assert await read_run_excerpt("t", "run-0") is None diff --git a/tests/unit/server/services/test_automation_settlement.py b/tests/unit/server/services/test_automation_settlement.py index d8c543827..248d49a55 100644 --- a/tests/unit/server/services/test_automation_settlement.py +++ b/tests/unit/server/services/test_automation_settlement.py @@ -370,13 +370,13 @@ def test_what_a_failed_model_call_records_decides_the_outcome( ): """The ledger metadata the run writes is what settlement reads, so the two halves are pinned together rather than each against a copy.""" - from src.server.services.runs import sse_producer + from src.server.services.runs import stream_errors - monkeypatch.setattr(sse_producer.app_settings, "HOST_MODE", host_mode) + monkeypatch.setattr(stream_errors.app_settings, "HOST_MODE", host_mode) exc = RuntimeError("refused") exc.__model_resilience__ = {"attempted_models": [{"status_code": status}]} - metadata = sse_producer.model_call_failure(exc, source) + metadata = stream_errors.model_call_failure(exc, source) assert ledger_outcome({"status": "error", "metadata": metadata}) is outcome diff --git a/tests/unit/server/services/test_transcript_sync.py b/tests/unit/server/services/test_transcript_sync.py index e9a6db361..c08b5e6f4 100644 --- a/tests/unit/server/services/test_transcript_sync.py +++ b/tests/unit/server/services/test_transcript_sync.py @@ -20,7 +20,7 @@ import pytest -from ptc_agent.agent.transcript import TranscriptTarget, load_manifest +from ptc_agent.agent.transcript import TranscriptTarget, Window, load_manifest from ptc_agent.core.paths import SandboxLayout, WorkspaceLayout from src.server.database.conversation import ThreadPrefixes from src.server.database.thread_transcripts import ( @@ -385,7 +385,7 @@ def test_a_carried_file_whose_row_moved_is_a_stale_copy(): def _job(messages) -> _Job: - return _Job(TranscriptTarget(T1), messages, {"thread_id": T1}, FP, "cp-1", None) + return _Job(TranscriptTarget(T1), messages, {"thread_id": T1}, FP, "cp-1", None, Window()) def _big(): @@ -846,7 +846,7 @@ async def save_stored(thread_id, prefix, user_id, copy): @pytest.mark.asyncio async def test_a_live_save_keeps_the_rendered_checkpoint_and_empties_the_fingerprint(live): live.stored[""] = _stored(manifest='{"thread_id": "%s", "checkpoint_id": "cp-1"}' % T1) - assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")]) + assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")], window=Window()) [(prefix, copy)] = live.saves assert prefix == "" assert copy.fingerprint == "" and copy.checkpoint_id == "cp-1" @@ -860,7 +860,7 @@ async def test_a_live_save_takes_the_threads_checkpoint_when_its_export_is_late( must not look older than it.""" live.stored[""] = _stored(checkpoint_id="cp-1") live.stamped.return_value = "cp-2" - assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")]) + assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")], window=Window()) [(_, copy)] = live.saves assert copy.checkpoint_id == "cp-2" @@ -871,7 +871,7 @@ async def test_a_first_turn_live_save_takes_the_checkpoint_tip(live): replaced: an export that read the checkpoint before this compaction then dropped what the live save had added until the turn ended.""" live.tip.return_value = "cp-3" - assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")]) + assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")], window=Window()) [(_, copy)] = live.saves assert copy.checkpoint_id == "cp-3" live.tip.assert_awaited_once_with(T1) @@ -880,7 +880,7 @@ async def test_a_first_turn_live_save_takes_the_checkpoint_tip(live): # end's render down too. live.saves.clear() live.stamped.return_value = "cp-2" - assert await transcripts.save_live(TranscriptTarget(T1), [_message("bye")]) + assert await transcripts.save_live(TranscriptTarget(T1), [_message("bye")], window=Window()) [(_, copy)] = live.saves assert copy.checkpoint_id == "cp-2" live.tip.assert_awaited_once() @@ -917,7 +917,7 @@ async def test_a_live_save_of_a_task_replaces_only_that_task(live): manifest='{"task_id": "k1", "description": "d"}', **{"tasks__k1__run-0001.jsonl": "b" * 64}, ) - assert await transcripts.save_live(TranscriptTarget(T1, "k1"), [_message("go")]) + assert await transcripts.save_live(TranscriptTarget(T1, "k1"), [_message("go")], window=Window()) [(prefix, copy)] = live.saves assert prefix == "tasks/k1/" assert copy.files.keys() == {"tasks/k1/meta.json", "tasks/k1/run-0001.jsonl"} @@ -931,7 +931,7 @@ async def test_a_live_save_of_a_task_takes_the_tasks_latest_checkpoint(live): live.stored["tasks/k1/"] = _stored(checkpoint_id="task-cp-1") live.stamped.return_value = "cp-9" live.tip.return_value = "task-cp-2" - assert await transcripts.save_live(TranscriptTarget(T1, "k1"), [_message("go")]) + assert await transcripts.save_live(TranscriptTarget(T1, "k1"), [_message("go")], window=Window()) [(_, copy)] = live.saves assert copy.fingerprint == "" and copy.checkpoint_id == "task-cp-2" live.tip.assert_awaited_once_with(T1, "task:k1") @@ -940,8 +940,8 @@ async def test_a_live_save_of_a_task_takes_the_tasks_latest_checkpoint(live): @pytest.mark.asyncio async def test_a_live_save_with_nothing_new_writes_nothing(live): - assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")]) + assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")], window=Window()) [(_, first)] = live.saves live.stored[""] = first - assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")]) + assert await transcripts.save_live(TranscriptTarget(T1), [_message("hi")], window=Window()) assert len(live.saves) == 1 diff --git a/tests/unit/tools/secretary/conftest.py b/tests/unit/tools/secretary/conftest.py index 3309f4615..f2fd15ce0 100644 --- a/tests/unit/tools/secretary/conftest.py +++ b/tests/unit/tools/secretary/conftest.py @@ -24,6 +24,17 @@ def _every_workspace_name_is_free(): yield +@pytest.fixture(autouse=True) +def _no_committed_texts(): + """No test reaches the stored slices or the checkpoints: no turn has a + committed text to read, so every run's comes from its stored events.""" + with patch( + "src.server.services.history.committed.committed_texts", + AsyncMock(return_value={}), + ): + yield + + def workspace_manager( row: dict | None = None, delete: AsyncMock | None = None ) -> MagicMock: diff --git a/tests/unit/tools/secretary/test_extract_from_db.py b/tests/unit/tools/secretary/test_extract_from_db.py index 6cb12cb11..7114d96e4 100644 --- a/tests/unit/tools/secretary/test_extract_from_db.py +++ b/tests/unit/tools/secretary/test_extract_from_db.py @@ -40,17 +40,38 @@ def _chunk(text: str) -> dict: return {"event": "message_chunk", "data": {"content_type": "text", "content": text}} +# Each run's stored events by response id, as ``_response`` built them. +_STORED: dict[str, list[dict]] = {} + + def _response(turn_index: int, *texts: str, **fields) -> dict: + events = _STORED[f"r-{turn_index}"] = [_chunk(t) for t in texts] return { "conversation_response_id": f"r-{turn_index}", "conversation_thread_id": "t-1", "turn_index": turn_index, "status": "completed", - "sse_events": [_chunk(t) for t in texts], + "sse_events": events, **fields, } +@pytest.fixture(autouse=True) +def _stored_events(): + """The reader's rows come without their stored events, so a run's are + read by id; here, from what ``_response`` built.""" + _STORED.clear() + + async def read(response_ids): + return {i: {"sse_events": _STORED[i]} for i in response_ids if i in _STORED} + + events = AsyncMock(side_effect=read) + with patch( + "src.server.database.conversation.replay_rows.get_replay_responses", events + ): + yield events + + def _ids(runs: list[dict]) -> list[str]: return [r["conversation_response_id"] for r in runs] @@ -302,3 +323,44 @@ def test_join_multi_turn_huge_newest_keeps_newest_head(): assert "NEWSTART " in out # newest turn's head survives assert "OLD answer" not in out # older turn dropped assert out.endswith("[truncated — full output available in workspace]") + + +# --- main_text: where a run's text comes from -------------------------------- + +_COMMITTED = "src.server.services.history.committed.committed_texts" + + +@pytest.mark.asyncio +async def test_a_run_reads_the_text_its_turn_committed(): + committed = AsyncMock(return_value={"r-1": "committed text"}) + with patch(_RECENT, AsyncMock(return_value=[_response(1, "stored text")])), patch( + _COMMITTED, committed + ): + runs = await _runs_to_read("t-1") + assert runs[0]["main_text"] == "committed text" + committed.assert_awaited_once_with("t-1", {"r-1": 1}) + + +@pytest.mark.asyncio +async def test_a_run_without_committed_text_reads_its_stored_main_lane(): + sub = { + "event": "message_chunk", + "data": {"content_type": "text", "content": "sub", "agent": "task:abc"}, + } + run = _response(1, "main") + run["sse_events"].append(sub) + with patch(_RECENT, AsyncMock(return_value=[run])): + runs = await _runs_to_read("t-1") + assert runs[0]["main_text"] == "main" + + +@pytest.mark.asyncio +async def test_only_runs_without_committed_text_read_their_stored_events( + _stored_events, +): + with patch( + _RECENT, AsyncMock(return_value=[_response(1), _response(2, "stored")]) + ), patch(_COMMITTED, AsyncMock(return_value={"r-1": "committed"})): + runs = await _runs_to_read("t-1", turns=2) + assert [r["main_text"] for r in runs] == ["committed", "stored"] + _stored_events.assert_awaited_once_with(["r-2"])