[Bug]: closing Agent.reply_stream leaves concurrent tool workers running #3218

Prerequisites

• I searched the existing issues, pull requests, and discussions. The related lifecycle fixes and the difference covered here are listed below.
• This is a reproducible bug in the public SDK, rather than a usage question.

Background / Description

Explicitly closing Agent.reply_stream() can return while a concurrency-safe tool started by that reply is still running. Releasing the tool afterward lets it perform an action even though the consumer has already awaited stream.aclose().

The case I wanted to check was a caller that reads part of a reply, decides it no longer needs the result, and closes the stream before leaving its request scope. That caller has a reason to wait for aclose(): it needs to know that the reply no longer owns live work before persisting state or releasing request-scoped resources. In the reproduction below, that boundary does not hold.

The model is deliberately non-streaming and offline. It returns a tool call immediately. The tool then waits on an asyncio.Event, which lets the caller close the public agent stream at a known point. No provider response, HTTP connection, middleware plugin, or external tool service is involved.

Steps to Reproduce

Use main at faf54fdd463b933cb29f46a80ae9f2a758d06bff (2.0.10dev), install the base package dependencies, and save the following as repro_agent.py:

python -m pip install -e .
python repro_agent.py

import asyncio

from agentscope.agent import Agent, InjectionConfig
from agentscope.credential import CredentialBase
from agentscope.event import ToolResultStartEvent
from agentscope.formatter import OpenAIChatFormatter
from agentscope.message import TextBlock, ToolCallBlock, UserMsg
from agentscope.model import ChatModelBase, ChatResponse
from agentscope.permission import PermissionBehavior, PermissionDecision
from agentscope.tool import ToolBase, ToolChunk, Toolkit


class OfflineModel(ChatModelBase):
    def __init__(self):
        super().__init__(CredentialBase(), "offline", self.Parameters(), stream=False)
        self.formatter = OpenAIChatFormatter()

    async def _call_api(self, *args, **kwargs):
        return ChatResponse(
            content=[ToolCallBlock(id="call-1", name="blocked", input="{}")],
            is_last=True,
        )


class BlockedTool(ToolBase):
    name = "blocked"
    description = "Wait before recording a local action."
    input_schema = {"type": "object", "properties": {}}
    is_concurrency_safe = True
    is_read_only = False

    def __init__(self):
        super().__init__()
        self.started = asyncio.Event()
        self.release = asyncio.Event()
        self.workers = []
        self.actions = []
        self.cancelled = False

    async def check_permissions(self, tool_input, context):
        return PermissionDecision(behavior=PermissionBehavior.ALLOW, message="offline test")

    async def __call__(self):
        self.workers.append(asyncio.current_task())
        self.started.set()
        try:
            await self.release.wait()
            self.actions.append("ran")
            yield ToolChunk(content=[TextBlock(text="done")], is_last=True)
        except asyncio.CancelledError:
            self.cancelled = True
            raise


async def main():
    tool = BlockedTool()
    agent = Agent(
        name="bot", system_prompt="Use blocked.", model=OfflineModel(),
        toolkit=Toolkit(tools=[tool]),
        injection_config=InjectionConfig(inject_runtime_state=False),
    )
    stream = agent.reply_stream(UserMsg("user", "Run the tool"))
    try:
        while not isinstance(await anext(stream), ToolResultStartEvent):
            pass
        await asyncio.wait_for(tool.started.wait(), 2)
        await stream.aclose()
        print("cancelled after close:", tool.cancelled)
        print("workers joined after close:", all(task.done() for task in tool.workers))
        tool.release.set()
        await asyncio.gather(*tool.workers, return_exceptions=True)
        print("actions after release:", tool.actions)
    finally:
        tool.release.set()
        await asyncio.gather(*tool.workers, return_exceptions=True)
        await stream.aclose()


asyncio.run(main())

The final cleanup joins the test workers, including on the original implementation. The local action is an append to a list; the script does not write files or call a remote service. Synchronization uses events rather than a timing-dependent sleep.

Observed behavior

cancelled after close: False
workers joined after close: False
actions after release: ['ran']

There is no exception at the public aclose() call. The unexpected behavior is the successful return while work owned by the reply is still live, followed by the local action after closure.

Expected behavior

Before the awaited public close returns, the agent should close the delegated streams and cancel and join its concurrent workers. Releasing the blocked test tool after that point should not run its action:

cancelled after close: True
workers joined after close: True
actions after release: []

An explicit close has no remaining event consumer. Cleanup should therefore finish without trying to yield interruption or final-answer events into the closing generator. A reply that is consumed normally should still finish its tools and emit the usual events.

Why this happens

The delegation chain is reply_stream() → _reply() → _reply_impl() → the tool-event generator. Each wrapper advances the next generator with async for, but that alone does not establish an awaited close relationship. A reply middleware can add another forwarding layer with the same behavior.

Separately, _execute_concurrent_tool_calls() creates a task that gathers the workers and handles CancelledError, but does not put worker ownership cleanup in a finally block. Consumer aclose() injects GeneratorExit at a suspended yield; it does not take the ordinary task-cancellation branch. Fixing only one wrapper would therefore leave other close boundaries, including middleware forwarding, uncovered.

The affected code is in src/agentscope/agent/_agent.py. The proposed fix makes downstream generator closure explicit, joins owned workers on exit, and finishes interrupted tool state without yielding during GeneratorExit.

Impact and scope

This affects SDK integrations that explicitly close an agent stream while concurrent tools are active. It can leave actions and state updates happening after the application considers the reply closed. The reproduction establishes that behavior for awaited public aclose(); it does not assume that every browser disconnect or HTTP handler uses that path.

Cancellation remains cooperative. The requested contract concerns stopping and joining owned asynchronous work, not rolling back a side effect that a tool already completed.

Related work

• [Bug]: Closing the public model stream does not await provider response cleanup #3210 / fix(model): await provider stream cleanup on public aclose #3216 address closure of a provider response through the public model stream. This report uses a non-streaming model and reaches the concurrent agent tool executor.
• [Bug]: cancellation before concurrent toolkit entry can still execute the tool #2767 / fix(agent): preserve cancellation before concurrent tool execution #2768 cover task cancellation before entry into the toolkit, when an interruption event does not yet exist. Here the tool has already entered and the consumer calls aclose() instead of task.cancel().
• [Bug]: ToolOffloadMiddleware keeps the tool running after an interrupt, and parallel calls run twice #3045 / fix(app): cancel the running tool when a reply is interrupted #3048 concern ToolOffloadMiddleware; that middleware is absent from this reproduction.
• [Bug]: Cancelling RealtimeAgent.reply_stream leaves a reader that consumes future events #2665 / fix(agent): release reply_stream queue getter on cancel #2671 concern the RealtimeAgent queue reader. This report uses the core Agent implementation.

Environment

• AgentScope: 2.0.10dev, main faf54fdd463b933cb29f46a80ae9f2a758d06bff.
• Python: 3.12.14.
• OS: Windows.
• No model-provider credentials or network calls are required.

A focused fix is available on dafyy321-pixel:fix/agent-stream-cleanup, with coverage for two concurrent tools, nested reply middleware, full consumption, and closing at normal and interrupted end-event boundaries.