plyngent/tests/test_agent/test_loop.py

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)