ag-ui-protocol / ag-ui-protocol/ag-ui
[adk-middleware] Per-request Runner.close() also closes caller-owned toolsets (MCP sessions dropped every request)
- Ngôn ngữ chính
- Python
- Star
- 15.9k
- Fork
- 1.4k
- Merge trung bình
- 1 ngày 17 giờ
- Pull request đã merge (30 ngày)
- 163
Mô tả
**TL;DR** — Same root cause as #2642, wider blast radius. The per-request `Runner` is closed over a shallow copy of the caller's agent tree, and `Runner.close()` closes every toolset it can reach. Tools are shared **by reference** on purpose, so a caller who puts an `McpToolset` in `agent.tools` has it closed on every single request. #2642 fixes the plugin half using ADK's `set_skip_closing_plugins`; there is no toolset analogue, so this half needs a maintainer decision.
---
## For humans
### What happens
`_run_adk_in_background` closes the per-request runner in its `finally`, with the comment *"Ensure the ADK runner releases any resources (e.g. toolsets)"*. `Runner.close()` calls `_cleanup_toolsets(_collect_toolset(self.agent))`, which walks `agent.tools` and closes every `BaseToolset` it finds.
The agent it walks is the per-request shallow copy. `_shallow_copy_agent_tree` shares tool objects by reference — deliberately, and its own comment explains why (an `McpToolset` holds an unpicklable `TextIOWrapper`, so deep-copying is not an option). So the toolsets the runner closes are the caller's, not the request's.
Measured on `main`, three requests against one long-lived `App`:
```
caller-owned toolset close() calls after 3 requests: 3
```
For `McpToolset` specifically, `close()` clears the tool-list cache and closes the session manager. For a stdio transport that is a subprocess. Every request.
### Why this is the more serious half
The plugin version of this bug cost us a BigQuery quota and 25 seconds of median latency (#2642, and google/adk-python#7017 for the amplification). This one drops a live MCP connection per request. Anyone running the middleware with an MCP toolset is reconnecting on every turn and probably reading it as MCP flakiness rather than a lifecycle bug.
It is also the case the `finally` block was actually written for. Nothing the middleware itself creates needs it:
| What `Runner.close()` does | Does the middleware need it? |
|---|---|
| `_cleanup_toolsets(...)` | No. The only toolset the middleware creates is `ClientProxyToolset`, whose `close()` is a bare log line. Everything else in that tree is the caller's. |
| `plugin_manager.close()` | No — that is #2642. |
| `session_service.flush()` | No. `flush()` is an unoverridden no-op on every stock ADK session service, and the service is long-lived and shared anyway. |
### Why I did not just fix it in #2642
`set_skip_closing_plugins` gave that half a one-line, version-gated, inert-on-old-ADK fix. There is no `set_skip_closing_toolsets`, so the options here all change behaviour rather than opt out of it:
1. **Close only what the request created.** Replace `runner.close()` with closing `client_proxy_toolsets` (already a parameter in scope) plus `await session_service.flush()`. Cleanest, and it dissolves both ownership violations — but it drops a cleanup path in a way that pre-commits the project to "the middleware never owns a toolset".
2. **Track ownership explicitly** — record which toolsets the request created and close only those.
3. **Ask ADK for a `skip_closing_toolsets` analogue**, and mirror #2642.
That is a design call, so I would rather it be yours than a drive-by from me. Happy to send the patch for whichever you prefer.
---
## For AI agents
```yaml
defect:
id: adk-middleware-closes-caller-owned-toolsets
component: ag_ui_adk.ADKAgent
package: ag-ui-adk
version: "0.7.0"
path: integrations/adk-middleware/python
status_at_head: present
class: [resource-lifecycle, ownership-violation]
sibling: "#2642 (same root cause, plugin half)"
severity_vs_sibling: higher # drops live MCP sessions / stdio subprocesses
root_cause: >-
The per-request Runner is built over a shallow copy of the caller's agent
tree (tools shared by reference by design) and is closed in a finally
block. Runner.close() closes every BaseToolset reachable from agent.tools,
which are the caller's objects, not the request's.
call_chain:
- file: src/ag_ui_adk/adk_agent.py
symbol: _shallow_copy_agent_tree
note: "shares tool objects by reference; deep copy rejected because McpToolset holds an unpicklable TextIOWrapper"
- file: src/ag_ui_adk/adk_agent.py
symbol: _run_adk_in_background
note: "DEFECT SITE — finally closes the per-request runner"
code: 'close_method = getattr(runner, "close", None)'
- repo: google/adk-python
file: src/google/adk/runners.py
symbol: Runner.close
note: "_cleanup_toolsets(_collect_toolset(self.agent)) -> toolset.close()"
reproduce: |
class RecordingToolset(BaseToolset):
def __init__(self): super().__init__(); self.close_count = 0
async def get_tools(self, readonly_context=None): return []
async def close(self): self.close_count += 1
toolset = RecordingToolset()
app = App(name="t", root_agent=LlmAgent(name="a", model=M, tools=[toolset]))
agent = ADKAgent.from_app(app, user_id="u", use_in_memory_services=True)
for _ in range(3):
runner = agent._create_runner(adk_agent=agent._adk_agent, user_id="u", app_name="t")
await runner.close()
assert toolset.close_count == 0 # fails: it is 3
observed: "close_count == 3 after 3 requests"
expected: "close_count == 0 (the caller owns it)"
what_runner_close_provides:
cleanup_toolsets: "not needed — only ClientProxyToolset is middleware-created, and its close() is a log line"
plugin_manager_close: "not needed — that is #2642"
session_service_flush: "not needed — unoverridden no-op on stock ADK services; service is shared and long-lived"
fix_options:
- id: close-only-what-the-request-created
change: "replace runner.close() with closing client_proxy_toolsets (already in scope) + await session_service.flush()"
pros: "dissolves both ownership violations; no ADK API needed"
cons: "removes a cleanup path; commits the project to 'middleware never owns a toolset'"
- id: explicit-ownership-tracking
change: "record request-created toolsets and close only those"
- id: upstream-analogue
change: "request set_skip_closing_toolsets in ADK and mirror #2642"
blocked_on: "google/adk-python"
decision_needed_from: maintainer
offer: "patch available for whichever option is preferred"
```
Hướng dẫn đóng góp
Đánh giá
Issue này chưa được đánh giá.