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)