agentscope-ai--agentscope/tests/service_agent_interrupt_test.py

# -*- coding: utf-8 -*-
# pylint: disable=missing-class-docstring,missing-function-docstring
"""Service-layer integration tests for the agent interruption pipeline.

Covers the plumbing that translates an external interrupt signal into a
local ``task.cancel()``:

    message bus publish
        → ``CancelDispatcher`` (background subscriber)
        → ``ChatRunRegistry`` lookup
        → ``task.cancel()``
        → agent exits with ``finished_reason='interrupted'``
        → the next reply on the same session works normally.

FastAPI is intentionally not started — the HTTP layer is a thin wrapper
around :meth:`ChatService.interrupt`. This suite validates the
service-level wiring using :class:`InMemoryMessageBus` so no Redis is
required.
"""
import asyncio
from typing import Any
from unittest.async_case import IsolatedAsyncioTestCase

from agentscope.agent import Agent
from agentscope.app._manager import (
    CancelDispatcher,
    ChatRunRegistry,
    BackgroundTaskManager,
)
from agentscope.app.message_bus import InMemoryMessageBus, MessageBusKeys
from agentscope.event import ReplyEndEvent
from agentscope.message import (
    TextBlock,
    UserMsg,
)
from agentscope.model import ChatResponse
from agentscope.tool import Toolkit


class ServiceAgentInterruptTest(IsolatedAsyncioTestCase):
    """Service-layer interrupt plumbing tests."""

    async def test_full_interrupt_flow(self) -> None:
        """Complete flow: user message → agent runs → interrupt published
        → CancelDispatcher cancels task → agent exits with interrupted
        → next message works.
        """
        bus = InMemoryMessageBus()
        registry = ChatRunRegistry()
        bg_manager = BackgroundTaskManager(message_bus=bus)

        session_id = "full-e2e-session"

        # ---- Build a model that streams slowly ----
        class SlowModel:
            """Streaming model with await points for testing."""

            def __init__(self) -> None:
                self.model = "slow-e2e"
                self.stream = True
                self.max_retries = 0
                self.context_size = 1000

            async def __call__(
                self,
                *_args: Any,
                **_kwargs: Any,
            ) -> Any:
                async def _stream() -> Any:
                    await asyncio.sleep(0.03)
                    yield ChatResponse(
                        content=[TextBlock(text="part1 ")],
                        is_last=False,
                    )
                    await asyncio.sleep(0.03)
                    yield ChatResponse(
                        content=[TextBlock(text="part2")],
                        is_last=False,
                    )
                    await asyncio.sleep(0.03)
                    yield ChatResponse(
                        content=[TextBlock(text="part1 part2 full")],
                        is_last=True,
                    )

                return _stream()

            async def count_tokens(
                self,
                *_args: Any,
                **_kwargs: Any,
            ) -> int:
                return 100

        # ---- Agent ----
        agent = Agent(
            name="FullE2EAgent",
            system_prompt="You are a test agent.",
            model=SlowModel(),
            toolkit=Toolkit(),
        )

        # ---- Start agent in background, register in ChatRunRegistry ----
        finished_reason_1 = None

        async def _chat_run() -> None:
            nonlocal finished_reason_1
            async for evt in agent.reply_stream(
                UserMsg(name="user", content="Hello"),
            ):
                if isinstance(evt, ReplyEndEvent):
                    finished_reason_1 = evt.finished_reason

        registry.spawn(_chat_run(), session_id=session_id)

        # ---- Start CancelDispatcher ----
        async with bus:
            async with CancelDispatcher(
                message_bus=bus,
                registry=registry,
                bg_manager=bg_manager,
            ):
                # Wait for agent to start streaming
                await asyncio.sleep(0.04)

                # Step 1: Publish interrupt (simulating API endpoint)
                await bus.publish(
                    MessageBusKeys.session_interrupt_channel(),
                    {"session_id": session_id},
                )

                # Wait for cancellation to propagate and agent to finish
                await asyncio.sleep(0.3)

                # Step 2: Verify agent was interrupted
                self.assertEqual(
                    finished_reason_1,
                    "interrupted",
                    "Full flow: agent should exit with interrupted",
                )

                # Step 3: Verify chat-run task is done
                task = registry.get(session_id)
                self.assertTrue(
                    task is None or task.done(),
                    "Chat-run task should be cleaned up after interrupt",
                )

        # ---- Step 4: Next conversation round should work ----
        model2 = SlowModel()
        agent2 = Agent(
            name="FullE2EAgent-Round2",
            system_prompt="You are a test agent.",
            model=model2,
            toolkit=Toolkit(),
        )
        # Copy the interrupted context
        agent2.state = agent.state

        finished_reason_2 = None
        async for evt in agent2.reply_stream(
            UserMsg(name="user", content="Continue please"),
        ):
            if isinstance(evt, ReplyEndEvent):
                finished_reason_2 = evt.finished_reason

        self.assertEqual(
            finished_reason_2,
            "completed",
            "Next round after full-flow interruption should complete normally",
        )