From 6abb2698f78ed0cb8e3fefc20df2be00f06bca71 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Thu, 7 May 2026 22:40:50 -0700 Subject: [PATCH] fix: apply repair=False to all display-read load_messages call sites --- tests/test_sessions.py | 65 ++++++++++++++++++ tests/test_workstream_endpoints.py | 91 +++++++++++++++++++++++++ turnstone/console/coordinator_client.py | 17 +++-- turnstone/console/server.py | 4 +- turnstone/core/memory.py | 4 +- turnstone/core/session_routes.py | 5 +- turnstone/core/storage/_postgresql.py | 6 +- turnstone/core/storage/_protocol.py | 13 +++- turnstone/core/storage/_sqlite.py | 6 +- turnstone/core/storage/_utils.py | 19 ++++++ 10 files changed, 217 insertions(+), 13 deletions(-) diff --git a/tests/test_sessions.py b/tests/test_sessions.py index 74eb8fc5..9b1b57e1 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -451,6 +451,71 @@ class TestInterruptedWorkstreamRepair: assert msgs[1]["role"] == "assistant" assert msgs[2]["role"] == "user" + def test_repair_false_preserves_partial_trailing_turn(self, tmp_db): + """``repair=False`` is the display-read contract for ``/history``. + + The default repair pass strips the trailing + ``assistant(tool_calls)`` when not all tool results are persisted + — correct for ``session.resume`` (LLM context), wrong for the + REST display read. A user refreshing the coordinator page mid- + tool-execution would otherwise lose the entire trailing turn + from the UI. ``repair=False`` returns the raw persisted state. + """ + import json + + tc_json = json.dumps( + [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"pwd"}'}, + }, + ] + ) + save_message("s1", "user", "hello") + save_message("s1", "assistant", "Checking", tool_calls=tc_json) + save_message("s1", "tool", "file.txt", tool_call_id="call_1") + # No call_2 result persisted — mid-execution refresh. + msgs = get_storage().load_messages("s1", repair=False) + # All three rows survive — the trailing partial turn is what the + # operator was actually watching live. + assert [m["role"] for m in msgs] == ["user", "assistant", "tool"] + assert msgs[1].get("tool_calls") and len(msgs[1]["tool_calls"]) == 2 + assert msgs[2]["tool_call_id"] == "call_1" + + def test_repair_false_does_not_synthesize_orphan_results(self, tmp_db): + """``repair=False`` must NOT splice synthetic ``"Tool execution + was cancelled."`` rows for mid-conversation orphans either — + the operator never saw those rows, and showing them would + invent UI content that doesn't reflect persisted state. + """ + import json + + tc_json = json.dumps( + [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + }, + ] + ) + save_message("s1", "user", "first") + save_message("s1", "assistant", "Working", tool_calls=tc_json) + # Cancel landed before any tool result — next turn happens. + save_message("s1", "user", "second") + save_message("s1", "assistant", "ok") + msgs = get_storage().load_messages("s1", repair=False) + roles = [m["role"] for m in msgs] + # No synthetic tool row spliced after the orphaned tool_calls. + assert roles == ["user", "assistant", "user", "assistant"] + assert all(m["role"] != "tool" for m in msgs) + # ── Workstream config persistence ───────────────────────────────────── diff --git a/tests/test_workstream_endpoints.py b/tests/test_workstream_endpoints.py index 6121a9b1..d8161466 100644 --- a/tests/test_workstream_endpoints.py +++ b/tests/test_workstream_endpoints.py @@ -898,6 +898,97 @@ class TestHistoryInteractive: # Above-cap → clamps to 500 (response is still 200; we have 4 rows). assert client.get(base, params={"limit": 999}).status_code == 200 + def test_returns_partial_trailing_turn_during_tool_execution(self, _inject_storage): + """The ``/history`` REST endpoint is a *display* read and must + surface partial state. When the operator refreshes the page + mid-tool-execution — assistant ``tool_calls`` saved, only some + results saved — the trailing turn must come back on the wire so + the UI can render what the operator was watching live. + + Storage's ``load_messages`` defaults to a repair pass that + strips this exact shape (correct for ``session.resume``, wrong + for display). ``make_history_handler`` must opt out via + ``repair=False``; flipping that flag back on breaks this test. + """ + import json + + ws_id = "ws-mid-exec" + _inject_storage.register_workstream(ws_id, kind="interactive", user_id="test-user") + _inject_storage.save_message(ws_id, "user", "kick off") + tc_json = json.dumps( + [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + }, + { + "id": "call_2", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"pwd"}'}, + }, + ] + ) + _inject_storage.save_message(ws_id, "assistant", "Working", tool_calls=tc_json) + _inject_storage.save_message(ws_id, "tool", "file.txt", tool_call_id="call_1") + # call_2 result not yet persisted — operator refreshes here. + mock_ws = MagicMock() + mock_ws.id = ws_id + mock_mgr = MagicMock() + mock_mgr.get.return_value = mock_ws + client = _build_history_app(mock_mgr, _inject_storage) + + r = client.get(f"/v1/api/workstreams/{ws_id}/history") + assert r.status_code == 200 + roles = [m.get("role") for m in r.json()["messages"]] + # All three rows survive — the trailing assistant + partial + # tool result are what the operator was watching live. The + # default-repair shape would have been just ``["user"]``. + assert roles == ["user", "assistant", "tool"] + + # Confirm the default-repair path collapses this to just the + # user message — locks in the regression contract. + with_repair = _inject_storage.load_messages(ws_id, repair=True) + assert [m.get("role") for m in with_repair] == ["user"] + + def test_history_does_not_synthesize_orphan_results(self, _inject_storage): + """``repair=False`` via ``/history`` must NOT splice synthetic + ``"Tool execution was cancelled."`` rows for mid-conversation + orphaned tool_calls — the operator never saw those rows, and + showing them would invent UI content that doesn't reflect + persisted state. + """ + import json + + ws_id = "ws-orphan-mid" + _inject_storage.register_workstream(ws_id, kind="interactive", user_id="test-user") + _inject_storage.save_message(ws_id, "user", "first") + tc_json = json.dumps( + [ + { + "id": "call_1", + "type": "function", + "function": {"name": "bash", "arguments": '{"command":"ls"}'}, + }, + ] + ) + _inject_storage.save_message(ws_id, "assistant", "Working", tool_calls=tc_json) + # Cancel landed before any tool result — next turn happens. + _inject_storage.save_message(ws_id, "user", "second") + _inject_storage.save_message(ws_id, "assistant", "ok") + mock_ws = MagicMock() + mock_ws.id = ws_id + mock_mgr = MagicMock() + mock_mgr.get.return_value = mock_ws + client = _build_history_app(mock_mgr, _inject_storage) + + r = client.get(f"/v1/api/workstreams/{ws_id}/history") + assert r.status_code == 200 + roles = [m.get("role") for m in r.json()["messages"]] + # No synthetic tool row spliced after the orphaned tool_calls. + assert roles == ["user", "assistant", "user", "assistant"] + assert all(m.get("role") != "tool" for m in r.json()["messages"]) + class TestBuildHistoryReminderPropagation: """``_build_history`` must surface the ``_reminders`` side-channel on diff --git a/turnstone/console/coordinator_client.py b/turnstone/console/coordinator_client.py index 3640f9c0..3cb038fb 100644 --- a/turnstone/console/coordinator_client.py +++ b/turnstone/console/coordinator_client.py @@ -1488,12 +1488,17 @@ class CoordinatorClient: is_own_child = full.get("parent_ws_id") == self._coord_ws_id if not (is_self or is_own_child): return miss - # load_messages returns the full history in chronological order - # (no limit param in the Protocol) — slice the tail here. Defensive + # load_messages returns the full history in chronological order. + # We slice the tail in Python because the SQL tail-N is + # approximate across conversation boundaries. Defensive # try/except: storage errors should not break inspect. messages: list[Any] = [] try: - all_msgs = self._storage.load_messages(ws_id) + # repair=False — inspect is a display read (admin viewing a + # child's history in the tree UI). The LLM-context repair + # pass would strip trailing partial turns the operator is + # watching. + all_msgs = self._storage.load_messages(ws_id, repair=False) if message_limit and message_limit > 0: messages = all_msgs[-message_limit:] else: @@ -1725,7 +1730,11 @@ def _last_assistant_text(storage: Any, ws_id: str) -> str | None: in just to surface its final turn. """ try: - rows = storage.load_messages(ws_id, limit=_WAIT_MESSAGE_TAIL_LIMIT) + # repair=False — this reads the tail for display ("waiting on" bubble). + # The repair pass would strip a trailing partial assistant turn, + # making us return the penultimate assistant message instead of the + # one the operator is watching. + rows = storage.load_messages(ws_id, limit=_WAIT_MESSAGE_TAIL_LIMIT, repair=False) except Exception: log.debug("coord_client.wait.load_messages_failed ws=%s", ws_id, exc_info=True) return None diff --git a/turnstone/console/server.py b/turnstone/console/server.py index da693443..76d9ba09 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -1328,7 +1328,9 @@ async def cluster_ws_detail(request: Request) -> JSONResponse: # Tail-N bound pushed into SQL (load_messages supports limit # on both backends). Offloaded to the default executor so # the async SSE loop stays unblocked under rapid fan-out. - messages = await asyncio.to_thread(storage.load_messages, ws_id, limit=limit) + messages = await asyncio.to_thread( + storage.load_messages, ws_id, limit=limit, repair=False + ) except Exception: log.debug("cluster_ws_detail.load_messages_failed", exc_info=True) diff --git a/turnstone/core/memory.py b/turnstone/core/memory.py index 8c7bd6c2..65374c1e 100644 --- a/turnstone/core/memory.py +++ b/turnstone/core/memory.py @@ -79,10 +79,10 @@ def save_messages_bulk(rows: list[dict[str, Any]]) -> None: log.warning("Failed to bulk-save %d messages", len(rows), exc_info=True) -def load_messages(ws_id: str) -> list[dict[str, Any]]: +def load_messages(ws_id: str, *, repair: bool = True) -> list[dict[str, Any]]: """Load messages for a workstream and reconstruct OpenAI message format.""" try: - return get_storage().load_messages(ws_id) + return get_storage().load_messages(ws_id, repair=repair) except Exception: log.warning("Failed to load messages for ws=%s", ws_id, exc_info=True) return [] diff --git a/turnstone/core/session_routes.py b/turnstone/core/session_routes.py index 811a9965..b1691d23 100644 --- a/turnstone/core/session_routes.py +++ b/turnstone/core/session_routes.py @@ -2269,7 +2269,10 @@ def make_history_handler(cfg: SessionEndpointConfig) -> Handler: messages: list[dict[str, Any]] = [] if storage is not None: try: - messages = await asyncio.to_thread(storage.load_messages, ws_id, limit=limit) + # repair=False — display read; see reconstruct_messages docstring. + messages = await asyncio.to_thread( + storage.load_messages, ws_id, limit=limit, repair=False + ) except Exception: log.debug("ws.history.load_failed ws=%s", ws_id[:8], exc_info=True) # Audit-trail decoration — attach persisted intent_verdict and diff --git a/turnstone/core/storage/_postgresql.py b/turnstone/core/storage/_postgresql.py index 1adbb424..f9eb3208 100644 --- a/turnstone/core/storage/_postgresql.py +++ b/turnstone/core/storage/_postgresql.py @@ -242,7 +242,9 @@ class PostgreSQLBackend: ) conn.commit() - def load_messages(self, ws_id: str, *, limit: int | None = None) -> list[dict[str, Any]]: + def load_messages( + self, ws_id: str, *, limit: int | None = None, repair: bool = True + ) -> list[dict[str, Any]]: with self._conn() as conn: if limit is not None and limit > 0: rows = conn.execute( @@ -286,7 +288,7 @@ class PostgreSQLBackend: if limit is not None and limit > 0: message_ids = [r[0] for r in rows] attachments = self.load_attachments_for_messages(ws_id, message_ids=message_ids) - return _reconstruct_messages(list(rows), ws_id, attachments or None) + return _reconstruct_messages(list(rows), ws_id, attachments or None, repair=repair) def delete_messages_after(self, ws_id: str, keep_count: int) -> int: with self._conn() as conn: diff --git a/turnstone/core/storage/_protocol.py b/turnstone/core/storage/_protocol.py index f81692a8..e92b273e 100644 --- a/turnstone/core/storage/_protocol.py +++ b/turnstone/core/storage/_protocol.py @@ -162,7 +162,9 @@ class StorageBackend(Protocol): """ ... - def load_messages(self, ws_id: str, *, limit: int | None = None) -> list[dict[str, Any]]: + def load_messages( + self, ws_id: str, *, limit: int | None = None, repair: bool = True + ) -> list[dict[str, Any]]: """Load messages for a workstream and reconstruct OpenAI message format. ``limit`` caps the number of underlying conversation rows fetched @@ -172,6 +174,15 @@ class StorageBackend(Protocol): entries than ``limit`` when a tool-call group splits across the boundary; callers that need strict tail-N semantics must slice again client-side. Default ``None`` fetches the full history. + + ``repair`` (default True) post-processes the result into a + wire-shape valid for an LLM round-trip — drops a trailing + ``assistant(tool_calls)`` whose results aren't all present and + fills mid-conversation orphans with synthetic cancellation + results. Display-only readers (``/history`` REST) should pass + ``repair=False`` so the user sees the actual partial state + instead of having the trailing turn silently stripped during + live tool execution. """ ... diff --git a/turnstone/core/storage/_sqlite.py b/turnstone/core/storage/_sqlite.py index ce05c1cc..f33e3ce7 100644 --- a/turnstone/core/storage/_sqlite.py +++ b/turnstone/core/storage/_sqlite.py @@ -306,7 +306,9 @@ class SQLiteBackend: self._fts5_available = False conn.commit() - def load_messages(self, ws_id: str, *, limit: int | None = None) -> list[dict[str, Any]]: + def load_messages( + self, ws_id: str, *, limit: int | None = None, repair: bool = True + ) -> list[dict[str, Any]]: with self._conn() as conn: if limit is not None and limit > 0: # Tail-N: fetch the last `limit` rows via DESC + LIMIT @@ -354,7 +356,7 @@ class SQLiteBackend: if limit is not None and limit > 0: message_ids = [r[0] for r in rows] attachments = self.load_attachments_for_messages(ws_id, message_ids=message_ids) - return _reconstruct_messages(list(rows), ws_id, attachments or None) + return _reconstruct_messages(list(rows), ws_id, attachments or None, repair=repair) def delete_messages_after(self, ws_id: str, keep_count: int) -> int: with self._conn() as conn: diff --git a/turnstone/core/storage/_utils.py b/turnstone/core/storage/_utils.py index d1cd7542..7061f591 100644 --- a/turnstone/core/storage/_utils.py +++ b/turnstone/core/storage/_utils.py @@ -286,6 +286,8 @@ def reconstruct_messages( rows: list[Any], ws_id: str, attachments_by_msg: dict[int, list[dict[str, Any]]] | None = None, + *, + repair: bool = True, ) -> list[dict[str, Any]]: """Reconstruct OpenAI message format from stored conversation rows. @@ -299,6 +301,17 @@ def reconstruct_messages( When ``attachments_by_msg`` is provided, any user row whose id has attachments is rebuilt with multipart list content (text + image_url/document parts). + + When ``repair`` is True (default) the result is post-processed to + produce a wire-shape valid for an LLM round-trip: the trailing + ``assistant(tool_calls)`` turn is dropped if not all tool_call ids + have a matching tool result, and any mid-conversation orphaned + tool_calls are filled with synthetic cancellation results. Callers + that consume the messages as LLM context (e.g. ``session.resume``) + must keep this on. Callers reading for *display* (the ``/history`` + REST endpoint) should pass ``repair=False`` so the user sees the + actual partial state — refreshing during tool execution otherwise + silently drops the trailing turn from the UI. """ messages: list[dict[str, Any]] = [] for row in rows: @@ -374,6 +387,12 @@ def reconstruct_messages( tmsg["_reminders"] = json.loads(reminders_json) messages.append(tmsg) + if not repair: + # Both passes below are LLM-context corrections — trailing-turn + # strip and orphan synthesis. Display callers want neither; see + # the reconstruct_messages docstring. + return messages + # Repair: strip trailing incomplete tool call turns while messages: tail_tools = 0