scriptoriumadmin/strix
1
0
Fork 0
mirror of https://github.com/usestrix/strix.git synced 2026-09-24 00:51:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(sessions): make failed-run history recovery idempotent

Browse source
This commit is contained in:
Genius229 2026-09-14 13:38:39 +04:00
parent 84f4108195
commit dbd078b3ec
4 changed files with 376 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
strix/core/execution.py
View file

@ -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)

88
strix/core/sessions.py
View file

@ -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:

1
tests/test_execution.py
View file

@ -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"},
]

300
tests/test_tool_history_recovery.py Normal file
View file

@ -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()