diff --git a/prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.md b/prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.md new file mode 100644 index 000000000..52b29c337 --- /dev/null +++ b/prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.md @@ -0,0 +1,135 @@ +# Fix standard MCP session deadlock + +[中文版](2026-10-26-standard-mcp-deadlock.zh.md) + +- Change ID: `trusted-mcp-session-routing` +- Created/revised: 2026-10-26 +- Status: `approved` +- Component contract: [Trusted MCP sessions](../../../specs/trusted-mcp-sessions/README.md) +- Baseline revision: `c662cc8d`; verification covers the working-tree session manager and two affected test files. + +## Background and evidence + +The supplied report describes VeADK 1.0.2 with Google ADK 2.1.0, MCP 1.26.0, and Python 3.12. A first standard streamable HTTP session, including Bearer authentication without `x-trusted-mcp`, waits indefinitely before connecting. + +In `veadk/tools/mcp_tool/trusted_mcp_session_manager.py`, `create_session` acquires `_session_lock`, then calls `super().create_session(headers)` for a standard connection. The parent also acquires the same non-reentrant lock. A cache miss therefore waits for a lock that its own caller holds. A healthy cached session can return before reaching this defect. The trusted branch creates its own session and avoids the nested parent call. + +The initially installed local parent implementation confirmed the nested-lock mechanism; that environment had Google ADK 1.18.0 and MCP 1.20.0. Repository metadata instead requires `google-adk>=2.7.0,<2.8.0` and `mcp==1.26.0`. Initial inspection was source evidence, not a runtime reproduction or verification of the reporter's dependency combination. Runtime checks now use the repository environment resolved by `uv`, as recorded below. + +Existing `tests/tools/mcp_tool/test_trusted_mcp_components.py` covers trusted creation and healthy cache reuse; it lacks standard first-connection coverage. + +## Goals and non-goals + +Restore standard MCP connection creation through `TrustedMcpToolset`, retain trusted routing and authentication-based pooling, and prevent nested locking. + +Dependency upgrades, new locks, transport redesign, broader trusted-session cancellation cleanup changes, release/version publication, cloud calls, and unrelated toolset configuration forwarding are outside scope. Public imports and signatures remain compatible. + +## Scenarios and requirements + +| ID | Given / when | Required result | +| --- | --- | --- | +| `FR-1` | Standard streamable HTTP, SSE, or stdio connection; first session requested | Delegate to the parent before acquiring the subclass session lock; connection creation can progress. | +| `FR-2` | Streamable HTTP with merged `x-trusted-mcp` exactly equal to `"true"` | Use `trusted_mcp_client`; keep the existing trusted session pool and cleanup behavior. | +| `FR-3` | Connection headers plus per-call headers, including a trust override | Use parent merge precedence. Routing inspection does not mutate either caller dictionary. Preserve authentication headers passed to the selected transport. | +| `FR-4` | Standard concurrent requests, repeated requests, or a disconnected cached session | Parent owns serialization, reuse, replacement, and resource cleanup; do not duplicate those operations. | +| `FR-5` | Standard connection failure or cancellation while waiting for the lock | Propagate the failure/cancellation. Release the lock so a later request can proceed; do not add silent success or new timeout behavior. | + +An absent marker, `"false"`, or another nonmatching value selects the standard route. A marker on SSE or stdio never enables the trusted transport. Tests use synthetic authentication values and local transport doubles. + +## Design and contract impact + +Recommended approach: after `_merge_headers(headers)`, inspect connection type and the marker with `.get`, before generating a session key or acquiring `_session_lock`. For a standard route, return `await super().create_session(headers)` immediately, passing the original per-call headers so the parent applies its normal merge and lifecycle. + +Only the trusted route executes the existing subclass session-key, lock, cache, disconnected cleanup, and `AsyncExitStack` operations. Generate the trusted session key before removing the trusted marker from the private merged copy for transport use. Remove the now-unreachable standard fallback inside the subclass lock. + +Supported-baseline implementation review on 2026-10-26 found that Google ADK 2.7.1 calls `_create_client` with additional `mtls_transport` and `session_key` arguments. Accept and forward parent arguments through the existing override so ordinary creation can progress beyond the fixed lock. Trusted routing behavior remains unchanged. This necessary private-interface adaptation fulfills `FR-1`; no new public contract or product scope is introduced. Update the existing ordinary-cache test to use the current parent pool tuple `(session, exit_stack, loop)` rather than its obsolete two-element fixture. + +Alternatives: + +1. Replace the lock with a reentrant async lock: adds custom lock behavior and conflicts with parent lock ownership; unnecessary. +2. Copy parent session creation into the subclass: avoids nesting but duplicates transport and lifecycle behavior, increasing upstream compatibility risk. + +This changes concurrency ownership for standard sessions, so the narrow bilingual component contract is required. No Studio, generated code, dependency metadata, permissions, event schemas, or deployment contracts change. No consumer migration is required. Existing user documentation is assessed during reconciliation; changing inaccurate general documentation beyond this routing defect requires separate scope. + +## Tasks and affected files + +| Task | Work | Requirements | +| --- | --- | --- | +| `T-1` | Review this bilingual PRD and `specs/trusted-mcp-sessions/README{,.zh}.md`; record user approval. | All | +| `T-2` | Add `tests/tools/mcp_tool/test_trusted_mcp_session_routing.py`. First prove standard creation times out before the fix, using the real parent session method with transport/session doubles and a bounded `asyncio.wait_for`. | `FR-1` | +| `T-3` | Apply early routing and parent client-argument forwarding in `veadk/tools/mcp_tool/trusted_mcp_session_manager.py`. | `FR-1`, `FR-2`, `FR-3` | +| `T-4` | Cover standard transport types, marker values/overrides, synthetic Bearer headers, trusted creation/reuse, concurrent creation, disconnected replacement, failure retry, and cancellation while waiting for the lock. Retain existing component tests, updating the ordinary-cache fixture for the parent pool tuple. | `FR-2`–`FR-5` | +| `T-5` | Run targeted and broader regression checks; review diff, resource ownership, and bilingual equivalence; reconcile documents with actual evidence. | All | + +Task order: `T-1` → failing regression in `T-2` → `T-3` → remaining coverage in `T-4` → `T-5`. Record failed pre-fix evidence before implementation. + +Execution plan (approved scope): + +- [x] `T-1`: User approved code and test changes on 2026-10-26: “按这个方案修改代码和测试”. +- [x] `T-2`: Build an isolated transport fixture with real parent locking and session pooling. Run `uv run --extra dev pytest tests/tools/mcp_tool/test_trusted_mcp_session_routing.py -k standard_first_session`; expect a bounded timeout before any transport entry. +- [x] `T-3`: Insert `if not (isinstance(self._connection_params, StreamableHTTPConnectionParams) and merged_headers and merged_headers.get("x-trusted-mcp") == "true"): return await super().create_session(headers)` before key generation. Keep the trusted key, remove its marker from the merged copy before transport entry, and eliminate the locked parent fallback. Run the same first-session test; expect success. +- [x] `T-4`: Run all routing/component tests with isolated transport/session contexts, checking creation counts, identity reuse, header values, cleanup, failure propagation, and post-cancellation progress. +- [x] `T-5`: Execute the verification commands below, review the final source/test diff against `FR-1`–`FR-5`, update both languages and component status, and record actual results. Subsequent user instructions authorize a scoped commit and synchronization of local history to `origin/main`. + +## Verification and acceptance + +Use `uv run --extra dev` with current repository metadata. Do not treat the stale local environment as the supported baseline. Doubles replace network/process transports, not the parent lock implementation. + +| Requirements | Tasks | Acceptance | Check | Current result | +| --- | --- | --- | --- | --- | +| `FR-1` | `T-2`, `T-3` | `AC-1`: first standard creation fails with bounded timeout before the fix and succeeds after it | `uv run --extra dev pytest tests/tools/mcp_tool/test_trusted_mcp_session_routing.py -k standard_first_session -q` | `pass`: two pre-fix timeouts at nested locking; two post-fix successes | +| `FR-2`, `FR-3` | `T-3`, `T-4` | `AC-2`: trusted routing, merge precedence, caller headers, and authentication remain correct | `uv run --extra dev pytest tests/tools/mcp_tool -q` | `pass`: included in 24 passing tests | +| `FR-4`, `FR-5` | `T-4` | `AC-3`: standard lifecycle, concurrent reuse, failure retry, and lock-wait cancellation pass | `uv run --extra dev pytest tests/tools/mcp_tool -q` | `pass`: 24 tests passed in 3.05s | +| All | `T-5` | `AC-4`: broader affected regression passes | `uv run --extra dev pytest -n 2 -m "not codex_smoke and not piagent_smoke"` | `fail`: 7556 passed, 17 failed, 58 skipped, 4 xfailed, 2 collection errors; 436.40s; two workers | +| All | `T-5` | `AC-5`: changed Python and document checks pass | Ruff 0.11.12 lint/format commands in the record below; `git diff --check`; paired-language/path/link review | `pass`: three Python files and both document pairs | + +The user subsequently authorized commit and push to `origin/main`, including local history synchronization. Pre-commit is therefore required before committing. PR creation, package publication, and deployment are not authorized. Live cloud/provider and real trusted attestation checks are `not_applicable`: the bug is lock ownership before transport entry, and verification is isolated. Existing trusted component tests establish simulated compatibility, not real attestation. + +## Risks and open questions + +- Google ADK private methods can evolve; verify against the declared 2.7.x baseline. Testing the reporter's 2.1.0 combination would be additional compatibility evidence, not a dependency-policy change or release guarantee. +- Parent methods re-merge headers on standard delegation. The first merge must remain a side-effect-free routing inspection. +- Parent transport failures and cancellations after resources have been entered retain upstream semantics. This change guarantees no additional subclass-held lock; it does not promise to repair upstream cleanup issues. +- The existing trusted branch catches `Exception`, not all cancellation cases. Broadening that cleanup is outside this focused deadlock fix. +- Google ADK 2.7.1 uses three-element parent pool entries; the unchanged trusted branch stores two elements. Its inherited final `close()` compatibility is a pre-existing limitation outside this ordinary-routing fix. Trusted simulated tests explicitly close their owned stacks and do not claim inherited final-close compatibility. +- No unresolved product choice blocks this design. Supported-environment setup and targeted behavior are verified; the reporter's Google ADK 2.1.0/Python 3.12 combination remains untested. + +## Review and delivery record + +2026-10-26, baseline `c662cc8d`, actual diff scope: the session manager, new isolated routing tests, the existing ordinary-cache test fixture, and this PRD/component-spec pair. + +Direct design review used because no `review-spec` skill is available in the supplied skill catalog. Reviewed ownership, feasibility, interfaces, errors, security, compatibility, testability, and bilingual equivalence. Findings resolved in design: route before locking; inspect headers without mutation; retain trusted key generation before marker removal; use the declared dependency baseline; distinguish source inspection from runtime evidence. No outstanding design blockers. + +User approval: granted on 2026-10-26, “按这个方案修改代码和测试”. The scoped source/test changes and document reconciliation are complete. Final implementation review found no blocking in-scope correctness, concurrency, header-isolation, resource-ownership, or compatibility issues. `AC-1`, `AC-2`, `AC-3`, and `AC-5` pass. `AC-4` is not satisfied because the full repository gate fails outside this routing change; retain `approved` rather than claim all acceptance criteria are implemented. No unrelated fixes or package release are included. + +Commit authorization was subsequently granted: “提交并推送到origin/main上”, “将本地历史一并同步到 origin/main”, and “提交代码”. The initial synchronization attempt was `blocked`: Git requires GPG signing, but the filesystem sandbox denied GPG's write to `~/.gnupg`. The rebase was aborted and all seven task files restored; no commit or push occurred. The backup stash remains available. Git hooks and signing configuration were not changed. + +Commit preparation on 2026-10-26, unchanged baseline `c662cc8d` and the seven-file staged diff: `uv run --extra dev pytest tests/tools/mcp_tool -q`: `pass`, 24 tests in 4.09s. The initial `uv run --extra dev pre-commit run --all-files` was `blocked` by virtualenv cache writes outside the sandbox. Re-running with temporary cache configuration, `VIRTUALENV_OVERRIDE_APP_DATA=/tmp/veadk-mcp-precommit-virtualenv uv run --extra dev pre-commit run --all-files`: `pass` for Ruff lint, Ruff format, Gitleaks, runtime dependency policy, and YAML secret scanning. All hooks remained enabled. `git diff --cached --check`: `pass`. These are preparation checks before successful synchronization; required post-synchronization checks remain pending. The previously recorded full-regression failures remain unresolved. + +Documentation verification on 2026-10-26: `pass` for both PRD/spec language pairs, matching requirement/task/acceptance/contract IDs, existing relative link targets, placeholder scan, and trailing whitespace; `git diff --check`: `pass`. These checks cover the four new design files and do not establish runtime behavior. + +Implementation evidence on 2026-10-26: + +- Environment resolved by `uv`: Python 3.11.14, Google ADK 2.7.1, MCP 1.26.0. +- Pre-fix first-session tests: `fail`, two `TimeoutError` results after 2s each, no transport entry, stack at the parent's lock acquisition. +- Post-fix first-session tests after fixture isolation correction: `pass`, two tests in 1.93s. The initial post-fix fixture exposed Google auth credential autodiscovery; the final fixture explicitly disables it and corrects memory-stream directions. +- Targeted final checks: `uv run --extra dev pytest tests/tools/mcp_tool -q`: `pass`, 24 tests, including 19 new tests and five retained component tests. The new fixture replaces external transports and trusted attestation with memory streams; ADK `SessionContext`, MCP initialization, local MCP server, discovery, and tool invocation are real. +- Initial combined run found an obsolete ordinary-cache tuple and a missing timeout argument in the trusted test fixture. Both test fixtures were corrected; no production workaround was added. +- `uv run --extra dev ruff check ...`: `blocked`, Ruff executable absent from the dev extra. Resolved without dependency metadata changes using the repository pre-commit version: `uvx --from ruff==0.11.12 ruff check veadk/tools/mcp_tool/trusted_mcp_session_manager.py tests/tools/mcp_tool/test_trusted_mcp_components.py tests/tools/mcp_tool/test_trusted_mcp_session_routing.py`: `pass`. The same files with `ruff format --check`: `pass`, all three already formatted. +- Existing upstream deprecation warnings remain; they do not indicate failed tests. +- Baseline comparison with the final isolated fixture: loaded the session manager source from `git show HEAD:veadk/tools/mcp_tool/trusted_mcp_session_manager.py` into a separate Python process and substituted only the toolset's manager class. `pytest.main(["tests/tools/mcp_tool/test_trusted_mcp_session_routing.py", "-k", "standard_first_session", "-q", "--tb=short"])`: `fail`, two expected lock timeouts, 17 deselected, 4.31s. This confirms the final fixture detects the original defect without changing working-tree files. + +Full-regression findings: + +- `uv run --extra dev pytest -n 2 -m "not codex_smoke and not piagent_smoke"`: `fail`, 7556 passed, 17 failed, 58 skipped, 4 xfailed, 52 warnings, 2 collection errors in 436.40s. All affected MCP tests passed. Skipped and xfailed tests are not evidence of passed behavior. +- Ten Codex runtime failures require the absent `openai_codex` optional package. Two self-host sandbox collection errors require the absent `anthropic` optional package. These dependencies are unrelated to MCP routing. +- Three failures attempted writes outside the workspace (Studio state and sibling application output); filesystem sandboxing rejected those paths. No broader filesystem access or machine configuration changes were requested. +- Four assertions failed in unchanged tests: Markdown artifact content type, two harness model configuration checks, and a sidecar child environment check. Avoid recording actual local model identifiers or raw environment values. +- Baseline confirmation: substitute the original manager class into both manager/toolset modules in a separate process, then run the following node IDs with `pytest.main([... , "-q", "--tb=no", "-o", "cache_dir=/tmp/veadk-mcp-baseline-pytest-cache"])`: all five failed in 6.87s, reproducing the corresponding failures with the pre-change manager: + - `tests/frontend/server/runtime_artifacts/test_writer.py::test_success_uses_trusted_scope_and_hashes_utf8_bytes` + - `tests/runtime/codex/test_codex_runtime.py::test_turn_usage_reports_this_turn_not_the_resumed_threads_total` + - `tests/cloud/test_harness_app_http.py::test_harness_app_exposes_agent_info` + - `tests/cloud/test_harness_app_http.py::test_harness_session_create_accepts_id_and_get_agent_config` + - `tests/extensions/harness/sidecar_runtime/test_sidecar.py::test_run_wraps_child_with_sidecar_environment` +- A terminal polling retry restarted the full command after its sandbox-rejected result; the redundant run was interrupted rather than repeated. The completed 436.40s run above is the recorded gate result. +- User docs/examples were assessed: no API or configuration migration is needed for the existing trusted examples. No published version is claimed to contain this working-tree fix. Real trusted attestation and the reporter's exact dependency combination remain unverified; trusted final-close compatibility remains the documented pre-existing limitation. diff --git a/prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.zh.md b/prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.zh.md new file mode 100644 index 000000000..584739d0f --- /dev/null +++ b/prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.zh.md @@ -0,0 +1,135 @@ +# 修复普通 MCP 会话死锁 + +[English](2026-10-26-standard-mcp-deadlock.md) + +- 变更 ID:`trusted-mcp-session-routing` +- 创建/修订日期:2026-10-26 +- 状态:`approved` +- 组件契约:[可信 MCP 会话](../../../specs/trusted-mcp-sessions/README.zh.md) +- 基线版本:`c662cc8d`;验证覆盖工作区会话管理器及两个相关测试文件。 + +## 背景与证据 + +用户报告的环境为 VeADK 1.0.2、Google ADK 2.1.0、MCP 1.26.0、Python 3.12。普通 streamable HTTP 首次创建会话时,包括使用 Bearer 鉴权但未带 `x-trusted-mcp` 的情况,会在连接前无限等待。 + +`veadk/tools/mcp_tool/trusted_mcp_session_manager.py` 的 `create_session` 获取 `_session_lock` 后,在普通连接分支调用 `super().create_session(headers)`。父类也会获取同一把不可重入锁。因此,缓存未命中时,父类等待的锁正由调用它的子类持有。已有健康缓存的会话可能在到达缺陷前返回。可信分支自行创建会话,不会嵌套调用父类。 + +最初安装的本地父类实现确认了重复获取锁的机制;当时环境为 Google ADK 1.18.0、MCP 1.20.0。仓库元数据要求 `google-adk>=2.7.0,<2.8.0` 和 `mcp==1.26.0`。最初的代码检查是根因证据,不代表当时已经运行复现,也不代表验证了报告中的依赖组合。运行检查现已使用 `uv` 按仓库要求解析的环境,结果见下文。 + +现有 `tests/tools/mcp_tool/test_trusted_mcp_components.py` 覆盖可信会话创建和健康缓存复用,缺少普通连接首次创建的测试。 + +## 目标与非目标 + +恢复 `TrustedMcpToolset` 的普通 MCP 连接创建,保留可信路由及按鉴权请求头划分的会话池,消除嵌套加锁。 + +本次不涉及依赖升级、新锁实现、传输重构、扩大可信会话取消清理范围、发布或版本更新、云端调用,以及无关的工具集配置转发。公共导入路径和函数签名保持兼容。 + +## 场景与需求 + +| ID | 前提/触发 | 必须达到的结果 | +| --- | --- | --- | +| `FR-1` | 普通 streamable HTTP、SSE 或 stdio 首次请求会话 | 在获取子类会话锁之前委托父类,连接创建能够继续执行。 | +| `FR-2` | Streamable HTTP 且合并后的 `x-trusted-mcp` 严格等于 `"true"` | 使用 `trusted_mcp_client`,保留现有可信会话池和清理行为。 | +| `FR-3` | 同时存在连接请求头和调用请求头,包括可信标记覆盖 | 遵循父类合并优先级。路由检查不修改调用方的两个字典。传给所选传输的鉴权请求头保持正确。 | +| `FR-4` | 普通连接并发请求、重复请求或已有缓存断开 | 由父类负责串行化、复用、替换和资源清理,不重复执行这些操作。 | +| `FR-5` | 普通连接失败,或等待锁时被取消 | 传播失败或取消;释放锁,使后续请求能够继续。不吞错,不新增超时行为。 | + +标记缺失、为 `"false"` 或其他不匹配值时选择普通路径。SSE 和 stdio 即使带标记,也不能启用可信传输。测试使用模拟鉴权值和本地传输替身。 + +## 设计与契约影响 + +推荐方案:在 `_merge_headers(headers)` 之后,用 `.get` 检查连接类型和标记,检查位置位于生成会话键和获取 `_session_lock` 之前。普通路径立即返回 `await super().create_session(headers)`,传入原始调用请求头,由父类完成正常合并和生命周期管理。 + +仅可信路径执行子类现有的会话键、锁、缓存、断开清理和 `AsyncExitStack` 操作。先生成可信会话键,再从内部合并副本中移除可信标记,供传输使用。移除子类锁内已经不可达的普通分支。 + +2026-10-26 按受支持基线复核实现时发现,Google ADK 2.7.1 调用 `_create_client` 会额外传入 `mtls_transport` 和 `session_key`。现有重写方法需要接收并转发父类参数,使普通连接在修复加锁后能够继续创建。可信路由行为保持不变。这是满足 `FR-1` 所需的私有接口适配,不引入新的公共契约或产品范围。现有普通缓存测试也需要用当前父类的 `(session, exit_stack, loop)` 元组替换过期的双元素测试数据。 + +其他方案: + +1. 把锁换成异步重入锁:引入自定义锁行为,并与父类锁的管理职责冲突,没有必要。 +2. 在子类复制父类会话创建流程:能避开嵌套加锁,但会重复传输及生命周期代码,增加上游兼容风险。 + +本次调整普通会话的并发管理职责,因此需要范围明确的双语组件契约。Studio、生成代码、依赖元数据、权限、事件结构和部署契约均不受影响。调用方无需迁移。交付核对时评估现有用户文档;修正超出本次路由缺陷的通用文档错误,需要另行确定范围。 + +## 任务与受影响文件 + +| 任务 | 工作 | 对应需求 | +| --- | --- | --- | +| `T-1` | 审查双语 PRD 与 `specs/trusted-mcp-sessions/README{,.zh}.md`,记录用户批准。 | 全部 | +| `T-2` | 新增 `tests/tools/mcp_tool/test_trusted_mcp_session_routing.py`。先使用真实父类会话方法、传输/会话替身和有界 `asyncio.wait_for`,证明修复前普通连接创建会超时。 | `FR-1` | +| `T-3` | 在 `veadk/tools/mcp_tool/trusted_mcp_session_manager.py` 中实施入口分流及父类客户端参数转发。 | `FR-1`、`FR-2`、`FR-3` | +| `T-4` | 覆盖普通传输类型、标记值/覆盖、模拟 Bearer 请求头、可信创建/复用、并发创建、断开替换、失败重试和等待锁时取消。保留现有组件测试,并按父类会话池元组更新普通缓存测试数据。 | `FR-2`–`FR-5` | +| `T-5` | 执行定向及扩大回归检查,审查差异、资源职责和双语一致性,按实际证据核对文档。 | 全部 | + +顺序:`T-1` → `T-2` 的失败回归 → `T-3` → `T-4` 的其余覆盖 → `T-5`。实施前记录修复前的失败证据。 + +执行步骤(已批准范围): + +- [x] `T-1`:用户于 2026-10-26 批准代码和测试修改:“按这个方案修改代码和测试”。 +- [x] `T-2`:构建隔离传输夹具,保留真实父类加锁和会话池。运行 `uv run --extra dev pytest tests/tools/mcp_tool/test_trusted_mcp_session_routing.py -k standard_first_session`,预期在进入传输前有界超时。 +- [x] `T-3`:在生成会话键前插入 `if not (isinstance(self._connection_params, StreamableHTTPConnectionParams) and merged_headers and merged_headers.get("x-trusted-mcp") == "true"): return await super().create_session(headers)`。保留可信会话键,在进入传输前从合并副本移除标记,删除锁内父类回退。再次运行首次会话测试,预期成功。 +- [x] `T-4`:使用隔离传输/会话上下文运行全部路由和组件测试,检查创建次数、对象复用、请求头值、清理、失败传播和取消后继续执行。 +- [x] `T-5`:执行下文验证命令,按 `FR-1`–`FR-5` 审查最终源码/测试差异,同步双语和组件状态,记录实际结果。用户后续已授权提交本次修改,并将本地历史同步到 `origin/main`。 + +## 验证与验收 + +使用 `uv run --extra dev` 和当前仓库元数据。不能把过期的本地环境当成受支持基线。测试替换网络/进程传输,不替换父类加锁实现。 + +| 需求 | 任务 | 验收条件 | 检查 | 当前结果 | +| --- | --- | --- | --- | --- | +| `FR-1` | `T-2`、`T-3` | `AC-1`:普通连接首次创建在修复前有界超时,修复后成功 | `uv run --extra dev pytest tests/tools/mcp_tool/test_trusted_mcp_session_routing.py -k standard_first_session -q` | `pass`:修复前两条在嵌套加锁处超时,修复后两条成功 | +| `FR-2`、`FR-3` | `T-3`、`T-4` | `AC-2`:可信路由、合并优先级、调用方请求头和鉴权保持正确 | `uv run --extra dev pytest tests/tools/mcp_tool -q` | `pass`:包含在 24 条通过的测试中 | +| `FR-4`、`FR-5` | `T-4` | `AC-3`:普通生命周期、并发复用、失败重试和等待锁时取消通过 | `uv run --extra dev pytest tests/tools/mcp_tool -q` | `pass`:24 条测试通过,耗时 3.05s | +| 全部 | `T-5` | `AC-4`:扩大后的相关回归通过 | `uv run --extra dev pytest -n 2 -m "not codex_smoke and not piagent_smoke"` | `fail`:7556 条通过、17 条失败、58 条跳过、4 条 xfailed、2 条收集错误;436.40s;两个 worker | +| 全部 | `T-5` | `AC-5`:修改的 Python 及文档检查通过 | 下文记录中的 Ruff 0.11.12 lint/格式命令;`git diff --check`;双语/路径/链接审查 | `pass`:三个 Python 文件和两组双语文档 | + +用户后续已授权提交并推送到 `origin/main`,包括同步本地历史,因此提交前必须运行 pre-commit。创建 PR、软件包发布和部署尚未获授权。真实云端/服务商及可信证明检查为 `not_applicable`:缺陷发生在进入传输之前的锁管理中,验证使用隔离环境。现有可信组件测试只能证明模拟兼容,不能证明真实可信证明流程。 + +## 风险与未决问题 + +- Google ADK 私有方法可能变化,需要按声明的 2.7.x 基线验证。额外测试报告中的 2.1.0 组合只能补充兼容证据,不代表改变依赖策略或作出发布保证。 +- 普通路径委托后,父类会再次合并请求头。第一次合并只用于路由检查,必须没有副作用。 +- 父类在资源进入后的传输失败及取消,保留上游语义。本次保证子类不会额外持锁,不承诺修复上游清理问题。 +- 现有可信分支捕获 `Exception`,未覆盖所有取消情况。扩大清理范围不属于本次死锁修复。 +- Google ADK 2.7.1 的父类会话池使用三个元素,未改动的可信分支仍保存两个元素。可信分支继承的最终 `close()` 兼容问题是既有局限,不属于本次普通路由修复。可信模拟测试会显式关闭它们持有的栈,不声称已验证继承的最终关闭兼容性。 +- 没有未决产品选择阻塞设计。受支持环境准备及定向行为已验证;报告中的 Google ADK 2.1.0/Python 3.12 组合尚未测试。 + +## 审查与交付记录 + +2026-10-26,基线 `c662cc8d`,实际差异范围:会话管理器、新增隔离路由测试、现有普通缓存测试数据,以及本次 PRD/组件规格双语文件。 + +提供的技能目录没有 `review-spec`,因此直接完成同等设计审查。已检查职责、可行性、接口、错误、安全、兼容、可测试性和双语一致性。设计中已解决的问题:先分流后加锁;检查请求头时不修改它;移除标记前生成可信会话键;使用声明的依赖基线;区分源码检查和运行证据。没有未解决的设计阻塞项。 + +用户已于 2026-10-26 批准:“按这个方案修改代码和测试”。范围内的源码/测试修改和文档核对已完成。最终实现审查未发现范围内的正确性、并发、请求头隔离、资源职责或兼容阻塞项。`AC-1`、`AC-2`、`AC-3`、`AC-5` 通过。由于完整仓库检查在本次路由变更之外失败,`AC-4` 未满足;保留 `approved`,不声称所有验收条件均已实现。未包含无关修复或软件包发布。 + +用户随后授权提交:“提交并推送到origin/main上”、“将本地历史一并同步到 origin/main”和“提交代码”。首次同步尝试为 `blocked`:Git 要求 GPG 签名,但文件系统沙箱拒绝 GPG 写入 `~/.gnupg`。已撤销 rebase,恢复全部七个任务文件;尚未提交或推送。备份 stash 仍保留。Git hooks 和签名配置均未更改。 + +2026-10-26 提交准备,基线仍为 `c662cc8d`,范围为暂存的七个文件:`uv run --extra dev pytest tests/tools/mcp_tool -q` 为 `pass`,24 条测试通过,耗时 4.09s。首次 `uv run --extra dev pre-commit run --all-files` 因 virtualenv 尝试写入沙箱之外的缓存目录而 `blocked`。改用临时缓存后,`VIRTUALENV_OVERRIDE_APP_DATA=/tmp/veadk-mcp-precommit-virtualenv uv run --extra dev pre-commit run --all-files` 为 `pass`:Ruff 检查、Ruff 格式、Gitleaks、运行时依赖策略和 YAML 密钥扫描全部通过。所有 hooks 均保留。`git diff --cached --check` 为 `pass`。这些是成功同步之前的准备检查;同步后的必需检查仍待执行。此前记录的完整回归失败仍未解决。 + +2026-10-26 文档验证:PRD/规格双语配对、需求/任务/验收/契约 ID 一致性、相对链接目标存在性、占位符扫描和行尾空白检查均为 `pass`;`git diff --check` 为 `pass`。这些检查覆盖四个新增设计文件,不证明运行行为。 + +2026-10-26 实施证据: + +- `uv` 解析后的环境:Python 3.11.14、Google ADK 2.7.1、MCP 1.26.0。 +- 修复前首次会话测试:`fail`,两条各在 2s 后得到 `TimeoutError`,未进入传输,堆栈停在父类获取锁处。 +- 修复后首次会话测试,在修正夹具隔离后:`pass`,两条测试耗时 1.93s。首次修复后运行暴露了 Google auth 自动发现凭据的行为;最终夹具显式禁用此行为,并修正内存流方向。 +- 最终定向检查:`uv run --extra dev pytest tests/tools/mcp_tool -q` 为 `pass`,24 条测试,包括 19 条新增测试及保留的五条组件测试。新夹具用内存流替换外部传输和可信证明;ADK `SessionContext`、MCP 初始化、本地 MCP 服务、工具发现和调用均使用真实实现。 +- 首次组合运行发现普通缓存元组过期,以及可信测试夹具漏传超时参数。均已修正测试夹具,没有增加生产绕过逻辑。 +- `uv run --extra dev ruff check ...` 为 `blocked`,dev extra 未提供 Ruff 可执行文件。未修改依赖元数据,改用仓库 pre-commit 的版本:`uvx --from ruff==0.11.12 ruff check veadk/tools/mcp_tool/trusted_mcp_session_manager.py tests/tools/mcp_tool/test_trusted_mcp_components.py tests/tools/mcp_tool/test_trusted_mcp_session_routing.py` 为 `pass`。相同文件的 `ruff format --check` 为 `pass`,三个文件均已符合格式。 +- 仍有上游既有弃用警告,不代表测试失败。 +- 使用最终隔离夹具对照基线:在独立 Python 进程中加载 `git show HEAD:veadk/tools/mcp_tool/trusted_mcp_session_manager.py` 的会话管理器源码,仅替换工具集使用的管理器类。`pytest.main(["tests/tools/mcp_tool/test_trusted_mcp_session_routing.py", "-k", "standard_first_session", "-q", "--tb=short"])` 为 `fail`,两条均出现预期的锁超时,17 条未选中,耗时 4.31s。这确认最终夹具能检测原始缺陷,且没有修改工作区文件。 + +完整回归发现: + +- `uv run --extra dev pytest -n 2 -m "not codex_smoke and not piagent_smoke"` 为 `fail`,7556 条通过、17 条失败、58 条跳过、4 条 xfailed、52 条警告、2 条收集错误,耗时 436.40s。相关 MCP 测试全部通过。跳过及 xfailed 的测试不证明行为通过。 +- 十条 Codex 运行时失败需要尚未安装的可选包 `openai_codex`;两条自托管沙箱收集错误需要尚未安装的可选包 `anthropic`。这些依赖与 MCP 路由无关。 +- 三条失败尝试写入工作区之外的路径,包括 Studio 状态和同级应用输出;文件系统沙箱拒绝了这些路径。未请求扩大文件系统访问范围,也未修改机器配置。 +- 未改动测试中的四个断言失败:Markdown 工件内容类型、两条 harness 模型配置检查,以及一条 sidecar 子进程环境检查。不记录真实本地模型标识符或原始环境值。 +- 基线确认:在独立进程中,把原始管理器类替换进管理器/工具集两个模块,再用 `pytest.main([... , "-q", "--tb=no", "-o", "cache_dir=/tmp/veadk-mcp-baseline-pytest-cache"])` 运行下列节点。五条全部失败,耗时 6.87s,在变更前管理器下复现了对应失败: + - `tests/frontend/server/runtime_artifacts/test_writer.py::test_success_uses_trusted_scope_and_hashes_utf8_bytes` + - `tests/runtime/codex/test_codex_runtime.py::test_turn_usage_reports_this_turn_not_the_resumed_threads_total` + - `tests/cloud/test_harness_app_http.py::test_harness_app_exposes_agent_info` + - `tests/cloud/test_harness_app_http.py::test_harness_session_create_accepts_id_and_get_agent_config` + - `tests/extensions/harness/sidecar_runtime/test_sidecar.py::test_run_wraps_child_with_sidecar_environment` +- 终端轮询重试在沙箱拒绝结果后重新启动了完整命令;已中断这次重复运行,不重复执行。记录的检查结果来自上文已经完成的 436.40s 运行。 +- 已评估用户文档/示例:现有可信示例无需 API 或配置迁移。不声称任何已发布版本包含本次工作区修复。真实可信证明及报告中的精确依赖组合仍未验证;可信最终关闭兼容性仍是已记录的既有局限。 diff --git a/specs/trusted-mcp-sessions/README.md b/specs/trusted-mcp-sessions/README.md new file mode 100644 index 000000000..4c72e919f --- /dev/null +++ b/specs/trusted-mcp-sessions/README.md @@ -0,0 +1,57 @@ +# Trusted MCP sessions + +[中文版](README.zh.md) + +- Component ID: `trusted-mcp-sessions` +- Status: `active`; scoped routing contracts verified by isolated component tests. +- Revision: 2026-10-26 +- Related PRD: [Standard MCP deadlock](../../prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.md) +- Owned code: `veadk/tools/mcp_tool/trusted_mcp_session_manager.py`. +- Consumer: `veadk/tools/mcp_tool/trusted_mcp_toolset.py`. +- Tests: `tests/tools/mcp_tool/test_trusted_mcp_components.py`; `tests/tools/mcp_tool/test_trusted_mcp_session_routing.py`. + +## Responsibility and entry points + +`TrustedMcpSessionManager` extends Google ADK's `MCPSessionManager`. It owns trusted transport selection and trusted session creation/pooling; standard sessions delegate to Google ADK. `TrustedMcpToolset` installs this manager for tool discovery and execution. + +The session entry point is `async create_session(headers: Optional[Dict[str, str]] = None) -> ClientSession`. Dependencies are Google ADK connection parameters and parent lifecycle methods, MCP `ClientSession`, and `bytedance.jeddak_trusted_mcp.trusted_mcp_client`. Inherited `close()` owns standard pool cleanup; the existing trusted-entry compatibility limitation is recorded below. Toolset argument forwarding, trusted attestation internals, cloud deployment, and general dependency policy are outside this component's scope. + +## Contracts + +- `CON-1`: Trusted selection requires `StreamableHTTPConnectionParams` and merged `x-trusted-mcp == "true"`. Other connection types and marker values use standard transport. +- `CON-2`: Standard `create_session` delegates before acquiring the subclass session lock or touching the pool. The parent owns standard synchronization, initialization, reuse, replacement, errors, and cleanup. The `_create_client` override accepts and forwards parent positional/keyword arguments on the standard route. +- `CON-3`: Header merge precedence is inherited: per-call values override connection values. Routing inspection must not mutate input dictionaries. Authentication headers remain available to the selected transport. The trusted marker is removed from the private merged copy only after computing the trusted session key. +- `CON-4`: Trusted creation uses `trusted_mcp_client` and stores the initialized session with its `AsyncExitStack`. Healthy sessions reuse the same key; disconnected sessions are removed and replaced under the session lock. Existing cleanup failures are reported by the current logger. +- `CON-5`: Standard errors and cancellations propagate with parent semantics. Cancellation while waiting for the lock must not leave the subclass holding it. No new timeout, retry policy, or success fallback is introduced. + +Example inputs: a streamable HTTP connection with `headers={"Authorization": "Bearer test-placeholder"}` selects standard MCP. Adding `x-trusted-mcp: "true"` selects trusted MCP; a per-call `"false"` override selects standard MCP. + +## State, data, and concurrency + +Session state is process-local in inherited `_sessions`: the current parent stores key → `(session, exit_stack, loop)` for standard sessions; the unchanged trusted branch stores key → `(session, exit_stack)`. There is no persistent schema or new environment variable. Session keys and connection timeouts use the parent implementation. + +Relevant transitions are absent → initialized/cached → reused, disconnected → cleaned/removed → replaced, and cached → closed/removed through inherited `close()`. Standard transitions belong to the parent. Trusted transitions use the subclass lock exactly once. A failed creation does not cache a successful session. Lock-wait cancellation releases any context-managed lock already acquired. + +This change does not adjust event-loop lock allocation, cross-loop session handling, late transport events, or cleanup on trusted cancellation after resource entry; these remain upstream/existing implementation limitations. + +On the verified Google ADK 2.7.1 baseline, inherited `close()` expects three-element pool entries and is incompatible with the unchanged trusted two-element entries. The scoped fix preserves trusted creation/reuse; tests explicitly close trusted stacks without claiming inherited final-close compatibility. Resolving this existing trusted lifecycle issue requires a separate approved change. + +## Security, failures, and compatibility + +Trust is selected by explicit connection type and marker, not by a Bearer header alone. Do not add credential values to diagnostics, specifications, or fixtures. Verification replaces real transports and does not read local AICC configuration or invoke attestation/cloud services. + +Public imports/signatures and authentication configuration remain compatible. Dependency support follows [runtime dependency compatibility](../runtime-dependency-compatibility/README.md), currently Google ADK `>=2.7.0,<2.8.0` and repository MCP `1.26.0`. The reported Google ADK 2.1.0 combination is source context, not verified current support. + +Existing session creation/cleanup logs remain the available diagnostics. The standard path removes the subclass-held lock that previously prevented connection entry; transport failures remain visible as parent exceptions. It does not guarantee successful connection to an unavailable server or repair unrelated upstream cancellation cleanup. + +## Verification and change record + +| Contract | Validation | Evidence | +| --- | --- | --- | +| `CON-1`, `CON-3`, `CON-4` | Routing, header, trusted creation/reuse tests plus existing component tests | `pass`, included in 24 targeted tests; related PRD | +| `CON-2` | Bounded pre-fix failure and post-fix success using real parent locking; standard transport/lifecycle tests | `pass`, two pre-fix timeouts and post-fix successes; related PRD | +| `CON-5` | Standard failure retry and cancellation while waiting for the lock | `pass`, included in 24 targeted tests; related PRD | + +Command: `uv run --extra dev pytest tests/tools/mcp_tool`. Broader checks and final results belong to the PRD. Real trusted attestation is `not_applicable` to the isolated lock regression. + +2026-10-26: User approved the scoped PRD; implemented standard-session delegation and parent factory-argument forwarding while preserving trusted routing. Verified on Python 3.11.14, Google ADK 2.7.1, and MCP 1.26.0 against baseline `c662cc8d` plus the working-tree routing diff. Targeted tests passed; broader verification and limitations are recorded in the PRD. diff --git a/specs/trusted-mcp-sessions/README.zh.md b/specs/trusted-mcp-sessions/README.zh.md new file mode 100644 index 000000000..1facc2db0 --- /dev/null +++ b/specs/trusted-mcp-sessions/README.zh.md @@ -0,0 +1,57 @@ +# 可信 MCP 会话 + +[English](README.md) + +- 组件 ID:`trusted-mcp-sessions` +- 状态:`active`;范围内路由契约已通过隔离组件测试验证。 +- 修订日期:2026-10-26 +- 相关 PRD:[普通 MCP 死锁](../../prd-spec/bugfixes/trusted-mcp-session-routing/2026-10-26-standard-mcp-deadlock.zh.md) +- 所属代码:`veadk/tools/mcp_tool/trusted_mcp_session_manager.py`。 +- 调用方:`veadk/tools/mcp_tool/trusted_mcp_toolset.py`。 +- 测试:`tests/tools/mcp_tool/test_trusted_mcp_components.py`;`tests/tools/mcp_tool/test_trusted_mcp_session_routing.py`。 + +## 职责与入口 + +`TrustedMcpSessionManager` 继承 Google ADK 的 `MCPSessionManager`,负责可信传输选择、可信会话创建和会话池;普通会话委托 Google ADK。`TrustedMcpToolset` 安装此管理器,用于工具发现和执行。 + +会话入口为 `async create_session(headers: Optional[Dict[str, str]] = None) -> ClientSession`。依赖包括 Google ADK 连接参数和父类生命周期方法、MCP `ClientSession`,以及 `bytedance.jeddak_trusted_mcp.trusted_mcp_client`。继承的 `close()` 负责普通会话池清理;可信条目的既有兼容局限见下文。工具集参数转发、可信证明内部实现、云端部署和通用依赖策略不属于本组件范围。 + +## 契约 + +- `CON-1`:选择可信路径必须同时满足 `StreamableHTTPConnectionParams` 和合并后 `x-trusted-mcp == "true"`。其他连接类型及标记值使用普通传输。 +- `CON-2`:普通 `create_session` 在获取子类会话锁或访问会话池之前委托父类。普通会话的同步、初始化、复用、替换、错误和清理由父类负责。`_create_client` 重写方法在普通路径接收并转发父类位置/关键字参数。 +- `CON-3`:继承父类请求头合并优先级,调用请求头覆盖连接请求头。路由检查不得修改输入字典。所选传输须能获得鉴权请求头。可信标记只在生成可信会话键后从内部合并副本移除。 +- `CON-4`:可信创建使用 `trusted_mcp_client`,缓存已初始化会话及其 `AsyncExitStack`。同一键复用健康会话;断开的会话在会话锁内移除并替换。清理失败沿用现有日志报告。 +- `CON-5`:普通错误及取消按父类语义传播。等待锁时取消,不得使子类继续持锁。不新增超时、重试策略或成功回退。 + +输入示例:streamable HTTP 连接使用 `headers={"Authorization": "Bearer test-placeholder"}` 时选择普通 MCP。增加 `x-trusted-mcp: "true"` 后选择可信 MCP;调用请求头用 `"false"` 覆盖时选择普通 MCP。 + +## 状态、数据与并发 + +会话状态仅保存在进程内继承的 `_sessions`:当前父类为普通会话保存键 → `(session, exit_stack, loop)`;未改动的可信分支保存键 → `(session, exit_stack)`。没有持久化结构或新增环境变量。会话键及连接超时沿用父类实现。 + +相关转换为:不存在 → 已初始化/缓存 → 复用;断开 → 清理/移除 → 替换;以及通过继承的 `close()` 从缓存 → 关闭/移除。普通转换由父类管理。可信转换只获取一次子类锁。创建失败不能缓存为成功会话。等待锁时取消会释放已经通过上下文管理器获取的锁。 + +本次不调整事件循环锁分配、跨循环会话处理、迟到的传输事件,以及可信资源进入后的取消清理;这些仍受上游或现有实现限制。 + +在已验证的 Google ADK 2.7.1 基线上,继承的 `close()` 要求三元素会话池条目,与未改动的可信双元素条目不兼容。本次修复保留可信创建/复用;测试显式关闭可信栈,不声称继承的最终关闭已兼容。解决这个既有可信生命周期问题,需要另行批准变更。 + +## 安全、失败与兼容 + +可信选择依据明确的连接类型和标记,不能仅凭 Bearer 请求头启用。诊断、规格和测试数据不得新增真实凭据。验证替换真实传输,不读取本地 AICC 配置,不调用可信证明或云端服务。 + +公共导入路径/签名及鉴权配置保持兼容。依赖支持遵循[运行时依赖兼容](../runtime-dependency-compatibility/README.zh.md),当前为 Google ADK `>=2.7.0,<2.8.0` 和仓库 MCP `1.26.0`。报告中的 Google ADK 2.1.0 组合只是问题背景,不代表当前已验证支持。 + +诊断继续使用现有会话创建/清理日志。普通路径已移除原先阻止连接开始的子类持锁;传输失败仍表现为父类异常。本次不保证连接不可用服务器成功,也不修复无关的上游取消清理问题。 + +## 验证与变更记录 + +| 契约 | 验证 | 证据 | +| --- | --- | --- | +| `CON-1`、`CON-3`、`CON-4` | 路由、请求头、可信创建/复用测试及现有组件测试 | `pass`,包含在 24 条定向测试中;见相关 PRD | +| `CON-2` | 使用真实父类加锁,验证修复前有界失败、修复后成功;普通传输/生命周期测试 | `pass`,修复前两条超时,修复后成功;见相关 PRD | +| `CON-5` | 普通失败重试及等待锁时取消 | `pass`,包含在 24 条定向测试中;见相关 PRD | + +命令:`uv run --extra dev pytest tests/tools/mcp_tool`。扩大检查及最终结果记录在 PRD。真实可信证明对本次隔离锁回归为 `not_applicable`。 + +2026-10-26:用户批准范围明确的 PRD;已实现普通会话委托和父类工厂参数转发,保留可信路由。验证环境为 Python 3.11.14、Google ADK 2.7.1、MCP 1.26.0,检查范围为基线 `c662cc8d` 加工作区路由差异。定向测试通过;扩大验证及局限记录在 PRD。 diff --git a/tests/tools/mcp_tool/test_trusted_mcp_components.py b/tests/tools/mcp_tool/test_trusted_mcp_components.py index 243c283e4..7bdb2894f 100644 --- a/tests/tools/mcp_tool/test_trusted_mcp_components.py +++ b/tests/tools/mcp_tool/test_trusted_mcp_components.py @@ -213,7 +213,11 @@ async def run_test(): existing_session = mock.MagicMock() existing_exit_stack = mock.MagicMock() manager._sessions = { - "session-key": (existing_session, existing_exit_stack) + "session-key": ( + existing_session, + existing_exit_stack, + asyncio.get_running_loop(), + ) } # Call create_session diff --git a/tests/tools/mcp_tool/test_trusted_mcp_session_routing.py b/tests/tools/mcp_tool/test_trusted_mcp_session_routing.py new file mode 100644 index 000000000..3d81fab9b --- /dev/null +++ b/tests/tools/mcp_tool/test_trusted_mcp_session_routing.py @@ -0,0 +1,340 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from contextlib import asynccontextmanager +from functools import partial +from types import SimpleNamespace + +import anyio +import pytest +from google.adk.tools.mcp_tool import mcp_session_manager as adk_sessions +from google.auth.exceptions import DefaultCredentialsError +from mcp import StdioServerParameters +from mcp.server.lowlevel import Server +from mcp.types import TextContent, Tool + +from veadk.tools.mcp_tool import trusted_mcp_session_manager as trusted_sessions +from veadk.tools.mcp_tool.trusted_mcp_session_manager import TrustedMcpSessionManager +from veadk.tools.mcp_tool.trusted_mcp_toolset import TrustedMcpToolset + +try: + from asyncio import timeout +except ImportError: # Python 3.10 + from async_timeout import timeout + +SESSION_TIMEOUT = 2 + + +@pytest.fixture +def memory_transports(monkeypatch): + """Replace external transports, retaining real ADK and MCP sessions.""" + state = SimpleNamespace(opened=[], closed=0, failure=None) + server = Server("session-routing-test") + + def no_credentials(**kwargs): + raise DefaultCredentialsError("Credentials are disabled in this test") + + monkeypatch.setattr(adk_sessions.google.auth, "default", no_credentials) + + @server.list_tools() + async def list_tools(): + return [ + Tool( + name="echo", + inputSchema={ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + ) + ] + + @server.call_tool() + async def call_tool(name, arguments): + return [TextContent(type="text", text=arguments["value"])] + + @asynccontextmanager + async def transport(**kwargs): + state.opened.append(kwargs) + try: + if state.failure is not None: + raise state.failure + server_write, client_read = anyio.create_memory_object_stream(10) + client_write, server_read = anyio.create_memory_object_stream(10) + async with ( + client_read, + client_write, + server_read, + server_write, + anyio.create_task_group() as tasks, + ): + tasks.start_soon( + server.run, + server_read, + server_write, + server.create_initialization_options(), + ) + try: + yield client_read, client_write, lambda: "test-session" + finally: + tasks.cancel_scope.cancel() + finally: + state.closed += 1 + if "http_client" in kwargs: + await kwargs["http_client"].aclose() + + def trusted_client(**kwargs): + return adk_sessions.SessionContext( + client=transport(kind="trusted", **kwargs), + timeout=SESSION_TIMEOUT, + sse_read_timeout=SESSION_TIMEOUT, + ) + + monkeypatch.setattr( + adk_sessions, "_StreamableHttpClientWrapper", partial(transport, kind="http") + ) + monkeypatch.setattr(adk_sessions, "sse_client", partial(transport, kind="sse")) + monkeypatch.setattr(adk_sessions, "stdio_client", partial(transport, kind="stdio")) + monkeypatch.setattr(trusted_sessions, "trusted_mcp_client", trusted_client) + return state + + +@asynccontextmanager +async def managed_manager(params): + manager = TrustedMcpSessionManager(params) + try: + yield manager + finally: + if any(len(entry) == 2 for entry in manager._sessions.values()): + # Trusted entries retain their existing two-element pool shape. + for entry in list(manager._sessions.values()): + await entry[1].aclose() + manager._sessions.clear() + else: + await manager.close() + + +async def create_session(manager, headers=None): + async with timeout(SESSION_TIMEOUT): + return await manager.create_session(headers) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "headers", [None, {"Authorization": "Bearer test-placeholder"}] +) +async def test_standard_first_session(memory_transports, headers): + toolset = TrustedMcpToolset( + connection_params=adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp", headers=headers + ) + ) + manager = toolset._mcp_session_manager + try: + session = await asyncio.wait_for( + manager.create_session(), timeout=SESSION_TIMEOUT + ) + assert not manager._is_session_disconnected(session) + assert len(memory_transports.opened) == 1 + finally: + await manager.close() + assert memory_transports.closed == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["http", "sse", "stdio"]) +async def test_standard_transport_types(memory_transports, kind): + if kind == "stdio": + params = adk_sessions.StdioConnectionParams( + server_params=StdioServerParameters(command="unused-test-command") + ) + else: + params_type = ( + adk_sessions.StreamableHTTPConnectionParams + if kind == "http" + else adk_sessions.SseConnectionParams + ) + params = params_type(url="http://mcp.example.test/mcp") + async with managed_manager(params) as manager: + headers = {"x-trusted-mcp": "true"} if kind != "http" else None + session = await create_session(manager, headers) + assert not manager._is_session_disconnected(session) + assert memory_transports.opened[0]["kind"] == kind + assert memory_transports.closed == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "base_marker,call_marker,expected_kind", + [ + (None, None, "http"), + ("false", None, "http"), + ("TRUE", None, "http"), + ("true", None, "trusted"), + (None, "true", "trusted"), + ("true", "false", "http"), + ("false", "true", "trusted"), + ], +) +async def test_merged_headers_select_route( + memory_transports, base_marker, call_marker, expected_kind +): + base_headers = {"Authorization": "Bearer base-placeholder", "X-Base": "base"} + call_headers = {"Authorization": "Bearer call-placeholder", "X-Call": "call"} + if base_marker is not None: + base_headers["x-trusted-mcp"] = base_marker + if call_marker is not None: + call_headers["x-trusted-mcp"] = call_marker + original_base, original_call = base_headers.copy(), call_headers.copy() + params = adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp", headers=base_headers + ) + async with managed_manager(params) as manager: + session = await create_session(manager, call_headers) + assert await create_session(manager, call_headers) is session + assert len(memory_transports.opened) == 1 + opened = memory_transports.opened[0] + assert opened["kind"] == expected_kind + transport_headers = ( + opened["headers"] + if expected_kind == "trusted" + else opened["http_client"].headers + ) + assert transport_headers["Authorization"] == "Bearer call-placeholder" + assert transport_headers["X-Base"] == "base" + assert transport_headers["X-Call"] == "call" + assert "x-trusted-mcp" not in transport_headers + assert base_headers == original_base + assert call_headers == original_call + assert params.headers == original_base + expected_key = manager._generate_session_key({**base_headers, **call_headers}) + assert manager._sessions[expected_key][0] is session + assert memory_transports.closed == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("trusted", [False, True]) +async def test_authentication_headers_isolate_sessions(memory_transports, trusted): + params = adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp", + headers={"x-trusted-mcp": "true"} if trusted else None, + ) + first_headers = {"Authorization": "Bearer first-placeholder"} + second_headers = {"Authorization": "Bearer second-placeholder"} + async with managed_manager(params) as manager: + first = await create_session(manager, first_headers) + second = await create_session(manager, second_headers) + assert first is not second + assert await create_session(manager, first_headers) is first + assert await create_session(manager, second_headers) is second + assert len(memory_transports.opened) == 2 + assert memory_transports.closed == 2 + + +@pytest.mark.asyncio +async def test_standard_concurrent_requests_reuse_session(memory_transports): + params = adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp" + ) + async with managed_manager(params) as manager: + sessions = await asyncio.wait_for( + asyncio.gather(*(manager.create_session() for _ in range(4))), + timeout=SESSION_TIMEOUT, + ) + assert all(session is sessions[0] for session in sessions) + assert len(memory_transports.opened) == 1 + assert memory_transports.closed == 1 + + +@pytest.mark.asyncio +async def test_standard_disconnected_session_is_replaced(memory_transports): + params = adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp" + ) + async with managed_manager(params) as manager: + first = await create_session(manager) + await first._write_stream.aclose() + second = await create_session(manager) + assert second is not first + assert not manager._is_session_disconnected(second) + assert len(memory_transports.opened) == 2 + assert memory_transports.closed == 1 + assert memory_transports.closed == 2 + + +@pytest.mark.asyncio +async def test_standard_failure_allows_retry(memory_transports): + params = adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp" + ) + async with managed_manager(params) as manager: + memory_transports.failure = RuntimeError("test transport failed") + with pytest.raises(ConnectionError, match="test transport failed"): + await create_session(manager) + assert not manager._sessions + assert not manager._session_lock.locked() + assert memory_transports.closed == 1 + memory_transports.failure = None + session = await create_session(manager) + assert not manager._is_session_disconnected(session) + assert len(memory_transports.opened) == 2 + assert memory_transports.closed == 2 + + +@pytest.mark.asyncio +async def test_standard_lock_wait_cancellation_allows_retry(memory_transports): + params = adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp" + ) + async with managed_manager(params) as manager: + started = asyncio.Event() + + async def request_session(): + started.set() + return await manager.create_session() + + async with manager._session_lock: + task = asyncio.create_task(request_session()) + await asyncio.wait_for(started.wait(), timeout=SESSION_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not memory_transports.opened + assert not manager._sessions + assert not manager._session_lock.locked() + session = await create_session(manager) + assert not manager._is_session_disconnected(session) + assert memory_transports.closed == 1 + + +@pytest.mark.asyncio +async def test_toolset_discovers_and_calls_standard_tool(memory_transports): + toolset = TrustedMcpToolset( + connection_params=adk_sessions.StreamableHTTPConnectionParams( + url="http://mcp.example.test/mcp" + ) + ) + try: + async with timeout(SESSION_TIMEOUT): + tools = await toolset.get_tools() + assert [tool.name for tool in tools] == ["echo"] + result = await tools[0].run_async( + args={"value": "test echo"}, tool_context=None + ) + assert result["content"][0]["text"] == "test echo" + assert len(memory_transports.opened) == 1 + finally: + await toolset.close() + assert memory_transports.closed == 1 diff --git a/veadk/tools/mcp_tool/trusted_mcp_session_manager.py b/veadk/tools/mcp_tool/trusted_mcp_session_manager.py index f6e13e5c8..e772a22ac 100644 --- a/veadk/tools/mcp_tool/trusted_mcp_session_manager.py +++ b/veadk/tools/mcp_tool/trusted_mcp_session_manager.py @@ -41,12 +41,16 @@ class TrustedMcpSessionManager(MCPSessionManager): session pooling based on authentication headers. """ - def _create_client(self, merged_headers: Optional[Dict[str, str]] = None): + def _create_client( + self, merged_headers: Optional[Dict[str, str]] = None, *args, **kwargs + ): """Creates an MCP client based on the connection parameters. Args: merged_headers: Optional headers to include in the connection. Only applicable for SSE and StreamableHTTP connections. + *args: Additional positional arguments for the parent client factory. + **kwargs: Additional keyword arguments for the parent client factory. Returns: The appropriate MCP client instance. @@ -74,7 +78,7 @@ def _create_client(self, merged_headers: Optional[Dict[str, str]] = None): ) else: logger.info("Creating MCP client") - client = super()._create_client(merged_headers) + client = super()._create_client(merged_headers, *args, **kwargs) return client async def create_session( @@ -94,8 +98,16 @@ async def create_session( Returns: ClientSession: The initialized MCP client session. """ - # Merge headers once at the beginning + # Inspect routing without mutating the caller's headers. merged_headers = self._merge_headers(headers) + if not ( + isinstance(self._connection_params, StreamableHTTPConnectionParams) + and merged_headers + and merged_headers.get("x-trusted-mcp") == "true" + ): + # The parent owns standard session locking and pooling. + logger.info("Initialize MCP session") + return await super().create_session(headers) # Generate session key using merged headers session_key = self._generate_session_key(merged_headers) @@ -127,32 +139,24 @@ async def create_session( try: # FIXME: reuse the normal procedure to create a session after updated trusted_mcp_client_context - if ( - isinstance(self._connection_params, StreamableHTTPConnectionParams) - and merged_headers - and merged_headers.pop("x-trusted-mcp", None) == "true" - ): - logger.info("Initialize TrustedMCP session via trusted_mcp_client") - session = await exit_stack.enter_async_context( - trusted_mcp_client( - url=self._connection_params.url, - headers=merged_headers, - timeout=timedelta(seconds=self._connection_params.timeout), - sse_read_timeout=timedelta( - seconds=self._connection_params.sse_read_timeout - ), - terminate_on_close=self._connection_params.terminate_on_close, - ) + merged_headers.pop("x-trusted-mcp", None) + logger.info("Initialize TrustedMCP session via trusted_mcp_client") + session = await exit_stack.enter_async_context( + trusted_mcp_client( + url=self._connection_params.url, + headers=merged_headers, + timeout=timedelta(seconds=self._connection_params.timeout), + sse_read_timeout=timedelta( + seconds=self._connection_params.sse_read_timeout + ), + terminate_on_close=self._connection_params.terminate_on_close, ) + ) - # Store session and exit stack in the pool - self._sessions[session_key] = (session, exit_stack) - logger.debug("Created new session: %s", session_key) - return session - else: - # Create a new session - logger.info("Initialize MCP session") - return await super().create_session(headers) + # Store session and exit stack in the pool + self._sessions[session_key] = (session, exit_stack) + logger.debug("Created new session: %s", session_key) + return session except Exception: # If session creation fails, clean up the exit stack if exit_stack: