diff --git a/navi/core/agent.py b/navi/core/agent.py index 0cc7a33..75aa9ce 100644 --- a/navi/core/agent.py +++ b/navi/core/agent.py @@ -561,6 +561,12 @@ else: yield _ev + # Persist planning output (plan_ctx_msg + prompt_msg appended in + # planning.py) immediately — otherwise it sits in memory until the + # first save inside the tool loop and is lost on a mid-turn crash / + # CancelledError (server restart). + await self._sessions.save(session) + # Planning auto-populates the todo (planning.set_tasks) when Phase 3 # produced steps. Push the fresh state to the UI so the side panel # reflects the plan before the tool-calling loop starts. @@ -741,6 +747,9 @@ ) session.messages.append(assistant_msg) session.context.append(assistant_msg) + # Persist the assistant's tool-call decision before executing tools, so + # a crash / CancelledError mid-tool-execution doesn't lose it. + await self._sessions.save(session) tool_ctx = ToolContext( session_id=session_id, @@ -1111,6 +1120,10 @@ if image_msg: session.messages.append(image_msg) session.context.append(image_msg) + # Persist each tool result as it completes, so a crash / + # CancelledError (server restart) between tool calls cannot lose + # already-finished work — only the in-flight single tool is at risk. + await self._sessions.save(session) finally: if not tool_task.done(): tool_task.cancel() diff --git a/navi/core/compressor.py b/navi/core/compressor.py index f92b1a5..7f215c6 100644 --- a/navi/core/compressor.py +++ b/navi/core/compressor.py @@ -140,7 +140,13 @@ base_keep = keep_recent recent_turns = turns[-base_keep:] old_turns = turns[:-base_keep] - if adaptive: + # Adaptive swap rescues important old turns in multi-turn conversations. It is + # disabled in midturn mode (keep_recent_messages set): there the priority is to + # compress the in-flight turn predictably (see the split below), and the swap + # could otherwise move a filler in-flight turn into old_turns and summarize it + # whole — losing the current request. No test exercises the swap with + # keep_recent_messages set, so this guard is safe. + if adaptive and keep_recent_messages is None: # Identify important old turns that should not be lost. important_old = [t for t in old_turns if _turn_importance(t) > 0] # Identify filler turns in the recent window that can be swapped out. @@ -155,7 +161,29 @@ old_turns = [t for t in turns if id(t) not in kept_set] to_summarize = [m for turn in old_turns for m in turn] - to_keep = [m for turn in recent_turns for m in turn] + # Midturn mode (keep_recent_messages set): the in-flight (last) turn can be very + # long (many tool iterations). When there are already > keep_recent turns, the + # turn-based keep window holds the whole in-flight turn verbatim, so midturn + # compression barely shrinks the context. Split the in-flight turn too: keep its + # head (the user request) + the newest keep_recent_messages verbatim, summarize + # the middle — on top of compressing older turns. Mirrors the head/tail shape of + # partition_current_turn_messages. Without this, the in-flight turn stays verbatim + # whenever turns > keep_recent, which is the 150->141 shallow-compression symptom. + if keep_recent_messages is not None and recent_turns: + in_flight = recent_turns[-1] + if len(in_flight) > keep_recent_messages + 1: + head = [in_flight[0]] if in_flight and in_flight[0].role == "user" else [] + tail_start = max(len(head), len(in_flight) - keep_recent_messages) + to_summarize = to_summarize + in_flight[len(head):tail_start] + to_keep = ( + [m for turn in recent_turns[:-1] for m in turn] + + head + + in_flight[tail_start:] + ) + else: + to_keep = [m for turn in recent_turns for m in turn] + else: + to_keep = [m for turn in recent_turns for m in turn] return to_summarize, to_keep diff --git a/tests/unit/core/test_agent.py b/tests/unit/core/test_agent.py index e10661e..e442c9d 100644 --- a/tests/unit/core/test_agent.py +++ b/tests/unit/core/test_agent.py @@ -5,6 +5,8 @@ """ import asyncio +import copy +from datetime import datetime, timezone import pytest import pytest_asyncio @@ -18,7 +20,7 @@ SubagentComplete, ) from navi.core.registry import BackendRegistry, ProfileRegistry -from navi.core.session import InMemorySessionStore +from navi.core.session import InMemorySessionStore, Session from navi.exceptions import MaxIterationsReached, NothingToCompactError, SessionNotFound from navi.llm.base import LLMChunk, Message, ToolCallRequest from tests.conftest_factory import FakeLLMBackend, make_profile, make_registry_with_tools @@ -578,3 +580,150 @@ assert not _is_casual_message("напиши скрипт для бэкапа") assert not _is_casual_message("проверь почему не работает ssh") assert not _is_casual_message("what is the weather in Berlin tomorrow") + + +# ─── Persistence-on-crash tests (B1: incremental flush) ───────────────────── + + +class _SnapshotSessionStore(InMemorySessionStore): + """Session store that mimics a real DB boundary. + + ``save()`` deep-copies the session into a snapshot; ``get()`` returns the last + saved snapshot (a fresh copy), not the live object the agent is mutating. So + mutations that were never ``save()``d are invisible to ``get()`` — exactly like + ``PgSessionStore`` where messages with ``sequence_number < 0`` only persist on + ``save()``. This lets us assert that a crash mid-turn does not lose work that + was incrementally flushed (B1). + """ + + def __init__(self) -> None: + super().__init__() + self._snapshots: dict[str, Session] = {} + + async def save(self, session: Session) -> None: + session.last_active = datetime.now(timezone.utc) + # Mirror PgSessionStore: assign sequence numbers to new messages on save. + for m in session.messages: + if m.sequence_number < 0: + m.sequence_number = session.db_next_sequence + session.db_next_sequence += 1 + self._snapshots[session.id] = copy.deepcopy(session) + self._sessions[session.id] = session + + async def get(self, session_id: str) -> Session | None: + snap = self._snapshots.get(session_id) + if snap is not None: + return copy.deepcopy(snap) + return self._sessions.get(session_id) + + +class _CrashBackend(FakeLLMBackend): + """FakeLLMBackend whose ``stream_complete`` raises on the Nth call. + + Simulates a server crash / CancelledError arriving mid-turn, so we can assert + that tool work completed before the crash was already persisted (B1). + """ + + def __init__( + self, + responses: list[str], + tool_calls: list[list[ToolCallRequest] | None], + raise_on_call: int, + exc: BaseException = RuntimeError("simulated crash"), + ) -> None: + super().__init__(responses=responses, tool_calls=tool_calls) + self._raise_on_call = raise_on_call + self._exc = exc + self._sc_idx = 0 + + async def stream_complete(self, messages, tools=None, temperature=0.7, model=None, + think=None, **kwargs): + idx = self._sc_idx + self._sc_idx += 1 + if idx == self._raise_on_call: + raise self._exc + content = self._responses[idx] if idx < len(self._responses) else "" + tcalls = self._tool_calls[idx] if idx < len(self._tool_calls) else None + if content: + yield LLMChunk(delta=content) + yield LLMChunk( + finish_reason="tool_calls" if tcalls else "stop", + tool_calls=tcalls, + ) + + +def _build_persistence_agent(store: InMemorySessionStore, backend) -> Agent: + profiles = ProfileRegistry() + profile = make_profile("test") + # Disable planning so it neither appends plan messages nor consumes LLM calls — + # the test controls backend calls precisely via stream_complete. + profile.planning_phase1_enabled = False + profile.planning_phase2_enabled = False + profile.planning_phase3_enabled = False + profiles.register(profile) + tools = make_registry_with_tools() + backends = BackendRegistry() + backends.register("ollama", backend) + return Agent( + session_store=store, + profile_registry=profiles, + tool_registry=tools, + backend_registry=backends, + ) + + +class TestAgentPersistenceOnCrash: + @pytest.mark.asyncio + async def test_tool_results_persisted_on_crash_midturn(self): + """A crash on a later LLM call must not lose tool results already completed + in earlier iterations — B1 persists each tool result as it completes.""" + store = _SnapshotSessionStore() + backend = _CrashBackend( + responses=["", "unused"], + tool_calls=[ + [ToolCallRequest(id="1", name="test_tool", arguments={})], + None, + ], + raise_on_call=1, + ) + agent = _build_persistence_agent(store, backend) + sess = await agent._sessions.create(profile_id="test") + + with pytest.raises(RuntimeError, match="simulated crash"): + async for _ in agent.run_stream(sess.id, "do work"): + pass + + saved = await agent._sessions.get(sess.id) + roles = [m.role for m in saved.messages] + # Iteration 0's assistant tool-call decision + tool result were flushed by + # B1 before the crash on iteration 1. Without B1 only the user message + # (saved at turn start) would survive. + assert "tool" in roles, f"tool result missing from saved state: {roles}" + assert any( + m.role == "assistant" and m.tool_calls for m in saved.messages + ), f"assistant tool-call msg missing: {roles}" + + @pytest.mark.asyncio + async def test_cancel_flushes_completed_tool_work(self): + """CancelledError mid-turn (server restart/shutdown) must not lose tool + results already completed — B1 flushed them incrementally before the cancel.""" + store = _SnapshotSessionStore() + backend = _CrashBackend( + responses=["", "unused"], + tool_calls=[ + [ToolCallRequest(id="1", name="test_tool", arguments={})], + None, + ], + raise_on_call=1, + exc=asyncio.CancelledError(), + ) + agent = _build_persistence_agent(store, backend) + sess = await agent._sessions.create(profile_id="test") + + with pytest.raises(asyncio.CancelledError): + async for _ in agent.run_stream(sess.id, "do work"): + pass + + saved = await agent._sessions.get(sess.id) + roles = [m.role for m in saved.messages] + assert "tool" in roles, f"tool result missing after cancel: {roles}" diff --git a/tests/unit/core/test_compressor.py b/tests/unit/core/test_compressor.py index 904a3c6..4237608 100644 --- a/tests/unit/core/test_compressor.py +++ b/tests/unit/core/test_compressor.py @@ -95,6 +95,40 @@ assert recent[0].content == "build a model" assert len(recent) == 5 # original user message + 4 newest in-flight messages + def test_midturn_splits_in_flight_turn_when_many_turns(self): + """When turns > keep_recent AND keep_recent_messages is set, the in-flight + (last) turn is also split — head (user request) + newest K messages kept, + middle summarized — on top of compressing older turns. Previously the + in-flight turn stayed verbatim in the keep window (the 150->141 shallow + midturn-compression symptom).""" + msgs = [] + # 14 older short turns (2 msgs each) = 28 msgs + for i in range(14): + msgs.append(Message(role="user", content=f"u{i}")) + msgs.append(Message(role="assistant", content=f"a{i}")) + # in-flight 15th turn: user + 10 assistant/tool pairs = 21 msgs + msgs.append(Message(role="user", content="inflight request")) + for i in range(10): + msgs.append(Message(role="assistant", content=f"step {i}")) + msgs.append(Message(role="tool", content=f"res {i}", name="fs", tool_call_id=str(i))) + + old, recent = partition_messages(msgs, keep_recent=12, keep_recent_messages=4) + + recent_contents = {m.content for m in recent} + old_contents = {m.content for m in old} + + # The in-flight user request is kept verbatim (head of the in-flight turn). + assert "inflight request" in recent_contents + # The newest in-flight messages are kept verbatim (tail). + assert "step 9" in recent_contents and "res 9" in recent_contents + # The oldest older turn is summarized. + assert "u0" in old_contents and "a0" in old_contents + # The middle of the in-flight turn is summarized, not kept verbatim. + assert "step 0" in old_contents and "res 0" in old_contents + assert "step 0" not in recent_contents + # Nothing is dropped: old + recent partition the whole non-system set. + assert len(old) + len(recent) == len(msgs) + class TestFormatForSummary: def test_user_message(self): @@ -382,6 +416,44 @@ assert len(new_context) == 6 # system + summary + 2 recent turns @pytest.mark.asyncio + async def test_midturn_compresses_in_flight_turn_with_many_turns(self): + """End-to-end midturn: with > keep_recent turns and a long in-flight turn, + compression shrinks the in-flight turn too (head + newest K verbatim, middle + summarized) — not just the older turns. Previously the in-flight turn stayed + verbatim and compression was shallow.""" + backend = FakeLLMBackend(responses=["consolidated summary"]) + compressor = ContextCompressor() + context = [] + for i in range(14): + context.append(Message(role="user", content=f"u{i}")) + context.append(Message(role="assistant", content=f"a{i}")) + context.append(Message(role="user", content="inflight request")) + for i in range(10): + context.append(Message(role="assistant", content=f"step {i}")) + context.append(Message(role="tool", content=f"res {i}", name="fs", tool_call_id=str(i))) + + result = await compressor.compress_session( + context=context, + llm=backend, + model="test", + temperature=0.3, + keep_recent=12, + keep_recent_messages=4, + ) + assert result is not None + new_context, summary = result + assert summary == "consolidated summary" + # 1 summary + head(user) + 11 kept older turns * 2 msgs + 4 in-flight tail + # = 1 + 1 + 22 + 4 = 28. Without A1 the in-flight turn (21 msgs) would stay + # verbatim and new_context would be ~1 + 24 + 21 = 46. + assert len(new_context) == 28 + assert new_context[0].is_summary is True + contents = {m.content for m in new_context} + assert "inflight request" in contents # head kept + assert "step 9" in contents and "res 9" in contents # tail kept + assert "step 0" not in contents # in-flight middle summarized away + + @pytest.mark.asyncio async def test_compress_session_retry_on_failure(self): """First attempt fails, second succeeds with keep_recent + 4.""" import navi.core.compressor as compressor_module