From 514b1ff8bb5e9f3c1c4cee005b6ab2c16956391f Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Wed, 3 Jun 2026 20:18:57 -0700 Subject: [PATCH] =?UTF-8?q?fix:=20pre-push=20review=20=E2=80=94=20page=20m?= =?UTF-8?q?igration=20060=20backfill,=20heal=20legacy=20orphan=20tool=5Fus?= =?UTF-8?q?e=20on=20load?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - B1: page _backfill_content_addressed_attachments via a composite keyset cursor (message_id, created, attachment_id) instead of one un-paged fetchall() of every blob's bytes — bounds peak migration memory regardless of stored blob volume. Validated on the dev-DB snapshot: upgrade + downgrade clean, 5431 conversations preserved, blobs deduped + content-addressed. - B2: _native_from_provider_data strips orphan client tool-call blocks on load when the row's tool_calls column is empty (the truncated-mid-tool_use legacy hole), so a same-provider resume can't replay an unanswered tool_use — closing the Anthropic 400 and the Google tool-call resurrection path. Healthy (mirror-holds) rows decode byte-identically. --- tests/test_migration_060.py | 86 +++++++++++++ tests/test_native_tool_calls_mirror.py | 119 ++++++++++++++++++ turnstone/core/storage/_utils.py | 33 ++++- .../versions/060_unwrap_tool_envelopes.py | 69 ++++++---- 4 files changed, 280 insertions(+), 27 deletions(-) diff --git a/tests/test_migration_060.py b/tests/test_migration_060.py index 470f79cd..d3561a4f 100644 --- a/tests/test_migration_060.py +++ b/tests/test_migration_060.py @@ -596,3 +596,89 @@ class TestMigration060AttachmentBackfill: assert json.loads(refs) == [h1, h2] finally: engine.dispose() + + def test_backfill_pages_across_batch_boundary(self, tmp_path: Path) -> None: + """The backfill reads consumed rows PAGED (keyset on (message_id, + created, attachment_id)) so a large blob corpus never materialises at + once. Seed more than one page of attachments and assert the + accumulators span the boundary: a single message's ref-list keeps its + order across the page split, and a duplicate that lands on a LATER page + than its canonical still dedups + bumps the refcount. Runs against the + real ``_BATCH`` so a regression to an un-paged ``.fetchall()`` (or a + cursor that stalls/skips at the boundary) is caught.""" + import importlib.util + + mig_path = ( + Path(__file__).resolve().parent.parent + / "turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py" + ) + spec = importlib.util.spec_from_file_location("_mig_060_batch", mig_path) + assert spec is not None and spec.loader is not None + mig = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mig) + batch: int = mig._BATCH + + # m1 carries one full page + 2 → its ref-list straddles the boundary. + # m2 carries a single attachment whose bytes duplicate m1's first blob + # but sorts onto page 2 (canonical seen on page 1, dup found on page 2). + n = batch + 2 + contents = [f"blob-{i}".encode() for i in range(n)] + hashes = [hashlib.sha256(c).hexdigest() for c in contents] + + db_path = tmp_path / "060-att-paging.db" + cfg = _alembic_cfg(db_path) + command.upgrade(cfg, "059") + engine = sa.create_engine(f"sqlite:///{db_path}") + try: + with engine.begin() as conn: + _seed_row(conn, role="user", content="m1", tool_call_id="p1") + _seed_row(conn, role="user", content="m2", tool_call_id="p2") + m1 = conn.execute( + sa.text("SELECT id FROM conversations WHERE tool_call_id = 'p1'") + ).scalar_one() + m2 = conn.execute( + sa.text("SELECT id FROM conversations WHERE tool_call_id = 'p2'") + ).scalar_one() + # Zero-padded ids so lexicographic order == insertion order + # (created is constant → attachment_id is the sort tiebreaker). + for i, c in enumerate(contents): + _seed_attachment( + conn, + attachment_id=f"a{i:05d}", + content=c, + size_bytes=len(c), + message_id=m1, + ) + _seed_attachment( + conn, + attachment_id="b00000", + content=contents[0], + size_bytes=len(contents[0]), + message_id=m2, + ) + + command.upgrade(cfg, "060") + + with engine.connect() as conn: + # m2's dup of blob-0 collapses → n distinct blobs (not n + 1). + blob_rows = conn.execute( + sa.text("SELECT attachment_id, refcount FROM workstream_attachments") + ).fetchall() + assert len(blob_rows) == n + by_id = {r[0]: r[1] for r in blob_rows} + # The cross-page duplicate (blob-0) is referenced by m1 + m2. + assert by_id[hashes[0]] == 2 + # A blob unique to the second page keeps refcount 1. + assert by_id[hashes[-1]] == 1 + # m1's ref-list preserves order ACROSS the page boundary. + refs1 = conn.execute( + sa.text("SELECT attachments FROM conversations WHERE id = :i"), {"i": m1} + ).scalar_one() + assert json.loads(refs1) == hashes + # m2 references the shared (cross-page-deduped) blob. + refs2 = conn.execute( + sa.text("SELECT attachments FROM conversations WHERE id = :i"), {"i": m2} + ).scalar_one() + assert json.loads(refs2) == [hashes[0]] + finally: + engine.dispose() diff --git a/tests/test_native_tool_calls_mirror.py b/tests/test_native_tool_calls_mirror.py index 15ba9297..37add941 100644 --- a/tests/test_native_tool_calls_mirror.py +++ b/tests/test_native_tool_calls_mirror.py @@ -16,10 +16,15 @@ from __future__ import annotations import json from typing import Any +from turnstone.core.providers._anthropic import AnthropicProvider +from turnstone.core.providers._google import GoogleProvider from turnstone.core.storage._utils import ( normalize_native_for_save, + reconstruct_messages, + reconstruct_turns, strip_orphan_client_tool_blocks, ) +from turnstone.core.trajectory import dicts_from_turns _THINKING = {"type": "thinking", "thinking": "reasoning text", "signature": "sig-1"} _TOOL_USE = {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Paris"}} @@ -158,3 +163,117 @@ def test_save_messages_bulk_drops_orphan_native_tool_use(backend: Any) -> None: asst = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant") types = [b["type"] for b in asst.get("_provider_content", [])] assert types == ["thinking"] + + +# --------------------------------------------------------------------------- # +# Load-side self-heal: ``reconstruct_turns`` re-enforces the mirror for LEGACY +# rows that predate the save chokepoint — an orphan client tool-call in the +# native lane with an empty ``tool_calls`` column. Without this, an Anthropic +# resume replays the orphan ``tool_use`` (400) and Google resurrects the +# ``function`` block into ``tool_calls`` (an unanswered call). Both heal once +# the block is stripped at load. +# --------------------------------------------------------------------------- # +def _assistant_row(provider_data: str, tool_calls: str | None) -> list[Any]: + """A legacy-shaped assistant conversation row for ``reconstruct_turns``. + + Positional tuple: (id, role, content, tool_name, tool_call_id, + provider_data, tool_calls, source) — the 8-col prefix; event_id / is_error / + meta are absent (older fixture), exercising the length-guarded unpack. + """ + return [1, "assistant", "truncated mid tool_use", None, None, provider_data, tool_calls, None] + + +def _native_types(turns: list[Any]) -> list[str]: + native = turns[0].native + if native is None: + return [] + return [b["type"] for b in native.blocks if isinstance(b, dict)] + + +def test_reconstruct_strips_orphan_tool_use_bare_list() -> None: + row = _assistant_row(json.dumps([_THINKING, _TOOL_USE]), None) + assert _native_types(reconstruct_turns([row], "ws")) == ["thinking"] + + +def test_reconstruct_strips_orphan_in_producer_envelope() -> None: + # The {producer, blocks} storage envelope (a 060-tagged legacy row), not + # just the bare-list shape, also heals. + row = _assistant_row( + json.dumps({"producer": "anthropic", "blocks": [_THINKING, _TOOL_USE]}), None + ) + assert _native_types(reconstruct_turns([row], "ws")) == ["thinking"] + + +def test_reconstruct_drops_native_when_only_orphan() -> None: + # All-orphan native lane collapses to None (mirrors normalize_native_for_save's + # ``None``), not an empty ProviderNative. + row = _assistant_row(json.dumps([_TOOL_USE]), None) + assert reconstruct_turns([row], "ws")[0].native is None + + +def test_reconstruct_keeps_native_when_mirror_holds() -> None: + # Healthy case (matching tool_calls) is untouched — no over-stripping. + row = _assistant_row(json.dumps([_THINKING, _TOOL_USE]), _TOOL_CALLS_JSON) + assert _native_types(reconstruct_turns([row], "ws")) == ["thinking", "tool_use"] + + +def test_healed_row_yields_no_orphan_on_anthropic_wire() -> None: + # End-to-end: the healed Turn projects to an Anthropic payload with NO orphan + # ``tool_use`` block in any assistant content (the resume 400 is gone). + row = _assistant_row(json.dumps([_THINKING, _TOOL_USE]), None) + msgs = dicts_from_turns(reconstruct_turns([row], "ws")) + _system, converted = AnthropicProvider()._convert_messages(msgs) + for m in converted: + if m.get("role") != "assistant": + continue + content = m.get("content") + blocks = content if isinstance(content, list) else [] + assert all(b.get("type") != "tool_use" for b in blocks if isinstance(b, dict)) + + +def test_healed_row_yields_no_resurrected_call_on_google_wire() -> None: + # End-to-end (the path the brief MISSED): Google's _prepare_messages + # resurrects ``function`` blocks from ``_provider_content`` into + # ``tool_calls``. With the orphan stripped at load there is nothing to + # resurrect, so the assistant turn carries no unanswered ``tool_calls``. + row = _assistant_row(json.dumps([_GOOGLE_FN]), None) + msgs = dicts_from_turns(reconstruct_turns([row], "ws")) + prepared = GoogleProvider()._prepare_messages(msgs) + assert all(not m.get("tool_calls") for m in prepared if m.get("role") == "assistant") + + +def test_inflight_toolcall_preserved_on_history_load() -> None: + """An IN-FLIGHT tool call (the assistant issued it; the tool result hasn't + landed yet) must survive a ``/history`` load untouched. + + The self-heal is gated on an EMPTY ``tool_calls`` column — which a + legitimately-issued call never has: the save-time mirror + (``normalize_native_for_save``, applied to every assistant row by both save + paths on both backends) keeps the native lane and ``tool_calls`` in lockstep. + So /history (``reconstruct_messages(repair=False)``, which deliberately + preserves the trailing partial turn during tool execution) shows the call, + and a later resume replays it intact. Only the broken truncated-mid-tool_use + legacy shape (native ``tool_use`` with empty ``tool_calls``) is stripped. + Guards against the heal ever being widened to misfire on live tool calls.""" + rows = [ + [1, "user", "do a thing", None, None, None, None, None], + # In-flight assistant turn: tool_use in native AND a matching tool_calls + # entry (mirror holds); no following tool-result row yet. + [ + 2, + "assistant", + "", + None, + None, + json.dumps([_THINKING, _TOOL_USE]), + _TOOL_CALLS_JSON, + None, + ], + ] + msgs = reconstruct_messages(rows, "ws", repair=False) + assert len(msgs) == 2 # repair=False keeps the trailing in-flight turn + asst = msgs[-1] + assert asst["role"] == "assistant" + assert asst.get("tool_calls") # the issued call survives the load + pc_types = [b["type"] for b in asst.get("_provider_content", []) if isinstance(b, dict)] + assert "tool_use" in pc_types # native lane intact — the heal did NOT strip it diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index 0fd350aa..7023e56d 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -606,12 +606,25 @@ def _content_blocks(text: str | None, refs: list[AttachmentRef]) -> tuple[Conten return () -def _native_from_provider_data(provider_data: str | None) -> ProviderNative | None: +def _native_from_provider_data( + provider_data: str | None, tool_calls_json: str | None +) -> ProviderNative | None: """Decode the stored ``provider_data`` lane into a :class:`ProviderNative`. The storage envelope is ``{producer, blocks}`` (new) or a bare block list (legacy, no producer). A decode failure or non-list/non-envelope payload yields ``None`` (the lane is dropped — matching the prior best-effort decode). + + Load-side mirror self-heal: when the row carries no ``tool_calls`` (empty or + absent column), any *client* tool-call block in the native lane is an orphan — + a legacy truncated-mid-``tool_use`` row predating the save-time + ``normalize_native_for_save`` chokepoint — that would replay on a same-provider + resume with no matching ``tool_result`` (Anthropic 400; Google resurrects it + into ``tool_calls`` via its fidelity lane). Strip those blocks here so every + legacy row heals on read regardless of whether migration 060 tagged it + (reasoning / server-tool / web-search blocks kept). Mirrors + ``normalize_native_for_save``'s gate, including its ``None`` when nothing + survives. The healthy (mirror-holds) path is byte-identical to a plain decode. """ if not provider_data: return None @@ -620,10 +633,18 @@ def _native_from_provider_data(provider_data: str | None) -> ProviderNative | No except (json.JSONDecodeError, TypeError): return None if isinstance(parsed, dict) and "blocks" in parsed: - return ProviderNative(producer=parsed.get("producer") or "", blocks=tuple(parsed["blocks"])) - if isinstance(parsed, list): - return ProviderNative(producer="", blocks=tuple(parsed)) - return None + producer = parsed.get("producer") or "" + blocks = list(parsed["blocks"]) + elif isinstance(parsed, list): + producer = "" + blocks = parsed + else: + return None + if not _has_tool_calls(tool_calls_json): + blocks = strip_orphan_client_tool_blocks(blocks) + if not blocks: + return None + return ProviderNative(producer=producer, blocks=tuple(blocks)) def _source_meta_from_json(meta_json: str | None) -> dict[str, Any] | None: @@ -705,7 +726,7 @@ def reconstruct_turns( Role.ASSISTANT, _content_blocks(content, []), tool_calls=_tool_calls_from_json(tool_calls_json), - native=_native_from_provider_data(provider_data), + native=_native_from_provider_data(provider_data, tool_calls_json), meta=meta, ) ) diff --git a/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py b/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py index 4c3c5496..5e2f0f1c 100644 --- a/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py +++ b/turnstone/core/storage/migrations/versions/060_unwrap_tool_envelopes.py @@ -361,14 +361,6 @@ def _backfill_content_addressed_attachments(bind: sa.engine.Connection) -> None: sa.column("attachments", sa.Text), ) - # Read every consumed legacy row in (message_id, created, attachment_id) - # order so each message's ref-list preserves the original attachment order. - rows = bind.execute( - sa.select(wa.c.attachment_id, wa.c.message_id, wa.c.content) - .where(wa.c.message_id.is_not(None)) - .order_by(wa.c.message_id, wa.c.created, wa.c.attachment_id) - ).fetchall() - # new_id (content hash) -> canonical old id kept as that blob's row. canonical_old_id: dict[str, str] = {} # new_id -> set of distinct message ids referencing it (refcount source). @@ -378,19 +370,54 @@ def _backfill_content_addressed_attachments(bind: sa.engine.Connection) -> None: # old ids to delete (duplicates that collapsed into a canonical row). drop_old_ids: list[str] = [] - for old_id, message_id, content in rows: - raw = content if isinstance(content, (bytes, bytearray)) else b"" - new_id = hashlib.sha256(bytes(raw)).hexdigest() - if new_id not in canonical_old_id: - canonical_old_id[new_id] = old_id - refcounting[new_id] = set() - elif old_id != canonical_old_id[new_id]: - # A distinct legacy row carrying identical bytes — collapse it. - drop_old_ids.append(old_id) - refcounting[new_id].add(int(message_id)) - bucket = per_message.setdefault(int(message_id), []) - if new_id not in bucket: - bucket.append(new_id) + # Read every consumed legacy row in (message_id, created, attachment_id) + # order so each message's ref-list preserves the original attachment order. + # PAGED via a composite keyset cursor (like steps 1/1b) so the LargeBinary + # ``content`` of the whole corpus is never materialised at once — only one + # _BATCH of blobs is resident per iteration (each is hashed then dropped), + # bounding peak memory regardless of total stored blob volume. The + # accumulators above span the full scan and flush after it, and the keyset + # order matches the ORDER BY exactly, so paging is byte-equivalent to a + # single ordered fetch. ``created`` is NOT NULL and ``attachment_id`` is the + # PK, so the (message_id, created, attachment_id) tuple is unique and the + # cursor never stalls or skips at a page boundary. + last_key: tuple[int, str, str] | None = None + while True: + page = sa.select(wa.c.attachment_id, wa.c.message_id, wa.c.content, wa.c.created).where( + wa.c.message_id.is_not(None) + ) + if last_key is not None: + last_msg, last_created, last_aid = last_key + page = page.where( + sa.or_( + wa.c.message_id > last_msg, + sa.and_(wa.c.message_id == last_msg, wa.c.created > last_created), + sa.and_( + wa.c.message_id == last_msg, + wa.c.created == last_created, + wa.c.attachment_id > last_aid, + ), + ) + ) + rows = bind.execute( + page.order_by(wa.c.message_id, wa.c.created, wa.c.attachment_id).limit(_BATCH) + ).fetchall() + if not rows: + break + for old_id, message_id, content, created in rows: + last_key = (int(message_id), created, old_id) + raw = content if isinstance(content, (bytes, bytearray)) else b"" + new_id = hashlib.sha256(bytes(raw)).hexdigest() + if new_id not in canonical_old_id: + canonical_old_id[new_id] = old_id + refcounting[new_id] = set() + elif old_id != canonical_old_id[new_id]: + # A distinct legacy row carrying identical bytes — collapse it. + drop_old_ids.append(old_id) + refcounting[new_id].add(int(message_id)) + bucket = per_message.setdefault(int(message_id), []) + if new_id not in bucket: + bucket.append(new_id) # Re-key each canonical row's PK to its content hash and set refcount/origin. # Re-key first (while the duplicates still hold their old PKs), then delete