From caa589e39b6e8e642f4e5f9612a42fdb7314a4d4 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 2 Jun 2026 21:29:03 -0700 Subject: [PATCH] feat(storage): tag the native lane with its producer ({producer, blocks}) Persist provider_data as a {producer, blocks} envelope (producer = the generating provider's name) so the lowering layer can later replay the native lane verbatim only to its producer and rebuild from neutral fields for any other. The envelope is storage-only: prepare_provider_data_for_save runs the P1 mirror on the bare block list and then wraps; reconstruct_messages dual-reads (new envelope OR legacy bare list), unwraps to a bare _provider_content list (every consumer's contract), and surfaces the producer on the stripped-before-wire _producer side channel. The producer threads the same four save layers as is_error (facade -> protocol -> SQLite + Postgres); the live assistant save tags it from self._provider.provider_name and the fork carries _producer. Legacy rows need no migration to keep working (dual-read); the one-shot backfill that tags them is a follow-up. Sub-commit 2a of the canonical-trajectory storage cut. --- tests/test_provider_data_envelope.py | 94 +++++++++++++++++++++++++++ turnstone/core/memory.py | 2 + turnstone/core/session.py | 2 + turnstone/core/storage/_postgresql.py | 14 ++-- turnstone/core/storage/_protocol.py | 1 + turnstone/core/storage/_sqlite.py | 14 ++-- turnstone/core/storage/_utils.py | 49 +++++++++++++- 7 files changed, 167 insertions(+), 9 deletions(-) create mode 100644 tests/test_provider_data_envelope.py diff --git a/tests/test_provider_data_envelope.py b/tests/test_provider_data_envelope.py new file mode 100644 index 00000000..875f6721 --- /dev/null +++ b/tests/test_provider_data_envelope.py @@ -0,0 +1,94 @@ +"""provider_data ``{producer, blocks}`` storage envelope (#5 sub-commit 2a). + +Native blocks persist wrapped with the generating provider's name so the lowering layer +can later replay them verbatim only to their producer. The envelope is storage-only: +``reconstruct_messages`` unwraps it back to a bare block list (every ``_provider_content`` +consumer requires a plain list) and surfaces the producer on the ``_producer`` side +channel. Legacy bare-list rows are dual-read unchanged. +""" + +from __future__ import annotations + +import json +from typing import Any + +from turnstone.core.storage._utils import prepare_provider_data_for_save, wrap_provider_data + +_BLOCKS = [{"type": "thinking", "thinking": "reasoning", "signature": "s"}] +_BLOCKS_JSON = json.dumps(_BLOCKS) + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # +def test_wrap_adds_envelope() -> None: + out = wrap_provider_data(_BLOCKS_JSON, "anthropic") + assert out is not None + assert json.loads(out) == {"producer": "anthropic", "blocks": _BLOCKS} + + +def test_wrap_no_producer_keeps_bare_list() -> None: + assert wrap_provider_data(_BLOCKS_JSON, None) == _BLOCKS_JSON + + +def test_wrap_none_passthrough() -> None: + assert wrap_provider_data(None, "anthropic") is None + + +def test_wrap_already_wrapped_is_unchanged() -> None: + wrapped = json.dumps({"producer": "x", "blocks": _BLOCKS}) + assert wrap_provider_data(wrapped, "anthropic") == wrapped + + +def test_prepare_strips_orphan_then_wraps() -> None: + # A tool_use with no matching tool_calls is an orphan (P1 mirror) → stripped, + # and what survives is wrapped with the producer. + pd = json.dumps( + [ + {"type": "thinking", "thinking": "t"}, + {"type": "tool_use", "id": "c1", "name": "x", "input": {}}, + ] + ) + out = prepare_provider_data_for_save("assistant", pd, None, "anthropic") + assert out is not None + env = json.loads(out) + assert env["producer"] == "anthropic" + assert [b["type"] for b in env["blocks"]] == ["thinking"] + + +# --------------------------------------------------------------------------- # +# round-trip via a real (ephemeral) backend: save → reconstruct dual-read +# --------------------------------------------------------------------------- # +def test_round_trip_surfaces_bare_blocks_and_producer(backend: Any) -> None: + ws = "ws-env-1" + backend.save_message(ws, "assistant", "hi", provider_data=_BLOCKS_JSON, producer="anthropic") + a = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant") + assert a["_provider_content"] == _BLOCKS # bare list — the consumer contract + assert a["_producer"] == "anthropic" + + +def test_legacy_bare_list_dual_read(backend: Any) -> None: + ws = "ws-env-2" + # producer omitted → stored as a bare list (legacy shape), read back unchanged. + backend.save_message(ws, "assistant", "hi", provider_data=_BLOCKS_JSON) + a = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant") + assert a["_provider_content"] == _BLOCKS + assert "_producer" not in a + + +def test_bulk_round_trip_with_producer(backend: Any) -> None: + ws = "ws-env-3" + backend.save_messages_bulk( + [ + { + "ws_id": ws, + "role": "assistant", + "content": "hi", + "provider_data": _BLOCKS_JSON, + "producer": "google", + } + ] + ) + a = next(m for m in backend.load_messages(ws, repair=False) if m["role"] == "assistant") + assert a["_provider_content"] == _BLOCKS + assert a["_producer"] == "google" diff --git a/turnstone/core/memory.py b/turnstone/core/memory.py index 8859481b..117f49ee 100644 --- a/turnstone/core/memory.py +++ b/turnstone/core/memory.py @@ -44,6 +44,7 @@ def save_message( source: str | None = None, event_id: int | None = None, is_error: bool = False, + producer: str | None = None, ) -> int: """Log a message to the conversations table. @@ -71,6 +72,7 @@ def save_message( source=source, event_id=event_id, is_error=is_error, + producer=producer, ) except Exception: log.warning("Failed to save message for ws=%s role=%s", ws_id, role, exc_info=True) diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 7e995dd0..dd613e5f 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -2591,6 +2591,7 @@ class ChatSession: "provider_data": pd_str, "source": src if isinstance(src, str) and src else None, "is_error": bool(msg.get("is_error", False)), + "producer": msg.get("_producer"), } ) save_messages_bulk(bulk_rows) @@ -3935,6 +3936,7 @@ class ChatSession: provider_data=provider_data, tool_calls=tool_calls_json, event_id=self._ui_event_id(), + producer=self._provider.provider_name if self._provider else None, ) tool_calls = assistant_msg.get("tool_calls") diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 13c3932b..d56f5d6a 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -109,10 +109,10 @@ from turnstone.core.storage._utils import ( from turnstone.core.storage._utils import ( escape_like as _escape_like, ) -from turnstone.core.storage._utils import normalize_native_for_save, sanitize_text from turnstone.core.storage._utils import ( normalize_search_terms as _normalize_search_terms, ) +from turnstone.core.storage._utils import prepare_provider_data_for_save, sanitize_text from turnstone.core.storage._utils import ( reconstruct_messages as _reconstruct_messages, ) @@ -292,10 +292,13 @@ class PostgreSQLBackend: source: str | None = None, event_id: int | None = None, is_error: bool = False, + producer: str | None = None, ) -> int: now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") content = sanitize_text(content) - provider_data = normalize_native_for_save(role, sanitize_text(provider_data), tool_calls) + provider_data = prepare_provider_data_for_save( + role, sanitize_text(provider_data), tool_calls, producer + ) source = sanitize_text(source) with self._conn() as conn: result = conn.execute( @@ -339,8 +342,11 @@ class PostgreSQLBackend: "content": sanitize_text(row["content"]), "tool_name": row.get("tool_name"), "tool_call_id": row.get("tool_call_id"), - "provider_data": normalize_native_for_save( - row["role"], sanitize_text(row.get("provider_data")), row.get("tool_calls") + "provider_data": prepare_provider_data_for_save( + row["role"], + sanitize_text(row.get("provider_data")), + row.get("tool_calls"), + row.get("producer"), ), "tool_calls": row.get("tool_calls"), "_source": sanitize_text(row.get("source")), diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index 5893105f..79603663 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -160,6 +160,7 @@ class StorageBackend(Protocol): source: str | None = None, event_id: int | None = None, is_error: bool = False, + producer: str | None = None, ) -> int: """Log a message to the conversations table. diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index 19b3e4eb..84d0dc2e 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -109,10 +109,10 @@ from turnstone.core.storage._utils import ( from turnstone.core.storage._utils import ( escape_like as _escape_like, ) -from turnstone.core.storage._utils import normalize_native_for_save, sanitize_text from turnstone.core.storage._utils import ( normalize_search_terms as _normalize_search_terms, ) +from turnstone.core.storage._utils import prepare_provider_data_for_save, sanitize_text from turnstone.core.storage._utils import ( reconstruct_messages as _reconstruct_messages, ) @@ -326,10 +326,13 @@ class SQLiteBackend: source: str | None = None, event_id: int | None = None, is_error: bool = False, + producer: str | None = None, ) -> int: now = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S") content = sanitize_text(content) - provider_data = normalize_native_for_save(role, sanitize_text(provider_data), tool_calls) + provider_data = prepare_provider_data_for_save( + role, sanitize_text(provider_data), tool_calls, producer + ) source = sanitize_text(source) with self._conn() as conn: result = conn.execute( @@ -387,8 +390,11 @@ class SQLiteBackend: "content": sanitize_text(row["content"]), "tool_name": row.get("tool_name"), "tool_call_id": row.get("tool_call_id"), - "provider_data": normalize_native_for_save( - row["role"], sanitize_text(row.get("provider_data")), row.get("tool_calls") + "provider_data": prepare_provider_data_for_save( + row["role"], + sanitize_text(row.get("provider_data")), + row.get("tool_calls"), + row.get("producer"), ), "tool_calls": row.get("tool_calls"), "_source": sanitize_text(row.get("source")), diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index a67d864b..7ccdf0fb 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -83,6 +83,44 @@ def normalize_native_for_save( return json.dumps(kept) if kept else None +def wrap_provider_data(provider_data: str | None, producer: str | None) -> str | None: + """Wrap a bare native-block list in the storage envelope ``{producer, blocks}``. + + ``producer`` (the provider that generated the turn) lets the lowering layer replay the + native lane verbatim only to its producing provider. Storage-only: the envelope is + unwrapped back to a bare block list by :func:`reconstruct_messages`, so every + ``_provider_content`` consumer still sees a plain list. Input that is empty, already + wrapped, unparseable, or has no ``producer`` is returned unchanged (the last keeps the + legacy bare-list shape, which reconstruct dual-reads). + """ + if not provider_data or not producer: + return provider_data + try: + blocks = json.loads(provider_data) + except (json.JSONDecodeError, TypeError): + return provider_data + if isinstance(blocks, dict) and "blocks" in blocks: + return provider_data + return json.dumps({"producer": producer, "blocks": blocks}) + + +def prepare_provider_data_for_save( + role: str | None, + provider_data: str | None, + tool_calls_json: str | None, + producer: str | None, +) -> str | None: + """Save-boundary preparation of the native lane: enforce the mirror, then wrap. + + The single entry point both backends' save paths call: + :func:`normalize_native_for_save` (the native↔tool_calls mirror, on the bare block + list) followed by :func:`wrap_provider_data` (the ``{producer, blocks}`` envelope). + """ + return wrap_provider_data( + normalize_native_for_save(role, provider_data, tool_calls_json), producer + ) + + def _attachment_to_content_part(att: dict[str, Any]) -> dict[str, Any] | None: """Convert a stored attachment row into an OpenAI-style content part. @@ -491,7 +529,16 @@ def reconstruct_messages( msg: dict[str, Any] = {"role": "assistant", "content": content or ""} if provider_data: with contextlib.suppress(json.JSONDecodeError, TypeError): - msg["_provider_content"] = json.loads(provider_data) + parsed = json.loads(provider_data) + # Storage envelope ``{producer, blocks}`` (new) vs bare list (legacy): + # surface bare blocks as ``_provider_content`` (every consumer expects a + # plain list) and carry the producer on a stripped-before-wire side channel. + if isinstance(parsed, dict) and "blocks" in parsed: + msg["_provider_content"] = parsed["blocks"] + if parsed.get("producer"): + msg["_producer"] = parsed["producer"] + else: + msg["_provider_content"] = parsed if tool_calls_json: with contextlib.suppress(json.JSONDecodeError, TypeError): msg["tool_calls"] = json.loads(tool_calls_json)