ag-ui-protocol / ag-ui-protocol/ag-ui

[adk-middleware] Per-request Runner.close() also closes caller-owned toolsets (MCP sessions dropped every request)

Đang mở
#2,659 0 bình luận 0 reaction 0 người được giao Xem trên GitHub
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

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.