plyngent/tests/test_agent/test_run_aside.py

"""ChatAgent.run_aside: side turns leave main transcript/memory alone."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal, cast, overload

import pytest

from plyngent.agent import ChatAgent, ToolRegistry, ToolTag, aside_notice, tool
from plyngent.agent.todo_stack import TodoStack
from plyngent.config.models import DatabaseConfig
from plyngent.lmproto.openai_compatible.model import (
    AnyChatMessage,
    AssistantChatMessage,
    ChatCompletionChoice,
    ChatCompletionChunk,
    ChatCompletionResponse,
    ChatCompletionsParam,
    DeveloperChatMessage,
    UserChatMessage,
)
from plyngent.memory import MemoryStore
from plyngent.tools.context import InstanceState, SessionState

if TYPE_CHECKING:
    from collections.abc import AsyncIterator


class ScriptedClient:
    """Returns a fixed assistant string each call."""

    def __init__(self, replies: list[str] | None = None) -> None:
        self.replies = list(replies or ["aside-answer"])
        self.calls = 0
        self.payloads: list[list[AnyChatMessage]] = []

    @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.calls += 1
        self.payloads.append(list(param.messages))
        text = self.replies[(self.calls - 1) % len(self.replies)]
        return ChatCompletionResponse(
            id="1",
            object="chat.completion",
            created=0,
            model="m",
            choices=[
                ChatCompletionChoice(
                    index=0,
                    message=AssistantChatMessage(content=text),
                    logprobs={},
                    finish_reason="stop",
                )
            ],
            system_fingerprint="",
            usage={},
        )


@pytest.mark.asyncio
async def test_run_aside_does_not_mutate_main_or_memory() -> None:
    memory = await MemoryStore.open(DatabaseConfig())
    try:
        session = await memory.create_session(name="main")
        client = ScriptedClient(["main-reply", "aside-reply", "main-again"])
        agent = ChatAgent(
            cast("Any", client),
            model="m",
            memory=memory,
            session_id=session.sid,
            stream=False,
            todo_stack=TodoStack(),
        )
        async for _ in agent.run("hello main"):
            pass
        main_len = len(agent.messages)
        main_snapshot = list(agent.messages)
        db_before = await memory.list_messages(session.sid)

        texts: list[str] = []
        async for event in agent.run_aside("side question?", include_history=True, tools=False):
            from plyngent.agent.events import AssistantMessageEvent

            if isinstance(event, AssistantMessageEvent) and event.message.content:
                texts.append(str(event.message.content))

        assert any("aside-reply" in t for t in texts)
        assert len(agent.messages) == main_len
        assert agent.messages == main_snapshot
        db_after = await memory.list_messages(session.sid)
        assert len(db_after) == len(db_before)
        # Aside request saw main history + side user.
        aside_payload = client.payloads[1]
        assert any(isinstance(m, UserChatMessage) and "side question" in m.content for m in aside_payload)
        assert any(isinstance(m, UserChatMessage) and "hello main" in m.content for m in aside_payload)
    finally:
        await memory.close()


@pytest.mark.asyncio
async def test_run_aside_fresh_skips_history() -> None:
    client = ScriptedClient(["only"])
    agent = ChatAgent(cast("Any", client), model="m", stream=False)
    agent.messages.append(UserChatMessage(content="prior"))
    async for _ in agent.run_aside("q", include_history=False, tools=False):
        pass
    payload = client.payloads[0]
    users = [m.content for m in payload if isinstance(m, UserChatMessage)]
    assert users == ["q"]


@pytest.mark.asyncio
async def test_run_aside_notice_reaches_model_but_not_history() -> None:
    """The aside notice is the model's last instruction; the host keeps none."""
    client = ScriptedClient(["ok"])
    agent = ChatAgent(cast("Any", client), model="m", stream=False)
    agent.messages.append(UserChatMessage(content="prior main"))
    before = list(agent.messages)
    notice = aside_notice(tools_mode="read")

    async for _ in agent.run_aside("side?", include_history=True, tools=False, notice=notice):
        pass

    payload = client.payloads[0]
    assert isinstance(payload[-2], DeveloperChatMessage)
    assert payload[-2].content == notice.to_message().content
    assert "read-only" in payload[-2].content
    assert isinstance(payload[-1], UserChatMessage)
    assert payload[-1].content == "side?"
    assert agent.messages == before


@pytest.mark.asyncio
async def test_run_aside_fresh_history_still_carries_notice() -> None:
    client = ScriptedClient(["ok"])
    agent = ChatAgent(cast("Any", client), model="m", stream=False)

    async for _ in agent.run_aside(
        "q",
        include_history=False,
        tools=False,
        notice=aside_notice(tools_mode="no"),
    ):
        pass

    payload = client.payloads[0]
    assert isinstance(payload[0], DeveloperChatMessage)
    assert "Tools are disabled" in payload[0].content
    assert isinstance(payload[-1], UserChatMessage)
    assert payload[-1].content == "q"


@pytest.mark.asyncio
async def test_run_aside_tools_clones_registry_with_fresh_session() -> None:
    hits: list[str] = []

    @tool(tags=ToolTag.LOCAL | ToolTag.SESSION_STATE, register=False)
    async def note_session() -> str:
        from plyngent.tools.context import get_session

        session = get_session()
        hits.append("ok" if session is not None else "none")
        if session is not None:
            session.extras["aside"] = True
        return "noted"

    main_session = SessionState(session_id="main")
    instance = InstanceState()
    registry = ToolRegistry(
        [note_session],
        auto_bind_state=True,
        instance_state=instance,
        session_state=main_session,
    )

    class ToolThenStop:
        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 stream, param
            self.calls += 1
            from plyngent.lmproto.openai_compatible.model import (
                AssistantFunctionTool,
                AssistantFunctionToolCall,
            )

            if self.calls == 1:
                message = AssistantChatMessage(
                    content="",
                    tool_calls=[
                        AssistantFunctionToolCall(
                            id="c1",
                            function=AssistantFunctionTool(name="note_session", arguments="{}"),
                        )
                    ],
                )
                finish = "tool_calls"
            else:
                message = AssistantChatMessage(content="done")
                finish = "stop"
            return ChatCompletionResponse(
                id="1",
                object="chat.completion",
                created=0,
                model="m",
                choices=[
                    ChatCompletionChoice(
                        index=0,
                        message=message,
                        logprobs={},
                        finish_reason=finish,
                    )
                ],
                system_fingerprint="",
                usage={},
            )

    client = ToolThenStop()
    agent = ChatAgent(cast("Any", client), model="m", tools=registry, stream=False)
    async for _ in agent.run_aside(
        "use tool",
        tools=True,
        instance_state=instance,
        session_state=SessionState(session_id="aside"),
    ):
        pass
    assert hits == ["ok"]
    assert "aside" not in main_session.extras


async def test_run_aside_read_only_clone_excludes_write_tools() -> None:
    """A read-only-cloned registry passed to run_aside exposes no write tools."""
    hits: list[str] = []

    @tool(tags=ToolTag.LOCAL | ToolTag.READ_ONLY, register=False)
    async def peek() -> str:
        hits.append("peek")
        return "peeked"

    @tool(tags=ToolTag.LOCAL, register=False)
    async def write_thing() -> str:
        hits.append("write")
        return "wrote"

    registry = ToolRegistry([peek, write_thing], auto_bind_state=True)

    def make_client(first_tool: str | None) -> Any:
        calls = 0

        class _Client:
            @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, param
                nonlocal calls
                calls += 1
                if calls == 1 and first_tool is not None:
                    from plyngent.lmproto.openai_compatible.model import (
                        AssistantFunctionTool,
                        AssistantFunctionToolCall,
                    )

                    message = AssistantChatMessage(
                        content="",
                        tool_calls=[
                            AssistantFunctionToolCall(
                                id="c1",
                                function=AssistantFunctionTool(name=first_tool, arguments="{}"),
                            )
                        ],
                    )
                    finish = "tool_calls"
                else:
                    message = AssistantChatMessage(content="done")
                    finish = "stop"
                return ChatCompletionResponse(
                    id="1",
                    object="chat.completion",
                    created=0,
                    model="m",
                    choices=[
                        ChatCompletionChoice(
                            index=0,
                            message=message,
                            logprobs={},
                            finish_reason=finish,
                        )
                    ],
                    system_fingerprint="",
                    usage={},
                )

        return _Client()

    # Read-only tool is available to the side turn.
    agent_peek = ChatAgent(make_client("peek"), model="m", tools=registry, stream=False)
    async for _ in agent_peek.run_aside("go", tools=registry.clone(read_only_only=True)):
        pass
    assert hits == ["peek"]

    # Write tool is not in the read-only clone: the call errors, never runs.
    hits.clear()
    agent_write = ChatAgent(make_client("write_thing"), model="m", tools=registry, stream=False)
    async for _ in agent_write.run_aside("go", tools=registry.clone(read_only_only=True)):
        pass
    assert hits == []