agentscope-ai / agentscope-ai/agentscope
[Bug]: ReplyBudgetControlMiddleware raises KeyError when an outer middleware swallows the ReplyEndEvent
- 主要語言
- Python
- 星號
- 31.5k
- 分支
- 3.5k
- 平均合併
- 1 天 23 小時
- 30 天內合併 PR
- 95
描述
### Prerequisites
- [x] I have searched the existing [issues](https://github.com/agentscope-ai/agentscope/issues) and [discussions](https://github.com/agentscope-ai/agentscope/discussions), and this is not a duplicate.
- [x] This is a bug, not a usage question. (For questions, please use [Discussions](https://github.com/agentscope-ai/agentscope/discussions/new?category=general) instead.)
### Background / Description
`ReplyBudgetControlMiddleware.on_reply` keeps a per-reply cost counter in `agent.state.middle_context[middleware_key][reply_id]`: it creates the entry on `ReplyStartEvent`, accumulates on every `ModelCallEndEvent`, and pops it on `ReplyEndEvent`.
Swallowing the `ReplyEndEvent` from an outer `on_reply` middleware is a supported pattern — the `MiddlewareBase.on_reply` docstring describes it, and `tests/middleware_test.py::test_on_reply_middleware_swallow_reply_end` verifies that the agent then runs another reasoning round without emitting a new `ReplyStartEvent`.
Combining the two crashes the reply. With `middlewares=[SwallowOnceMiddleware(), ReplyBudgetControlMiddleware(...)]` the budget middleware is the inner layer, so events reach it first on the way out:
1. Round 1 finishes and `_reply_impl` emits `ReplyEndEvent`.
2. The budget middleware sees it first and pops the counter for this `reply_id`.
3. The outer middleware swallows the event, so the agent starts a redo round. No new `ReplyStartEvent` is emitted, so nothing re-creates the counter.
4. The redo round's `ModelCallEndEvent` hits the accumulation branch in `src/agentscope/middleware/_budget.py` (lines 139-146). The defensive check there only covers a missing `middleware_key` — after the pop the outer dict still exists (now empty), so the check passes and `agent.state.middle_context[middleware_key][event.reply_id] += ...` raises `KeyError`, killing the whole reply.
I expected the reply to complete with the budget still enforced across the redo round.
Two things I checked while digging into this:
- The HITL pause path is fine: pausing at `REQUIRE_USER_CONFIRM` does not emit a `ReplyEndEvent`, so the counter survives, as `test_token_accumulation_persists_across_hitl` shows.
- With the reversed order `[Budget, Swallow]` the bug doesn't fire, because the inner swallow middleware eats the event before the budget middleware sees it. Nothing documents an ordering requirement, and a middleware can't know what sits outside it.
A minimal fix is to accumulate defensively on `ModelCallEndEvent`:
```python
bucket = agent.state.middle_context.setdefault(middleware_key, {})
bucket[event.reply_id] = bucket.get(event.reply_id, 0) + (
self.input_token_weight * event.input_tokens
+ self.output_token_weight * event.output_tokens
)
```
This stops the crash, though the count restarts from zero after a swallowed end, so earlier rounds no longer count against the budget. Keeping the full count would mean moving cleanup from `ReplyEndEvent` to the next `ReplyStartEvent` (replace the whole bucket there), at the price of one stale entry staying in `middle_context` between replies — that would also change what `test_state_cleanup_after_reply` asserts. I can send a PR for whichever direction you prefer.
### Error Messages
```shell
Traceback (most recent call last):
File "repro.py", line 82, in
asyncio.run(main())
...
File ".../src/agentscope/agent/_agent.py", line 833, in execute_chain
async for item in mw.on_reply(
File "repro.py", line 58, in on_reply
async for item in next_handler():
File ".../src/agentscope/agent/_agent.py", line 827, in next_handler
async for item in execute_chain(
File ".../src/agentscope/agent/_agent.py", line 833, in execute_chain
async for item in mw.on_reply(
File ".../src/agentscope/middleware/_budget.py", line 143, in on_reply
agent.state.middle_context[middleware_key][event.reply_id] += (
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^
KeyError: 'd54f36c93c064250b9cbcdceee368629'
```
### Steps to Reproduce
Self-contained, no API key needed — the mock model returns a fixed response, and the swallow middleware is the same pattern as `tests/middleware_test.py::test_on_reply_middleware_swallow_reply_end`.
1. Code:
```python
import asyncio
from typing import Any, AsyncGenerator, Type
from pydantic import BaseModel
from agentscope.agent import Agent
from agentscope.credential import CredentialBase
from agentscope.event import ReplyEndEvent
from agentscope.formatter import OpenAIChatFormatter
from agentscope.message import TextBlock, UserMsg
from agentscope.middleware import MiddlewareBase, ReplyBudgetControlMiddleware
from agentscope.model import ChatModelBase, ChatResponse, ChatUsage
from agentscope.tool import Toolkit
class MockCredential(CredentialBase):
@classmethod
def get_chat_model_class(cls) -> Type["ChatModelBase"]:
return MockModel
class MockModel(ChatModelBase):
"""Offline stub so the repro needs no API key."""
class Parameters(BaseModel):
pass
def __init__(self) -> None:
super().__init__(
credential=MockCredential(),
model="mock",
stream=False,
parameters=MockModel.Parameters(),
context_size=1000,
)
self.formatter = OpenAIChatFormatter()
self.cnt = 0
async def _call_api(self, *args: Any, **kwargs: Any) -> ChatResponse:
self.cnt += 1
return ChatResponse(
content=[TextBlock(text=f"answer {self.cnt}")],
is_last=True,
usage=ChatUsage(input_tokens=10, output_tokens=5, time=0.0),
)
class SwallowOnceMiddleware(MiddlewareBase):
"""Same pattern as the swallow test in tests/middleware_test.py."""
def __init__(self) -> None:
self.swallowed = False
async def on_reply(self, agent, input_kwargs, next_handler) -> AsyncGenerator:
async for item in next_handler():
if isinstance(item, ReplyEndEvent) and not self.swallowed:
self.swallowed = True
continue # force one more reasoning round
yield item
async def main() -> None:
agent = Agent(
name="assistant",
system_prompt="You are a helpful assistant.",
model=MockModel(),
toolkit=Toolkit(),
middlewares=[
SwallowOnceMiddleware(), # outer
ReplyBudgetControlMiddleware(token_budget=1_000_000), # inner
],
)
async for _ in agent.reply_stream(UserMsg("user", "hi")):
pass
print("no crash")
asyncio.run(main())
```
2. Run: `python repro.py`
Observed: the redo round crashes with the `KeyError` above.
Expected: the reply completes ("no crash" is printed) and the budget keeps counting.
Removing `SwallowOnceMiddleware` from the list makes it finish normally, and the budget is large enough that exhaustion plays no role here.
### Environment
- AgentScope Version: 2.0.7 (also reproduces on current main, e8c8d7b9)
- Python Version: 3.12
- OS: Linux
貢獻指南
評估
這個 Issue 還沒有評估資料。