from __future__ import annotations import asyncio from typing import TYPE_CHECKING, Literal, cast, overload import pytest from msgspec import UNSET from plyngent.agent import ( AssistantMessageEvent, ChatAgent, MaxRoundsEvent, ReasoningDeltaEvent, TextDeltaEvent, ToolCallEvent, ToolRegistry, ToolResultEvent, TruncatedEvent, UsageEvent, run_chat_loop, tool, ) from plyngent.agent.notices import TRUNCATED_NOTICE_KIND, notice_of_kind from plyngent.config.models import DatabaseConfig from plyngent.lmproto.openai_compatible.model import ( AnyChatMessage, AssistantChatMessage, AssistantFunctionTool, AssistantFunctionToolCall, ChatCompletionChoice, ChatCompletionChunk, ChatCompletionResponse, ChatCompletionsParam, ChunkChoice, DeltaMessage, FinishReason, StreamFunctionDelta, StreamToolCallDelta, ToolChatMessage, UserChatMessage, ) from plyngent.memory import MemoryStore if TYPE_CHECKING: from collections.abc import AsyncIterator def _chunks_from_response(response: ChatCompletionResponse) -> list[ChatCompletionChunk]: """Turn a full response into stream chunks (library-style stream=True path).""" message = response.choices[0].message chunks: list[ChatCompletionChunk] = [] if isinstance(message.content, str) and message.content: chunks.append( ChatCompletionChunk( id=response.id, object="chat.completion.chunk", created=response.created, model=response.model, choices=[ ChunkChoice( index=0, delta=DeltaMessage(content=message.content), finish_reason=None, ) ], ) ) # Emit terminal finish_reason when content was streamed without one. fr = response.choices[0].finish_reason if ( fr and (isinstance(message.content, str) and message.content) and (message.tool_calls is UNSET or not message.tool_calls) ): chunks.append( ChatCompletionChunk( id=response.id, object="chat.completion.chunk", created=response.created, model=response.model, choices=[ ChunkChoice( index=0, delta=DeltaMessage(), finish_reason=fr, ) ], ) ) tool_calls = message.tool_calls if tool_calls is not UNSET and tool_calls: deltas: list[StreamToolCallDelta] = [] for i, call in enumerate(tool_calls): if isinstance(call, AssistantFunctionToolCall): deltas.append( StreamToolCallDelta( index=i, id=call.id, type="function", function=StreamFunctionDelta( name=call.function.name, arguments=call.function.arguments, ), ) ) if deltas: chunks.append( ChatCompletionChunk( id=response.id, object="chat.completion.chunk", created=response.created, model=response.model, choices=[ ChunkChoice( index=0, delta=DeltaMessage(tool_calls=deltas), finish_reason="tool_calls", ) ], ) ) if not chunks: chunks.append( ChatCompletionChunk( id=response.id, object="chat.completion.chunk", created=response.created, model=response.model, choices=[ ChunkChoice( index=0, delta=DeltaMessage(), finish_reason=response.choices[0].finish_reason or "stop", ) ], ) ) return chunks class ScriptedClient: """Scripted chat completions; supports stream=True via chunked responses.""" _responses: list[ChatCompletionResponse] calls: list[ChatCompletionsParam] def __init__(self, responses: list[ChatCompletionResponse]) -> None: self._responses = list(responses) self.calls = [] @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]: self.calls.append(param) if not self._responses: msg = "no more scripted responses" raise RuntimeError(msg) response = self._responses.pop(0) if stream: return self._as_stream(response) return response async def _as_stream(self, response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]: for chunk in _chunks_from_response(response): yield chunk def _response( message: AssistantChatMessage, *, usage: dict[str, object] | None = None, finish_reason: FinishReason | None = "stop", ) -> ChatCompletionResponse: return ChatCompletionResponse( id="1", object="chat.completion", created=0, model="test", choices=[ ChatCompletionChoice( index=0, message=message, logprobs={}, finish_reason=finish_reason, ) ], system_fingerprint="", usage=usage if usage is not None else {}, ) async def test_run_chat_loop_text_only() -> None: client = ScriptedClient( [ _response(AssistantChatMessage(content="hello")), ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] assert isinstance(events[0], TextDeltaEvent) assert events[0].content == "hello" assert isinstance(events[1], AssistantMessageEvent) assert len(messages) == 2 assert len(client.calls) == 1 async def test_run_chat_loop_carries_reasoning_effort() -> None: client = ScriptedClient([_response(AssistantChatMessage(content="hello"))]) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] _ = [e async for e in run_chat_loop(client, messages, model="m", stream=False, reasoning_effort="high")] assert client.calls[0].reasoning_effort == "high" async def test_run_chat_loop_omits_unset_reasoning_effort() -> None: client = ScriptedClient([_response(AssistantChatMessage(content="hello"))]) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] _ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] assert client.calls[0].reasoning_effort is UNSET async def test_non_stream_emits_usage() -> None: client = ScriptedClient( [ _response( AssistantChatMessage(content="hi"), usage={"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11}, ), ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="x")] events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] usages = [e for e in events if isinstance(e, UsageEvent)] assert len(usages) == 1 assert usages[0].usage.prompt_tokens == 9 assert usages[0].usage.completion_tokens == 2 assert usages[0].usage.total_tokens == 11 assert usages[0].usage.source == "api" async def test_non_stream_estimates_usage_when_missing() -> None: client = ScriptedClient([_response(AssistantChatMessage(content="hello"))]) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] usages = [e for e in events if isinstance(e, UsageEvent)] assert len(usages) == 1 assert usages[0].usage.source == "estimate" assert usages[0].usage.total_tokens > 0 async def test_chat_agent_accumulates_session_usage() -> None: client = ScriptedClient( [ _response( AssistantChatMessage(content="a"), usage={"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, ), _response( AssistantChatMessage(content="b"), usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, ), ] ) agent = ChatAgent(client, model="m", stream=False) _ = [e async for e in agent.run("one")] assert agent.last_turn_usage.total_tokens == 6 assert agent.last_request_usage.total_tokens == 6 assert agent.last_turn_rounds == 1 assert agent.session_usage.total_tokens == 6 _ = [e async for e in agent.run("two")] assert agent.last_turn_usage.total_tokens == 10 assert agent.last_request_usage.total_tokens == 10 assert agent.session_usage.total_tokens == 16 async def test_chat_agent_turn_usage_sums_tool_rounds() -> None: """Multi-round tool loop: turn usage is billing sum; last_request is final call.""" @tool(register=False) def ping() -> str: return "pong" registry = ToolRegistry([ping]) client = ScriptedClient( [ _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="1", function=AssistantFunctionTool(name="ping", arguments="{}"), ) ], ), usage={"prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105}, ), _response( AssistantChatMessage(content="done"), usage={"prompt_tokens": 200, "completion_tokens": 3, "total_tokens": 203}, ), ] ) agent = ChatAgent(client, model="m", tools=registry, stream=False) _ = [e async for e in agent.run("go")] assert agent.last_turn_rounds == 2 assert agent.last_request_usage.prompt_tokens == 200 assert agent.last_turn_usage.prompt_tokens == 300 assert agent.last_turn_usage.total_tokens == 308 # Context size is last request prompt, not billed sum assert agent.context_tokens == 200 assert agent.context_tokens_source == "api" async def test_context_tokens_falls_back_to_message_estimate() -> None: agent = ChatAgent( ScriptedClient([]), model="m", stream=False, messages=[UserChatMessage(content="12345678")], ) assert agent.last_request_usage.is_zero() assert agent.context_tokens_source == "estimate" assert agent.context_tokens >= 1 async def test_stream_yields_deltas_incrementally() -> None: """Text deltas are yielded as chunks arrive, not only after the full stream.""" class ChunkClient: @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]: del param if not stream: return _response(AssistantChatMessage(content="ab")) async def chunks() -> AsyncIterator[ChatCompletionChunk]: for part in ("a", "b"): yield ChatCompletionChunk( id="1", object="chat.completion.chunk", created=0, model="t", choices=[ ChunkChoice( index=0, delta=DeltaMessage(content=part), finish_reason=None, ) ], ) return chunks() messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(ChunkClient(), messages, model="m", stream=True)] deltas = [e for e in events if isinstance(e, TextDeltaEvent)] assert [d.content for d in deltas] == ["a", "b"] assert any(isinstance(e, AssistantMessageEvent) for e in events) assert isinstance(messages[-1], AssistantChatMessage) assert messages[-1].content == "ab" async def test_stream_yields_reasoning_deltas() -> None: class ReasoningClient: @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]: del param if not stream: return _response(AssistantChatMessage(content="ans", reasoning_content="think")) async def chunks() -> AsyncIterator[ChatCompletionChunk]: for part in ("th", "ink"): yield ChatCompletionChunk( id="1", object="chat.completion.chunk", created=0, model="t", choices=[ ChunkChoice( index=0, delta=DeltaMessage(reasoning_content=part), finish_reason=None, ) ], ) yield ChatCompletionChunk( id="1", object="chat.completion.chunk", created=0, model="t", choices=[ ChunkChoice( index=0, delta=DeltaMessage(content="ans"), finish_reason="stop", ) ], ) return chunks() messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(ReasoningClient(), messages, model="m", stream=True)] reasoning = [e for e in events if isinstance(e, ReasoningDeltaEvent)] assert [r.content for r in reasoning] == ["th", "ink"] deltas = [e for e in events if isinstance(e, TextDeltaEvent)] assert [d.content for d in deltas] == ["ans"] assert isinstance(messages[-1], AssistantChatMessage) assert messages[-1].content == "ans" assert messages[-1].reasoning_content == "think" async def test_non_stream_yields_reasoning() -> None: client = ScriptedClient([_response(AssistantChatMessage(content="ok", reasoning_content="plan"))]) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] reasoning = [e for e in events if isinstance(e, ReasoningDeltaEvent)] assert len(reasoning) == 1 assert reasoning[0].content == "plan" assert isinstance(messages[-1], AssistantChatMessage) assert messages[-1].reasoning_content == "plan" async def test_run_chat_loop_with_tools() -> None: @tool(register=False) def add(a: int, b: int) -> int: return a + b registry = ToolRegistry([add]) client = ScriptedClient( [ _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c1", function=AssistantFunctionTool(name="add", arguments='{"a": 1, "b": 2}'), ) ], ) ), _response(AssistantChatMessage(content="3")), ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="1+2")] events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry)] types = [type(e) for e in events] assert ToolCallEvent in types assert ToolResultEvent in types assert any(isinstance(e, TextDeltaEvent) and e.content == "3" for e in events) assert len(client.calls) == 2 # second call includes tool result message assert any(getattr(m, "tool_call_id", None) == "c1" for m in client.calls[1].messages) async def test_max_rounds() -> None: @tool(register=False) def ping() -> str: return "pong" registry = ToolRegistry([ping]) forever = _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c", function=AssistantFunctionTool(name="ping", arguments="{}"), ) ], ) ) client = ScriptedClient([forever, forever, forever]) messages: list[AnyChatMessage] = [UserChatMessage(content="x")] events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry, max_rounds=2)] assert any(isinstance(e, MaxRoundsEvent) and e.rounds == 2 and not e.continued for e in events) assert len(client.calls) == 2 async def test_max_rounds_continue_hook() -> None: @tool(register=False) def ping() -> str: return "pong" registry = ToolRegistry([ping]) forever = _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c", function=AssistantFunctionTool(name="ping", arguments="{}"), ) ], ) ) final = _response(AssistantChatMessage(content="done")) client = ScriptedClient([forever, forever, final]) messages: list[AnyChatMessage] = [UserChatMessage(content="x")] asks: list[str] = [] def on_limit(reason: str) -> bool: asks.append(reason) return len(asks) == 1 events = [ e async for e in run_chat_loop( client, messages, model="m", tools=registry, max_rounds=2, on_limit=on_limit, ) ] assert len(asks) == 1 assert any(isinstance(e, MaxRoundsEvent) and e.continued for e in events) assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events) assert len(client.calls) == 3 async def test_max_rounds_async_continue_hook() -> None: @tool(register=False) def ping() -> str: return "pong" registry = ToolRegistry([ping]) forever = _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c", function=AssistantFunctionTool(name="ping", arguments="{}"), ) ], ) ) final = _response(AssistantChatMessage(content="done")) client = ScriptedClient([forever, forever, final]) messages: list[AnyChatMessage] = [UserChatMessage(content="x")] asks: list[str] = [] async def on_limit(reason: str) -> bool: asks.append(reason) return True events = [ e async for e in run_chat_loop( client, messages, model="m", tools=registry, max_rounds=2, on_limit=on_limit, ) ] assert len(asks) == 1 assert any(isinstance(e, MaxRoundsEvent) and e.continued for e in events) assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events) async def test_default_max_rounds_is_generous() -> None: from plyngent.agent.loop import DEFAULT_MAX_ROUNDS assert DEFAULT_MAX_ROUNDS >= 16 async def test_chat_agent_memory_roundtrip() -> None: store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") client = ScriptedClient([_response(AssistantChatMessage(content="yo"))]) agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) events = [e async for e in agent.run("hi")] assert any(isinstance(e, TextDeltaEvent) and e.content == "yo" for e in events) loaded = await store.list_messages(session.sid) assert len(loaded) == 2 assert isinstance(loaded[0], UserChatMessage) assert loaded[0].content == "hi" agent2 = ChatAgent(client, model="m", memory=store, session_id=session.sid) await agent2.load_history() assert len(agent2.messages) == 2 await store.close() async def test_chat_agent_failed_turn_keeps_user_in_db() -> None: store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") class BoomClient: @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]: del param, stream msg = "network down" raise RuntimeError(msg) agent = ChatAgent(BoomClient(), model="m", memory=store, session_id=session.sid) with pytest.raises(RuntimeError, match="network down"): _ = [e async for e in agent.run("hello")] assert len(agent.messages) == 1 assert isinstance(agent.messages[0], UserChatMessage) assert agent.pending_retry_text == "hello" loaded = await store.list_messages(session.sid) assert len(loaded) == 1 assert isinstance(loaded[0], UserChatMessage) assert loaded[0].content == "hello" await store.close() async def test_rollback_keeps_executed_tool_result_in_store() -> None: """Rolling back a cut-short batch keeps its executed result, in memory and on disk.""" store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") ran: list[str] = [] @tool(register=False) def first() -> str: ran.append("first") return "first ran" client = ScriptedClient([_response(AssistantChatMessage(content="finished"))]) agent = ChatAgent( client, model="m", tools=ToolRegistry([first]), memory=store, session_id=session.sid, stream=False, ) user = UserChatMessage(content="do both") await store.append_message(session.sid, user) agent.replace_messages( [ user, AssistantChatMessage( content="running both", tool_calls=[ AssistantFunctionToolCall( id="call_1", function=AssistantFunctionTool(name="first", arguments="{}") ), AssistantFunctionToolCall( id="call_2", function=AssistantFunctionTool(name="other", arguments="{}") ), ], ), ToolChatMessage(tool_call_id="call_1", content="first ran"), ], persist_from=1, ) await agent._rollback_uncommitted(0) # the failure path itself kept = [message for message in agent.messages if isinstance(message, ToolChatMessage)] assert [message.content for message in kept] == ["first ran"] assert agent.pending_retry_text == "do both" # What the model continues from is on disk too, or a restart re-runs it. assert agent.persist_from == len(agent.messages) stored = await store.list_messages(session.sid) assert [type(message).__name__ for message in stored] == [ "UserChatMessage", "AssistantChatMessage", "ToolChatMessage", ] record = stored[1] assert isinstance(record, AssistantChatMessage) assert record.content is UNSET calls = record.tool_calls assert calls is not UNSET assert [call.id for call in calls] == ["call_1"] # The retry answers from the stored result instead of running the tool again. events = [event async for event in agent.retry()] assert ran == [] assert any(isinstance(event, TextDeltaEvent) and event.content == "finished" for event in events) sent = client.calls[-1].messages assert any(isinstance(message, ToolChatMessage) and message.content == "first ran" for message in sent) await store.close() async def test_cancelled_batch_records_every_computed_result() -> None: """Leaving mid-batch keeps the results that were already computed.""" ran: list[str] = [] @tool(register=False) def first() -> str: ran.append("first") return "first ran" @tool(register=False) def second() -> str: ran.append("second") return "second ran" client = ScriptedClient( [ _response( AssistantChatMessage( content="running both", tool_calls=[ AssistantFunctionToolCall( id="call_1", function=AssistantFunctionTool(name="first", arguments="{}") ), AssistantFunctionToolCall( id="call_2", function=AssistantFunctionTool(name="second", arguments="{}") ), ], ) ) ] ) agent = ChatAgent(client, model="m", tools=ToolRegistry([first, second]), stream=False) stream = agent.run("do both") async for event in stream: if isinstance(event, ToolResultEvent): break await stream.aclose() results = [message for message in agent.messages if isinstance(message, ToolChatMessage)] assert [message.content for message in results] == ["first ran", "second ran"] assert sorted(ran) == ["first", "second"] async def test_chat_agent_retry_after_failure() -> None: store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") class FlakyClient: calls: int def __init__(self) -> None: self.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]: del param self.calls += 1 if self.calls == 1: msg = "temporary" raise RuntimeError(msg) response = _response(AssistantChatMessage(content="recovered")) if stream: async def as_stream() -> AsyncIterator[ChatCompletionChunk]: for chunk in _chunks_from_response(response): yield chunk return as_stream() return response client = FlakyClient() agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) with pytest.raises(RuntimeError, match="temporary"): _ = [e async for e in agent.run("ping")] assert agent.pending_retry_text == "ping" # User already in DB after first attempt assert len(await store.list_messages(session.sid)) == 1 events = [e async for e in agent.retry()] assert any(isinstance(e, TextDeltaEvent) and e.content == "recovered" for e in events) assert agent.pending_retry_text is None loaded = await store.list_messages(session.sid) assert len(loaded) == 2 assert isinstance(loaded[0], UserChatMessage) assert loaded[0].content == "ping" # Single user message (no duplicate on retry) assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1 await store.close() async def test_retry_after_resume_orphan_user() -> None: store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") _ = await store.append_message(session.sid, UserChatMessage(content="left hanging")) client = ScriptedClient([_response(AssistantChatMessage(content="ok now"))]) agent = ChatAgent(client, model="m", memory=store, session_id=session.sid) await agent.load_history() assert agent.pending_retry_text == "left hanging" events = [e async for e in agent.retry()] assert any(isinstance(e, TextDeltaEvent) and e.content == "ok now" for e in events) loaded = await store.list_messages(session.sid) assert len(loaded) == 2 assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1 await store.close() async def test_retry_keeps_committed_tools_after_second_round_fails() -> None: """Tools that already ran stay in history; retry continues without re-calling them.""" store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") calls: list[str] = [] @tool(register=False) def side_effect() -> str: calls.append("ran") return "done-once" registry = ToolRegistry([side_effect]) class ToolThenBoom: n: int def __init__(self) -> None: self.n = 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]: del stream self.n += 1 if self.n == 1: return _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c1", function=AssistantFunctionTool( name="side_effect", arguments="{}", ), ) ], ) ) if self.n == 2: msg = "api down after tools" raise RuntimeError(msg) # Retry continues: third call is next model round with tools in history. assert any(getattr(m, "tool_call_id", None) == "c1" for m in param.messages) return _response(AssistantChatMessage(content="finished")) client = ToolThenBoom() agent = ChatAgent( client, model="m", tools=registry, memory=store, session_id=session.sid, stream=False, ) with pytest.raises(RuntimeError, match="api down after tools"): _ = [e async for e in agent.run("do it")] assert calls == ["ran"] assert agent.pending_retry_text == "do it" # User + assistant(tool_calls) + tool result kept; incomplete second assistant dropped. assert isinstance(agent.messages[-1], ToolChatMessage) loaded = await store.list_messages(session.sid) assert len(loaded) == 3 # user, assistant, tool — committed before boom events = [e async for e in agent.retry()] assert any(isinstance(e, TextDeltaEvent) and e.content == "finished" for e in events) assert calls == ["ran"] # tool not re-executed assert agent.pending_retry_text is None await store.close() async def test_retry_keeps_lineno_reads_across_attempt_tasks(tmp_path: object) -> None: """A numbered read before a failed round still enables edits after retry. Each turn attempt runs in its own task (the host wraps turns in ``asyncio.create_task``), so the retry attempt's context does not carry the first attempt's ContextVar binding. Without ``retry`` re-binding the lineno tracker, ``read_file(with_lineno=true)`` silently records nothing and ``edit_lineno`` rejects every line edit — even after re-reading the file as its own error message instructs. """ from pathlib import Path from plyngent.tools.context import InstanceState from plyngent.tools.file import edit_lineno, read_file from plyngent.tools.workspace import ( clear_workspace_allowlist, clear_workspace_root, set_workspace_root, ) root = cast("Path", tmp_path).resolve() (root / "f.txt").write_text("alpha\nbravo\ncharlie\n", encoding="utf-8") instance = InstanceState(workspace_root=root) from plyngent.tools.context import bind_instance with bind_instance(instance): clear_workspace_allowlist() _ = set_workspace_root(root) try: registry = ToolRegistry([read_file, edit_lineno], instance_state=instance) class ReadThenBoom: """Round 1 reads with line numbers; round 2's model call fails.""" n: int def __init__(self) -> None: self.n = 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]: del stream self.n += 1 if self.n == 1: return _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c1", function=AssistantFunctionTool( name="read_file", arguments='{"path": "f.txt", "with_lineno": true}', ), ) ], ) ) if self.n == 2: msg = "api down after read" raise RuntimeError(msg) if self.n == 3: edit_args = '{"path": "f.txt", "start_line": 2, "end_line": 2, "new_content": "BRAVO\\n"}' return _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="c2", function=AssistantFunctionTool( name="edit_lineno", arguments=edit_args, ), ) ], ) ) return _response(AssistantChatMessage(content="done")) client = ReadThenBoom() agent = ChatAgent( client, model="m", tools=registry, stream=False, ) async def run_attempt() -> list[object]: return [e async for e in agent.run("edit line 2")] # First attempt in its own task (mirrors run_cancellable): commits # the read, then the next model round fails. with pytest.raises(RuntimeError, match="api down after read"): await asyncio.create_task(run_attempt()) assert isinstance(agent.messages[-1], ToolChatMessage) # Retry in a *fresh* task: its context copy has no tracker bound. async def run_retry() -> list[object]: return [e async for e in agent.retry()] events = await asyncio.create_task(run_retry()) results = [e for e in events if isinstance(e, ToolResultEvent)] assert results, "retry must run the edit_lineno call" assert "replaced lines 2-2" in results[0].message.content assert (root / "f.txt").read_text(encoding="utf-8") == "alpha\nBRAVO\ncharlie\n" assert agent.pending_retry_text is None finally: clear_workspace_allowlist() clear_workspace_root() async def test_chat_agent_system_prompt_prepended() -> None: from plyngent.lmproto.openai_compatible.model import SystemChatMessage client = ScriptedClient([_response(AssistantChatMessage(content="ok"))]) agent = ChatAgent(client, model="m", system_prompt="Be brief.", stream=False) _ = [e async for e in agent.run("hi")] assert isinstance(agent.messages[0], SystemChatMessage) assert agent.messages[0].content == "Be brief." assert isinstance(client.calls[0].messages[0], SystemChatMessage) async def test_tool_result_char_budget() -> None: @tool(register=False) def big() -> str: return "x" * 2000 registry = ToolRegistry([big]) client = ScriptedClient( [ _response( AssistantChatMessage( content="", tool_calls=[ AssistantFunctionToolCall( id="1", function=AssistantFunctionTool(name="big", arguments="{}"), ) ], ) ), _response(AssistantChatMessage(content="done")), ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="go")] _ = [ e async for e in run_chat_loop( client, messages, model="m", tools=registry, stream=False, max_tool_result_chars=200, parallel_tools=False, ) ] from plyngent.lmproto.openai_compatible.model import ToolChatMessage # The stored result stays whole; only the wire copy is capped, and it names # the document that reads the rest. stored = [m for m in messages if isinstance(m, ToolChatMessage)] assert len(stored) == 1 assert stored[0].content == "x" * 2000 wire = next(m for m in client.calls[-1].messages if isinstance(m, ToolChatMessage)) assert wire.content.startswith("x") assert "document_id=1" in wire.content assert len(wire.content) < 2000 async def test_empty_completion_raises() -> None: client = ScriptedClient([_response(AssistantChatMessage(content=None))]) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] with pytest.raises(RuntimeError, match="empty model completion"): _ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] async def test_length_finish_keeps_the_partial_answer_and_continues() -> None: client = ScriptedClient( [ _response(AssistantChatMessage(content="partial"), finish_reason="length"), _response(AssistantChatMessage(content="rest")), ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] truncated = [e for e in events if isinstance(e, TruncatedEvent)] assert len(truncated) == 1 assert truncated[0].kept_chars == len("partial") assert truncated[0].continued # The partial answer stays in history under a truncation notice, and the # continuation round is sent that history — the model continues its answer # rather than writing it again from the start. assert isinstance(messages[1], AssistantChatMessage) assert messages[1].content == "partial" assert notice_of_kind(messages[2], TRUNCATED_NOTICE_KIND) assert isinstance(messages[3], AssistantChatMessage) assert messages[3].content == "rest" assert any(notice_of_kind(m, TRUNCATED_NOTICE_KIND) for m in client.calls[1].messages) async def test_length_finish_stops_once_the_continuation_budget_is_spent() -> None: client = ScriptedClient([_response(AssistantChatMessage(content="partial"), finish_reason="length")] * 2) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [ e async for e in run_chat_loop( client, messages, model="m", stream=False, max_truncation_continues=1, ) ] assert [e.continued for e in events if isinstance(e, TruncatedEvent)] == [True, False] assert len(client.calls) == 2 # Nothing raised: the turn ends with the partial answer kept, so `/retry` # continues it instead of regenerating. assert isinstance(messages[-2], AssistantChatMessage) assert messages[-2].content == "partial" assert notice_of_kind(messages[-1], TRUNCATED_NOTICE_KIND) async def test_length_finish_drops_the_rounds_tool_calls() -> None: ran: list[str] = [] @tool(name="boom", register=False) async def _boom(argument: str) -> str: ran.append(argument) return "ran" call = AssistantFunctionToolCall( id="c1", function=AssistantFunctionTool(name="boom", arguments='{"argument": "x"'), ) client = ScriptedClient( [ _response( AssistantChatMessage(content="partial", tool_calls=[call]), finish_reason="length", ), _response(AssistantChatMessage(content="done")), ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(client, messages, model="m", tools=ToolRegistry([_boom]), stream=False)] # A call cut off mid-write is not replayed: the model re-issues it after # reading the notice. assert ran == [] assert not [e for e in events if isinstance(e, (ToolCallEvent, ToolResultEvent))] kept = messages[1] assert isinstance(kept, AssistantChatMessage) assert kept.content == "partial" assert kept.tool_calls is UNSET async def test_length_finish_without_content_still_raises() -> None: client = ScriptedClient([_response(AssistantChatMessage(content=None), finish_reason="length")]) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] with pytest.raises(RuntimeError, match="output limit reached with nothing written"): _ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] async def test_truncated_turn_is_retryable() -> None: store = await MemoryStore.open(DatabaseConfig()) session = await store.create_session(name="t") client = ScriptedClient( [ _response(AssistantChatMessage(content="partial"), finish_reason="length"), _response(AssistantChatMessage(content="rest")), ] ) agent = ChatAgent(client, model="m", memory=store, session_id=session.sid, max_truncation_continues=0, stream=False) _ = [e async for e in agent.run("hi")] # Kept, not rolled back: the turn is unfinished, so retry continues it. assert agent.pending_retry_text == "hi" _ = [e async for e in agent.retry()] assert any(notice_of_kind(m, TRUNCATED_NOTICE_KIND) for m in client.calls[1].messages) assert agent.pending_retry_text is None await store.close() async def test_content_filter_finish_raises() -> None: client = ScriptedClient( [ _response( AssistantChatMessage(content=""), finish_reason="content_filter", ) ] ) messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] with pytest.raises(RuntimeError, match="content filter"): _ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)] async def test_stream_missing_terminal_empty_raises() -> None: class EmptyStreamClient: @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]: del param if not stream: return _response(AssistantChatMessage(content="x")) async def chunks() -> AsyncIterator[ChatCompletionChunk]: # Usage-only style chunk with no finish_reason and no content. yield ChatCompletionChunk( id="1", object="chat.completion.chunk", created=0, model="t", choices=[], usage={"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1}, ) return chunks() messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] with pytest.raises(RuntimeError, match="stream ended without a terminal"): _ = [e async for e in run_chat_loop(EmptyStreamClient(), messages, model="m", stream=True)] async def test_stream_payload_without_finish_reason_ok() -> None: """Some providers omit finish_reason; non-empty text still counts as terminal.""" class PayloadNoFinishClient: @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]: del param if not stream: return _response(AssistantChatMessage(content="ok")) async def chunks() -> AsyncIterator[ChatCompletionChunk]: yield ChatCompletionChunk( id="1", object="chat.completion.chunk", created=0, model="t", choices=[ ChunkChoice( index=0, delta=DeltaMessage(content="ok"), finish_reason=None, ) ], ) return chunks() messages: list[AnyChatMessage] = [UserChatMessage(content="hi")] events = [e async for e in run_chat_loop(PayloadNoFinishClient(), messages, model="m", stream=True)] assert any(isinstance(e, TextDeltaEvent) and e.content == "ok" for e in events)