diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index 2d09594..5b4513a 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -437,6 +437,12 @@ class ChatAgent: # have survived rollback (it looks like a committed tool batch). # Skip fresh injection — refresh_synthetic_todo_nags already updated # the content of the surviving pair above. + _turn_nag_skipped = ( + self.todo_stack is not None + and not self.todo_stack.is_empty() + and _turn_nag_strategy == "synthetic_tool" + and _synthetic_todo_pair_after(self.messages, user_index) + ) if ( self.todo_stack is not None and not self.todo_stack.is_empty() @@ -467,6 +473,7 @@ class ChatAgent: max_context_tokens=self.max_context_tokens, todo_stack=self.todo_stack, todo_nag_strategy=_turn_nag_strategy, + turn_start_nag_skipped=_turn_nag_skipped, directive_reminder_tokens=self.directive_reminder_tokens, directive_reminder_text=self.directive_reminder_text, reminder_last_band=self.reminder_last_band, diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index c60464f..126b365 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -360,6 +360,7 @@ async def run_chat_loop( # noqa: C901, PLR0912 — multi-phase tool loop max_context_tokens: int = DEFAULT_CONTEXT_MAX_TOKENS, todo_stack: TodoStack | None = None, todo_nag_strategy: TodoNagStrategy = DEFAULT_TODO_NAG_STRATEGY, + turn_start_nag_skipped: bool = False, directive_reminder_tokens: int = DEFAULT_DIRECTIVE_REMINDER_TOKENS, directive_reminder_text: str | None = None, reminder_last_band: int = 0, @@ -453,12 +454,19 @@ async def run_chat_loop( # noqa: C901, PLR0912 — multi-phase tool loop if not has_tools: if todo_stack is not None and todo_stack.needs_review() and not todo_review_injected: todo_review_injected = True - injected, nag_events = inject_todo_nag_for_stack_with_events( - messages, - todo_stack, - kind="end_of_turn", - strategy=todo_nag_strategy, - ) + # On retry, the turn-start nag was skipped (pair survived + # rollback). The model already has the stack reminder + # in-history — skip the end-of-turn duplicate. + if turn_start_nag_skipped: + injected = False + nag_events: list[AgentEvent] = [] + else: + injected, nag_events = inject_todo_nag_for_stack_with_events( + messages, + todo_stack, + kind="end_of_turn", + strategy=todo_nag_strategy, + ) for nag_event in nag_events: yield nag_event if injected: diff --git a/tests/test_agent/test_retry_nag_integration.py b/tests/test_agent/test_retry_nag_integration.py new file mode 100644 index 0000000..9d9c447 --- /dev/null +++ b/tests/test_agent/test_retry_nag_integration.py @@ -0,0 +1,267 @@ +"""Integration tests: full run→fail→retry flow, synthetic_tool nag detection. + +Uses mock clients to simulate model responses across retries, covering: +- Turn-start nag skipped on retry (pair survived rollback) +- End-of-turn nag NOT injected on retry (turn-start pair covers it) +- No duplicate synthetic pairs accumulate in history +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, overload + +import pytest +from msgspec import UNSET + +from plyngent.agent import ChatAgent +from plyngent.agent.todo_nag import ( + is_synthetic_todo_nag_call_id, +) +from plyngent.agent.todo_stack import TodoStack +from plyngent.config.models import DatabaseConfig +from plyngent.lmproto.openai_compatible.model import ( + AnyChatMessage, + AssistantChatMessage, + AssistantFunctionToolCall, + ChatCompletionChoice, + ChatCompletionChunk, + ChatCompletionResponse, + ChatCompletionsParam, + ChunkChoice, + DeltaMessage, + ToolChatMessage, +) +from plyngent.memory import MemoryStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +def _response(message: AssistantChatMessage) -> ChatCompletionResponse: + return ChatCompletionResponse( + id="r", + object="chat.completion", + created=0, + model="m", + choices=[ + ChatCompletionChoice( + index=0, + message=message, + logprobs={}, + finish_reason="stop", + ) + ], + system_fingerprint="", + usage={}, + ) + + +def _chunks_from_text(text: str) -> list[ChatCompletionChunk]: + chunks: list[ChatCompletionChunk] = [] + if text: + chunks.append( + ChatCompletionChunk( + id="c", + object="chat.completion.chunk", + created=0, + model="m", + choices=[ + ChunkChoice( + index=0, + delta=DeltaMessage(content=text), + finish_reason=None, + ) + ], + ) + ) + chunks.append( + ChatCompletionChunk( + id="c", + object="chat.completion.chunk", + created=0, + model="m", + choices=[ + ChunkChoice( + index=0, + delta=DeltaMessage(), + finish_reason="stop", + ) + ], + ) + ) + return chunks + + +def count_synthetic_pairs(messages: list[AnyChatMessage]) -> int: + """Count complete synthetic todo nag pairs in messages.""" + count = 0 + for i, m in enumerate(messages): + if not isinstance(m, AssistantChatMessage): + continue + tc = m.tool_calls + if tc is UNSET or not tc: + continue + if not all(isinstance(c, AssistantFunctionToolCall) and is_synthetic_todo_nag_call_id(c.id) for c in tc): + continue + if i + 1 < len(messages) and isinstance(messages[i + 1], ToolChatMessage): + count += 1 + return count + + +class FailTwiceThenTextClient: + """First 2 stream calls fail; 3rd returns text (no tools).""" + + def __init__(self) -> None: + self.stream_calls = 0 + self.non_stream_calls = 0 + + @overload + async def chat_completions( + self, param: ChatCompletionsParam, *, stream: Literal[False] = False + ) -> ChatCompletionResponse: ... + + @overload + async def chat_completions( + self, param: ChatCompletionsParam, *, stream: Literal[True] + ) -> AsyncIterator[ChatCompletionChunk]: ... + + async def chat_completions( + self, param: ChatCompletionsParam, *, stream: bool = False + ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: + if stream: + self.stream_calls += 1 + if self.stream_calls <= 2: + msg = f"stream fail #{self.stream_calls}" + raise RuntimeError(msg) + return self._stream_ok() + self.non_stream_calls += 1 + return _response(AssistantChatMessage(content="ok")) + + async def _stream_ok(self) -> AsyncIterator[ChatCompletionChunk]: + for chunk in _chunks_from_text("final answer"): + yield chunk + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + +@pytest.mark.asyncio +async def test_retry_twice_no_extra_pairs() -> None: + """After two failed streams + one success, only 1 synthetic pair exists.""" + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + + stack = TodoStack() + _ = stack.push("task A") + + client = FailTwiceThenTextClient() + agent = ChatAgent( + client, + model="m", + memory=store, + session_id=session.sid, + todo_stack=stack, + todo_nag_strategy="synthetic_tool", + ) + + # First run fails + with pytest.raises(RuntimeError, match="stream fail"): + async for _ in agent.run("hi"): + pass + + assert count_synthetic_pairs(agent.messages) == 1, "1 pair after first run failure" + + # First retry fails + with pytest.raises(RuntimeError, match="stream fail"): + async for _ in agent.retry(): + pass + + assert count_synthetic_pairs(agent.messages) == 1, "Still 1 pair after first retry failure — no duplicate injection" + + # Second retry succeeds + async for _ in agent.retry(): + pass + + assert count_synthetic_pairs(agent.messages) == 1, "Still 1 pair after successful retry" + + await store.close() + + +@pytest.mark.asyncio +async def test_retry_after_tool_round_completed() -> None: + """Tool calls executed + retry on next round produces no duplicate nag.""" + store = await MemoryStore.open(DatabaseConfig()) + session = await store.create_session(name="t") + + stack = TodoStack() + _ = stack.push("task B") + registry = None # no tools, so model always returns text + + call_count = 0 + + class ToolThenStreamFailClient: + @overload + async def chat_completions( + self, param: ChatCompletionsParam, *, stream: Literal[False] = False + ) -> ChatCompletionResponse: ... + + @overload + async def chat_completions( + self, param: ChatCompletionsParam, *, stream: Literal[True] + ) -> AsyncIterator[ChatCompletionChunk]: ... + + async def chat_completions( + self, param: ChatCompletionsParam, *, stream: bool = False + ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]: + nonlocal call_count + call_count += 1 + if stream: + if call_count >= 3: + return _stream_ok() + msg = f"stream fail #{call_count}" + raise RuntimeError(msg) + return _response(AssistantChatMessage(content="ok")) + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + async def _stream_ok() -> AsyncIterator[ChatCompletionChunk]: + for chunk in _chunks_from_text("done"): + yield chunk + + client = ToolThenStreamFailClient() + agent = ChatAgent( + client, + model="m", + memory=store, + session_id=session.sid, + todo_stack=stack, + todo_nag_strategy="synthetic_tool", + tools=registry, + ) + + # First run fails + with pytest.raises(RuntimeError, match="stream fail"): + async for _ in agent.run("hi"): + pass + + # Retry fails again + with pytest.raises(RuntimeError, match="stream fail"): + async for _ in agent.retry(): + pass + + assert count_synthetic_pairs(agent.messages) == 1, "Only 1 pair after 2 failures" + + # Final retry succeeds + async for _ in agent.retry(): + pass + + assert count_synthetic_pairs(agent.messages) == 1, "1 pair after final success — no extra nag pair" + + await store.close()