From dbd078b3ec47b48bfe1087e9017a6cc511ff2c47 Mon Sep 17 00:00:00 2001 From: Genius229 <95580080+Genius229@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:38:39 +0400 Subject: [PATCH] fix(sessions): make failed-run history recovery idempotent --- strix/core/execution.py | 7 +- strix/core/sessions.py | 88 ++++++-- tests/test_execution.py | 1 + tests/test_tool_history_recovery.py | 300 ++++++++++++++++++++++++++++ 4 files changed, 376 insertions(+), 20 deletions(-) create mode 100644 tests/test_tool_history_recovery.py diff --git a/strix/core/execution.py b/strix/core/execution.py index dfcd39fa..568daa9a 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -29,7 +29,7 @@ from strix.core.inputs import child_initial_input from strix.core.sessions import ( enforce_image_budget, open_agent_session, - replace_session_items, + recover_session_items, seed_initial_input, strip_all_images_from_session, ) @@ -163,11 +163,10 @@ async def _salvage_stream_to_session( except Exception: logger.exception("could not build salvage history for %s", agent_id) return - desired = list(pre_run_items) + replay - if len(desired) <= len(pre_run_items): + if not replay: return try: - await replace_session_items(session, desired) + await recover_session_items(session, pre_run_items, replay) except Exception: logger.exception("salvaging crashed run history failed for %s", agent_id) diff --git a/strix/core/sessions.py b/strix/core/sessions.py index 9286b662..afc0fc1b 100644 --- a/strix/core/sessions.py +++ b/strix/core/sessions.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import json import logging import sqlite3 from contextlib import contextmanager @@ -102,11 +103,11 @@ async def _rewrite_session( session: Session, transform: Callable[[list[Any]], tuple[list[Any], bool]], ) -> bool: - """Read-modify-write a session under its write lock, restoring on failure.""" + """Rewrite atomically against SDK writes for SQLite; lock other backends.""" async with session_write_lock(session): + if isinstance(session, SQLiteSession): + return await asyncio.to_thread(_rewrite_sqlite_session, session, transform) items = await session.get_items() - if not items: - return False rebuilt, changed = transform(list(items)) if not changed: return False @@ -123,6 +124,37 @@ async def _rewrite_session( return True +def _rewrite_sqlite_session( + session: SQLiteSession, + transform: Callable[[list[Any]], tuple[list[Any], bool]], +) -> bool: + # The SDK writes without our asyncio lock. The SQLite write transaction + # covers read/compare/replace, including writers on another session object. + # Table names come from the SDK's configured schema, never tool/message input. + with session._locked_connection() as connection: # pyright: ignore[reportPrivateUsage] + try: + connection.execute("BEGIN IMMEDIATE") + rows = connection.execute( + f"SELECT message_data FROM {session.messages_table} " # noqa: S608 # nosec B608 + "WHERE session_id = ? ORDER BY id ASC", + (session.session_id,), + ).fetchall() + original = [json.loads(row[0]) for row in rows] + rebuilt, changed = transform(original) + if changed: + connection.execute( + f"DELETE FROM {session.messages_table} WHERE session_id = ?", # noqa: S608 # nosec B608 + (session.session_id,), + ) + session._insert_items(connection, cast("list[TResponseInputItem]", rebuilt)) # pyright: ignore[reportPrivateUsage] + connection.commit() + except BaseException: + connection.rollback() + raise + else: + return changed + + async def replace_session_items( session: Session, new_items: list[Any], @@ -135,25 +167,49 @@ async def replace_session_items( longer has that many items (a concurrent writer changed it), so a slow compaction summary can't clobber newer turns. """ - async with session_write_lock(session): - original = list(await session.get_items()) + + def _transform(original: list[Any]) -> tuple[list[Any], bool]: if expected_len is not None and len(original) != expected_len: logger.warning( "skipping session rewrite: expected %d items, found %d", expected_len, len(original), ) - return False - rebuilt = cast("list[TResponseInputItem]", new_items) - await session.clear_session() - try: - await session.add_items(rebuilt) - except Exception: - logger.exception("session rewrite failed; restoring original items") - await session.clear_session() - await session.add_items(original) - raise - return True + return original, False + return new_items, new_items != original + + return await _rewrite_session(session, _transform) + + +async def recover_session_items(session: Session, before: list[Any], replay: list[Any]) -> bool: + """Merge a full SDK replay once, retaining append-only incoming messages. + + Occurrences are compared by position, never globally deduplicated. If a + concurrent compaction or divergent SDK write changed the prefix, refuse the + rewrite rather than guess which events to delete or execute again. + """ + + def _transform(current: list[Any]) -> tuple[list[Any], bool]: + if current[: len(before)] != before or replay[: len(before)] != before: + raise ValueError("session history diverged from the pre-run snapshot") + common = 0 + for left, right in zip(current, replay, strict=False): + if left != right: + break + common += 1 + if common == len(replay): + return current, False + tail = current[common:] + if any( + not isinstance(item, dict) or cast("dict[str, Any]", item).get("role") != "user" + for item in tail + ): + raise ValueError("session contains divergent generated items; recovery is ambiguous") + if any(item in replay[common:] for item in tail): + raise ValueError("incoming message ownership is ambiguous; history retained") + return replay + tail, True + + return await _rewrite_session(session, _transform) async def strip_all_images_from_session(session: Session) -> bool: diff --git a/tests/test_execution.py b/tests/test_execution.py index 8fbf18ff..e1c8430d 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -793,6 +793,7 @@ async def test_salvage_stream_to_session_preserves_full_history(tmp_path: Any) - # A crash mid-run: the stream produced two turns the SDK never committed. stream = _SalvageStream( [ + *pre_run, {"role": "assistant", "content": "recon turn 1"}, {"role": "assistant", "content": "recon turn 2"}, ] diff --git a/tests/test_tool_history_recovery.py b/tests/test_tool_history_recovery.py new file mode 100644 index 00000000..6bd4a36a --- /dev/null +++ b/tests/test_tool_history_recovery.py @@ -0,0 +1,300 @@ +"""Crash recovery must preserve occurrences, not concatenate full SDK histories.""" + +from __future__ import annotations + +import asyncio +import json +import sqlite3 +import threading +from contextlib import contextmanager +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +import httpx +import pytest +from agents import Agent, Runner, function_tool +from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel +from agents.run import RunConfig +from openai import APIConnectionError, AsyncOpenAI + +from strix.config.models import _NonStreamingModel +from strix.core.execution import _salvage_stream_to_session +from strix.core.sessions import open_agent_session, replace_session_items + + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + + +def _call(call_id: str = "call_a") -> dict[str, Any]: + return {"type": "function_call", "call_id": call_id, "name": "noop", "arguments": "{}"} + + +def _output(call_id: str = "call_a") -> dict[str, Any]: + return {"type": "function_call_output", "call_id": call_id, "output": "ok"} + + +async def test_full_replay_without_new_events_does_not_duplicate(tmp_path: Path) -> None: + session = open_agent_session("a", tmp_path / "agents.db") + original = [{"role": "user", "content": "task"}, _call(), _output()] + try: + await session.add_items(original) + for _ in range(3): + await _salvage_stream_to_session( + session, original, SimpleNamespace(to_input_list=lambda: original), "a" + ) + assert await session.get_items() == original + finally: + session.close() + + +@pytest.mark.parametrize("persisted_count", [0, 1, 2]) +async def test_partial_sdk_persistence_and_late_message_are_preserved( + tmp_path: Path, persisted_count: int +) -> None: + session = open_agent_session("a", tmp_path / "agents.db") + original = [{"role": "user", "content": "task"}] + generated = [_call(), _output()] + late = {"role": "user", "content": "new instructions"} + replay = original + generated + try: + await session.add_items(original + generated[:persisted_count] + [late]) + stream = SimpleNamespace(to_input_list=lambda: replay) + await _salvage_stream_to_session(session, original, stream, "a") + await _salvage_stream_to_session(session, original, stream, "a") + assert await session.get_items() == [*replay, late] + finally: + session.close() + + +async def test_recovery_preserves_repeated_legitimate_messages(tmp_path: Path) -> None: + session = open_agent_session("a", tmp_path / "agents.db") + message = {"role": "user", "content": "continue"} + original = [message, message] + replay = [*original, _call(), _output()] + try: + await session.add_items([*original, message, message]) + await _salvage_stream_to_session( + session, original, SimpleNamespace(to_input_list=lambda: replay), "a" + ) + assert await session.get_items() == [*replay, message, message] + finally: + session.close() + + +async def test_divergent_compacted_session_is_not_overwritten( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + session = open_agent_session("a", tmp_path / "agents.db") + original = [{"role": "user", "content": "old"}] + current = [{"role": "user", "content": "compacted"}] + try: + await session.add_items(current) + await _salvage_stream_to_session( + session, original, SimpleNamespace(to_input_list=lambda: [*original, _call()]), "a" + ) + assert await session.get_items() == current + assert "salvaging crashed run history failed" in caplog.text + finally: + session.close() + + +async def test_parallel_sdk_writes_are_not_lost(tmp_path: Path) -> None: + db = tmp_path / "agents.db" + session = open_agent_session("a", db) + # A different SDK session instance bypasses Strix's per-instance asyncio lock. + writer = open_agent_session("a", db) + original = [{"role": "user", "content": "task"}] + replay = [*original, _call(), _output()] + late = [{"role": "user", "content": f"message {i}"} for i in range(20)] + try: + await session.add_items(original) + await asyncio.gather( + _salvage_stream_to_session( + session, original, SimpleNamespace(to_input_list=lambda: replay), "a" + ), + *(writer.add_items([item]) for item in late), + ) + stored = await session.get_items() + assert stored[:3] == replay + assert sorted(item["content"] for item in stored[3:]) == sorted( + item["content"] for item in late + ) + finally: + session.close() + writer.close() + + +async def test_real_sdk_failure_salvage_and_resume_never_repeat_a_tool(tmp_path: Path) -> None: + """Exercise SDK session persistence and Chat conversion, with no network.""" + requests: list[dict[str, Any]] = [] + executed: list[int] = [] + + @function_tool + def noop(n: int) -> str: + executed.append(n) + return f"result-{n}" + + def respond(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + requests.append(body) + turn = len(requests) + if turn == 3: + raise RuntimeError("synthetic stream failure") + message: dict[str, Any] = {"role": "assistant", "content": "finished"} + if turn < 3: + message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call_{turn}", + "type": "function", + "function": {"name": "noop", "arguments": json.dumps({"n": turn})}, + } + ], + } + return httpx.Response( + 200, + json={ + "id": f"completion-{turn}", + "object": "chat.completion", + "created": 0, + "model": "offline", + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": "tool_calls" if turn < 3 else "stop", + } + ], + }, + ) + + session = open_agent_session("a", tmp_path / "agents.db") + client = AsyncOpenAI( + api_key="offline", + base_url="http://offline.invalid/v1", + max_retries=0, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + model = _NonStreamingModel(OpenAIChatCompletionsModel(model="offline", openai_client=client)) + agent = Agent(name="offline", tools=[noop], model=model) + config = RunConfig(tracing_disabled=True) + try: + original = [{"role": "user", "content": "run two noops"}] + await session.add_items(original) + stream = Runner.run_streamed(agent, input=[], session=session, run_config=config) + with pytest.raises(APIConnectionError): + async for _ in stream.stream_events(): + pass + assert executed == [1, 2] + persisted = await session.get_items() + for _ in range(3): + await _salvage_stream_to_session(session, original, stream, "a") + assert await session.get_items() == persisted + result = await Runner.run(agent, input=[], session=session, run_config=config) + assert result.final_output == "finished" + assert executed == [1, 2] + sent = requests[-1]["messages"] + assert [m["tool_call_id"] for m in sent if m["role"] == "tool"] == ["call_1", "call_2"] + finally: + session.close() + await client.close() + + +async def test_sqlite_rewrite_rolls_back_partial_insertion( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + session = open_agent_session("a", tmp_path / "agents.db") + original = [{"role": "user", "content": "keep me"}] + insert = session._insert_items + + def fail_after_insert(connection: Any, items: Any) -> None: + insert(connection, items) + raise RuntimeError("write failed") + + try: + await session.add_items(original) + monkeypatch.setattr(session, "_insert_items", fail_after_insert) + with pytest.raises(RuntimeError, match="write failed"): + await replace_session_items(session, [_call()]) + assert await session.get_items() == original + finally: + session.close() + + +@pytest.mark.parametrize("later_turn", [False, True]) +async def test_recovery_never_overwrites_divergent_or_later_generated_items( + tmp_path: Path, later_turn: bool +) -> None: + session = open_agent_session("a", tmp_path / "agents.db") + before = [{"role": "user", "content": "task"}] + replay = [*before, _call(), _output()] + other = [_call("call_b"), _output("call_b")] + current = (replay if later_turn else before) + other + try: + await session.add_items(current) + await _salvage_stream_to_session( + session, before, SimpleNamespace(to_input_list=lambda: replay), "a" + ) + assert await session.get_items() == current + finally: + session.close() + + +async def test_recovery_transaction_blocks_writer_without_exposing_empty_history( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + db = tmp_path / "agents.db" + session, writer = [open_agent_session("a", db) for _ in range(2)] + entered, attempted, release = threading.Event(), threading.Event(), threading.Event() + original = [{"role": "user", "content": "task"}] + replay = [*original, _call(), _output()] + late = {"role": "user", "content": "late"} + insert, locked = session._insert_items, writer._locked_connection + + def hold_transaction(connection: Any, items: Any) -> None: + entered.set() + if not release.wait(5): + raise TimeoutError("test did not release transaction") + insert(connection, items) + + @contextmanager + def announce_writer() -> Iterator[sqlite3.Connection]: + attempted.set() + with locked() as connection: + yield connection + + tasks: list[asyncio.Task[Any]] = [] + try: + await session.add_items(original) + monkeypatch.setattr(session, "_insert_items", hold_transaction) + monkeypatch.setattr(writer, "_locked_connection", announce_writer) + tasks.append( + asyncio.create_task( + _salvage_stream_to_session( + session, original, SimpleNamespace(to_input_list=lambda: replay), "a" + ) + ) + ) + assert await asyncio.to_thread(entered.wait, 5) + tasks.append(asyncio.create_task(writer.add_items([late]))) + assert await asyncio.to_thread(attempted.wait, 5) + # DELETE is uncommitted: a separate reader still sees the original, + # while the SDK writer waits for the recovery transaction to commit. + with sqlite3.connect(db) as connection: + rows = connection.execute( + "SELECT message_data FROM agent_messages ORDER BY id" + ).fetchall() + assert [json.loads(row[0]) for row in rows] == original + assert not tasks[1].done() + release.set() + await asyncio.gather(*tasks) + assert await session.get_items() == [*replay, late] + finally: + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + for item in (session, writer): + item.close()