[Bug]: ToolOffloadMiddleware keeps the tool running after an interrupt, and parallel calls run twice #3045

import asyncio
import json
import time
import httpx
from agentscope.agent import Agent
from agentscope.app._manager import (
    BackgroundTaskManager,
    CancelDispatcher,
    ChatRunRegistry,
)
from agentscope.app.message_bus import InMemoryMessageBus, MessageBusKeys
from agentscope.app.middleware import ToolOffloadMiddleware
from agentscope.credential import OpenAICredential
from agentscope.event import ReplyEndEvent
from agentscope.message import TextBlock, UserMsg
from agentscope.model import OpenAIChatModel
from agentscope.tool import FunctionTool, ToolChunk, Toolkit

T0 = time.monotonic()
log = []

def note(text: str) -> None:
    log.append(f"{time.monotonic() - T0:.1f}s {text}")

async def search(query: str) -> ToolChunk:
    """Search the files for a query."""
    note(f"start {query}")
    try:
        await asyncio.sleep(1)
    except asyncio.CancelledError:
        note(f"cancelled {query}")
        raise
    note(f"finished {query}")
    return ToolChunk(content=[TextBlock(text=f"found {query}")])

def fake_openai(n_calls: int, requests: list):
    # The first reply calls the tool n_calls times, the next one is text
    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        message = {"role": "assistant", "content": "Done."}
        if len(requests) == 1:
            message = {"role": "assistant", "content": None, "tool_calls": [{
                "id": f"call_{i}", "type": "function",
                "function": {"name": "search",
                             "arguments": json.dumps({"query": f"q{i}"})},
            } for i in range(n_calls)]}
        return httpx.Response(200, json={
            "id": "chatcmpl-1", "object": "chat.completion",
            "created": 0, "model": "gpt-4.1",
            "choices": [{"index": 0, "message": message,
                         "finish_reason": "stop"}],
            "usage": {"prompt_tokens": 10, "completion_tokens": 10,
                      "total_tokens": 20},
        })
    return handler

async def run(n_calls: int, parallel: bool) -> None:
    global T0
    log.clear()
    requests = []
    model = OpenAIChatModel(
        credential=OpenAICredential(api_key="sk-test"),
        model="gpt-4.1",
        stream=False,
        client_kwargs={"http_client": httpx.AsyncClient(
            transport=httpx.MockTransport(fake_openai(n_calls, requests)))},
    )
    bus = InMemoryMessageBus()
    bg_manager = BackgroundTaskManager(message_bus=bus)
    registry = ChatRunRegistry()
    agent = Agent(
        name="Friday",
        system_prompt="You are a helpful assistant.",
        model=model,
        toolkit=Toolkit(tools=[FunctionTool(
            search, is_read_only=True, is_concurrency_safe=parallel)]),
        # ChatService.run installs this middleware on every chat run
        middlewares=[ToolOffloadMiddleware(
            bg_manager=bg_manager, message_bus=bus,
            user_id="u", agent_id="a")],
    )
    async def chat_run() -> None:
        async for event in agent.reply_stream(UserMsg("user", "Search.")):
            if isinstance(event, ReplyEndEvent):
                note(f"reply ended: {event.finished_reason}")

    async with bus:
        async with CancelDispatcher(
            message_bus=bus, registry=registry, bg_manager=bg_manager,
        ):
            T0 = time.monotonic()
            registry.spawn(chat_run(), session_id="s1")
            await asyncio.sleep(0.3)
            # What POST /sessions/s1/interrupt does for a running session
            note("interrupt")
            await bus.publish(
                MessageBusKeys.session_interrupt_channel(),
                {"session_id": "s1"},
            )
            await asyncio.sleep(2)

    kind = "parallel" if parallel else "sequential"
    print(f"{n_calls} {kind} call(s):")
    for line in log:
        print("  " + line)
    print(f"  model requests: {len(requests)}")


async def main() -> None:
    await run(1, parallel=False)
    await run(2, parallel=True)


asyncio.run(main())