spec: MCP session lifecycle fix — Phase 1 - #122
Conversation
Add OpenSpec change for fixing stale MCP toolset cache and session-scoped resource lifecycle bugs (#121). Includes: - proposal.md: What & why (6 lifecycle fixes, no config changes) - design.md: 8 design decisions (D1-D8) with Oracle + Momus review - specs/mcp-session-lifecycle: 7 requirements, 14 scenarios - specs/session-orchestration: Modified requirements for close path - specs/unified-session-lifecycle: WebSocket disconnect hook - tasks.md: 7 task groups, 46 tasks (P1a-P1f + E2E) - tests/mcp_server/test_stale_mcp_connection.py: 5 reproduction tests Reviewed by Momus (PASS) and Oracle (PASS) after 2 revision cycles. Closes #121 (spec phase)
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive redesign of the MCP-over-ACP session lifecycle to resolve resource leaks that occur during session resumption and WebSocket disconnects. The changes include implementing per-session MCP resource tracking in the MCPManager, adding a deterministic cleanup mechanism for session-scoped resources, and ensuring that all session closure paths—including resume and disconnect—properly handle active run cancellation before cleaning up resources. My review identified several critical design gaps regarding session key tracking in the AcpMcpConnectionManager, potential attribute errors in cleanup paths, and the need for more robust resource cleanup in both production code and test suites.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
4 accepted fixes from dialectical analysis with Oracle: 1. Task 3.1/3.2/3.3: Change _session_connections to dict[str, set[tuple[str, int]]] — store (connection_id, session_key) pairs so AcpMcpConnectionManager.cleanup_session() can look up SessionStreamPair via session_key 2. Task 5.2 (D6): Two-layer cleanup on resume — call both SessionController.close_session() (RunHandle lifecycle) AND ACPSession.close() (ACP env/signals/prompts). Neither alone is sufficient. 3. Task 6.4 (D7): Same two-layer cleanup for WebSocket disconnect 4. Task 2.8: Use try/finally or fixture teardown for test cleanup Rejected comments (2): - hasattr(self.agent, 'mcp'): violates AGENTS.md, mcp always set - hasattr(agent, 'mcp'): same, agent is not None check exists Already addressed (2): - Concurrency re-verify after lock: spec's lock-on-context design handles this implicitly - await on_disconnect: type signature makes it obvious
- Add _SessionContext dataclass to MCPManager with per-session state (connection_pool, toolset_cache, snapshot, acp_connection_ids, _cleanup_lock) - Add _session_contexts dict to MCPManager.__init__ - Add _session_connections reverse index to AcpMcpConnectionManager - Add register_session_connection() method for tracking session→connection mappings Implements T1 and T6 of fix-mcp-session-lifecycle plan.
- get_or_create_session() and update_session_snapshot() on MCPManager (T2) - add_acp_transport() on MCPManager for session-scoped ACP tracking (T3) - register_session() returns tuple[SessionStreamPair, int] (T7, GAP-1) - has_active_sessions() on AcpMcpConnection (T7) - cleanup_session() with _cleanup_lock on AcpMcpConnectionManager (T7, GAP-12) - Updated all callers of register_session() to unpack tuple return
…nnection - cleanup_session() with per-session _cleanup_lock on MCPManager (T4) - _acp_mcp_manager field added for ACP cleanup delegation - connect_acp_mcp_server() gains session_id parameter (T8, GAP-5) - Returns tuple[str, int] (connection_id, session_key) - Call site in session.py passes session_id and calls add_acp_transport - All test callers updated for new signature
- Change as_capability(snapshot=, session_pool=) to as_capability(session_id=) - Parameterize _make_capability with toolset_cache dict parameter (GAP-7) - Split _process_snapshot into _process_global_configs and _process_session_configs - GAP-11: KeyError fallback for concurrent cleanup_session race - Backward compat: session_id=None processes self.servers with self._toolset_cache
… (T12) - Replace as_capability(snapshot=, session_pool=) with as_capability(session_id=) - GAP-4: Use run_ctx.session_id from AgentRunContext instead of self._session_id - Remove if/else branching on _mcp_snapshot — as_capability handles internally - Keep _mcp_snapshot and _session_connection_pool field declarations for compat
- Update 6 tests in test_mcpmanager_caching.py for new as_capability(session_id) API - Update 15 failing tests in test_mcp_provider_lifecycle.py to use session context - Fix static source assertion in test_no_dedup_hack_in_get_agentlet - All 48 tests pass
- Rename test_session_resume_returns_stale_toolset → _returns_fresh_toolset - Rename test_multiple_acp_servers_all_go_stale → _get_fresh_toolsets - Rename test_disconnect_all_clears_cache → test_cleanup_session_clears_per_session_cache - All 5 tests now verify the fix instead of documenting the bug - All tests pass with new session_id API
…run (T21+T22+T23)
- Add on_disconnect: Callable[[AgentSideConnection], Awaitable[None]] | None parameter - Generate UUID4 connection_id on AgentSideConnection at accept time (GAP-3) - Call on_disconnect in ConnectionClosed handler before conn.close() - Backward compatible: on_disconnect defaults to None
- Add _connection_sessions reverse index on ACPSessionManager - Add connection_id parameter to create_session() and resume_session() - Implement close_all_sessions_for_connection() for WebSocket disconnect cleanup - Idempotent: pops connection_id, iterates sessions, closes via SessionController + ACPSession.close()
- Add on_disconnect parameter to serve(), _serve_websocket(), _serve_streamable_http() - Wire on_disconnect callback in ACPServer._start_async() closure - Add session_manager field to AgentPoolACPAgent for shared session tracking - Create shared ACPSessionManager in ACPServer for cross-connection session tracking - Add disconnect detection in _serve_streamable_http via recv_task completion - Fix test_resume_session_is_idempotent -> test_resume_session_closes_old_and_recreates (T20 changed resume_session from idempotent to close-then-recreate)
…(T27+T28) - T27: test_websocket_disconnect_closes_all_sessions — 2 sessions same conn, disconnect, both closed - T27: test_websocket_disconnect_preserves_other_connections — 2 conns, disconnect one, other survives - T28: test_websocket_disconnect_during_run — active run cancelled with 2s timeout on disconnect
…ssion lifecycle test (T33) - T32: Use cast() to type session_manager field as ACPSessionManager (not | None) for mypy - T33: test_e2e_session_lifecycle — full lifecycle: connect→session→MCP→disconnect→reconnect→resume→verify fresh
|
/gemini review |
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements Phase 1 of the MCP session lifecycle redesign, centralizing session-scoped state in _SessionContext on MCPManager, wiring cleanup into all close paths, and fixing resume_session() to close-then-recreate. The code review highlights several critical and high-severity issues, including a potential premature termination bug in resume_session() due to untracked session mappings in _connection_sessions, missing imports (asyncio and uuid), and potential memory leaks or unnecessary instantiations in as_capability() and cleanup_session(). Additionally, the reviewer suggests adding defensive None checks for agent.mcp and validating connection identifiers before registration.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
- ruff format: reformat 5 files (manager.py, session_controller.py, session.py, test_session_lifecycle.py, test_stale_mcp_connection.py) - ruff check: shorten docstring in test_acp_session_resume.py (E501)
…e leaks - Fix #1 (Critical): resume_session() now removes session_id from _connection_sessions before closing old session. Prevents stale connection disconnect from closing the newly resumed session. - Fix #2 (High): as_capability() uses _session_contexts.get() instead of get_or_create_session(). Prevents memory leak when context was already cleaned up. Removes dead try/except KeyError code. - Fix #3 (Medium): cleanup_session() uses _session_contexts.get() and returns early if None. Avoids creating throwaway SessionConnectionPool. - TDD: 3 tests in test_review_fixes.py verify all fixes.
|
/gemini review |
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements Phase 1 of the MCP session lifecycle redesign, centralizing session-scoped MCP state into _SessionContext on MCPManager and ensuring proper cleanup of resources (toolsets, transports, and ACP connections) when sessions close or WebSocket connections drop. The feedback identifies a critical issue where _acp_mcp_manager is never set in production, causing session-scoped ACP connections to leak, and a missing uuid import in transports.py that will cause a runtime NameError. Additionally, suggestions are provided to enhance idempotency in cleanup_session and reduce duplicated fallback logic in as_capability.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…apability Review round 2 fixes: Fix #4 (Critical): Wire _acp_mcp_manager in ACPSession.__post_init__ - MCPManager._acp_mcp_manager was initialized to None and never set - cleanup_session() could never delegate to AcpMcpConnectionManager - Per-session ACP stream pairs and reverse-index entries leaked - Fix: wire agent.mcp._acp_mcp_manager = acp_agent._mcp_manager in __post_init__ Fix #5 (Medium): Identity check after acquiring cleanup lock - Concurrent cleanup_session() callers could do redundant work - All ops were idempotent but wasteful (clearing empty dicts, etc.) - Fix: check if self._session_contexts.get(session_id) is not ctx after lock Fix #6 (Medium): Consolidate duplicated fallback in as_capability() - Three identical 'for server in self.servers:' loops consolidated to one - Pure readability refactor, zero behavior change TDD: 3 new tests (2 RED before fix, 3 GREEN after) - test_cleanup_session_delegates_to_acp_mcp_manager (unit) - test_acp_session_wires_acp_mcp_manager (integration) - test_cleanup_session_identity_check_prevents_redundant_work (unit) 213 tests pass, ruff clean.
Categories A-D from Oracle integration test plan: - A (4): Cross-component wiring — cleanup delegation, __post_init__ wiring, full close chain, close_all_sessions_for_connection - B (7): Lifecycle edge cases — full create/cleanup, close/recreate, shared connection isolation, WebSocket disconnect, resume, concurrent cleanup - C (4): State consistency — registry consistency after cleanup/close/resume, stream pair unregistration - D (5): Error paths — ACP manager raises, session close raises, MCP cleanup raises, resume old close raises, pool cleanup raises These tests would have caught the _acp_mcp_manager wiring bug (round 2 review comment #1) that unit tests missed due to component isolation.
…ll sites - Declare connection_id: str | None on AgentSideConnection (replaces monkey-patch) - Remove # type: ignore[attr-defined] from transports.py connection_id assignments - Add _get_connection_id() helper on AgentPoolACPAgent using isinstance check - Wire connection_id= into all 5 create_session/resume_session call sites: new_session, load_session, fork_session, resume_session, handler.py - Fix misleading GAP-11 comment: dict.get() returns None, never raises KeyError Without this fix, _connection_sessions dict was never populated, making close_all_sessions_for_connection() always return immediately — the entire WebSocket disconnect cleanup feature was dead code.
|
/gemini review |
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
Covers all 13 gap areas identified by Oracle analysis: - G1: Full create_session → get_or_create_session_agent → MCPManager chain - G2: as_capability with non-empty ACP snapshot → real MCPToolset - G3: initialize_mcp_servers → connect_acp_mcp_server → AcpMcpTransport - G4: Full tool execution through as_capability → MCPToolset → AcpMcpTransport - G5: SessionController.close_session with real agent + real MCP resources - G6: resume_session with real ACPSession (not patched) - G7: Full on_disconnect → close_all_sessions_for_connection chain - G8: connection_id propagation: create_session populates _connection_sessions - G9: as_capability during concurrent cleanup (GAP-11 race) - G10: ACP transport failure during tool execution + cleanup - G11: Multiple sessions on same connection with real ACPSessions - G12: Child session inherits parent's ACP transports - G13: Pool shutdown cleans all session MCP resources
|
/gemini review |
- server.py: Remove unused type: ignore, use None guard for connection_id - test_acp_session_resume.py: Add connection_id to expected resume_session call args
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
- Mark all 46 tasks as complete in tasks.md - Sync 3 delta specs to main specs: - mcp-session-lifecycle (new) - session-orchestration (updated) - unified-session-lifecycle (updated) - Archive to openspec/changes/archive/2026-07-07-fix-mcp-session-lifecycle/
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements a robust cleanup mechanism for MCP session-scoped resources, including toolsets, connection pools, and ACP connections, to prevent resource leaks. Key changes include centralizing session state in _SessionContext, wiring cleanup into session close and WebSocket disconnect paths, and refining session management logic. Feedback from the review addressed potential runtime errors (missing imports, AttributeError risks) and memory leaks (unnecessary context recreation), which have been incorporated to ensure the lifecycle management is both safe and efficient.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
- session_controller.py: Replace get_or_create_session() with _session_contexts.get() when reading parent snapshot/pool. Prevents phantom _SessionContext creation when parent was already cleaned up. - transports.py: Move on_disconnect callback from except ConnectionClosed to finally block. Ensures callback fires on any exception path. - 3 TDD tests: leak detection, regression guard, disconnect coverage.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements Phase 1 of the MCP session lifecycle redesign, resolving resource leaks by centralizing session-scoped state into _SessionContext on MCPManager and ensuring coordinated cleanup across all session close and WebSocket disconnect paths. It also updates as_capability() to use a session-isolated toolset cache and replaces the early-return in resume_session() with a close-then-recreate flow. The review feedback correctly identifies a potential AttributeError in SessionController._close_session_run_turn when accessing agent.mcp on non-native or mock agents, suggesting an explicit type check to ensure safety.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Architecture Review: MCP Lifecycle AnalysisReviewed the full PR from a lifecycle architecture perspective. The core design is sound — three-tier model (Pool → Agent → Session), ✅ What's Correct
❌ Structural Gaps1. Subagent inherits ACP transport without registering on connection manager (High)This is the most significant structural gap. Child sessions copy ACP transport references via Failure chain: Deeper issue: In # Current — only copies transport reference
await child_ctx.connection_pool.copy_pre_created_transports(
parent_ctx.connection_pool
)
# Missing — register child on ACP connection manager
for (conn_id, session_key) in parent_ctx.acp_connection_ids:
conn = parent_agent.mcp._acp_mcp_manager._connections.get(conn_id)
if conn is not None:
_pair, child_session_key = conn.register_session()
parent_agent.mcp._acp_mcp_manager.register_session_connection(
child_session_id, conn_id, child_session_key
)
await child_agent.mcp.add_acp_transport(
child_session_id, client_id, transport, conn_id, child_session_key
)2.
|
| Item | Status | Priority |
|---|---|---|
| Three-tier lifecycle model | ✅ Correct | — |
_SessionContext centralization |
✅ Correct | — |
| Per-session toolset isolation | ✅ Correct | — |
| Cleanup chain structure | ✅ Correct | — |
| Subagent ACP transport registration | ❌ Missing | High |
GlobalConnectionPool lock type |
Medium-High | |
| Global toolset invalidation | ❌ Missing | Medium |
| Snapshot concurrency protection | Medium | |
| Agent-level MCPManager redundancy | Low | |
| Transport health check | ❌ Missing | Low |
#1 (subagent ACP transport registration) is the only item that could cause silent data corruption in production. Recommend addressing it immediately after this PR merges. The rest can be tackled systematically in Phase 2.
… toolset __aexit__
Three fixes for child session ACP transport registration gaps:
1. Wire _acp_mcp_manager on child agent from parent (session_controller.py)
- Child sessions created via get_or_create_session_agent() don't go
through ACPSession.__post_init__, so _acp_mcp_manager stayed None.
Now copied from parent after copy_pre_created_transports().
2. Add on_session_registered callback to AcpMcpTransport (acp_mcp_transport.py)
- Optional callback invoked after register_session() with (connection_id,
session_key). Enables callers to register ACP connections for cleanup
tracking via register_session_connection().
3. Fix toolset_cache.clear() to call __aexit__ first (manager.py)
- cleanup_session() called .clear() without closing MCPToolset instances,
leaking stream pairs and forwarder tasks. Now mirrors disconnect_all()
pattern: iterate values, call __aexit__(None, None, None) with
contextlib.suppress(ValueError), then clear.
TDD: 3 tests in test_child_session_acp_fix.py (all GREEN).
252 MCP+ACP tests pass, 0 regressions, ruff clean.
|
感谢这份详细的架构评审！逐条回复： Gap 1: Subagent inherits ACP transport without registering on connection manager (High)已在 commit
关于评审中提到的"message cross-talk"风险： 关于"parent cleanup kills child transport"：子的 Gap 2:
|
| Gap | Status | Commit |
|---|---|---|
| 1. Subagent ACP transport registration | ✅ Fixed | b6cafcd6a |
2. GlobalConnectionPool lock type |
— | |
| 3. Global toolset_cache invalidation | — | |
| 4. Snapshot concurrency | — | |
| 5. Agent-level MCPManager redundancy | — | |
| 6. Transport health check | ℹ️ Intentional design | — |
感谢审查！Gap 1 是唯一可能导致 silent data corruption 的问题，已在本 PR 中修复。其余均为 pre-existing 或 Phase 2 范畴。
2 participants
Summary
OpenSpec change for fixing stale MCP toolset cache and session-scoped resource lifecycle bugs identified in #121.
Problem
Session-scoped MCP resources (toolsets, transports, ACP connections) are never cleaned up when sessions close or WebSocket connections drop. The root cause is that session-scoped MCP state is scattered across 4 different objects (
MCPManager._toolset_cache,Agent._session_connection_pool,Agent._mcp_snapshot,AcpMcpConnectionManager._connections) with no coordinated cleanup. This causes silent session failures on resume — the agentlet tries to initialize MCP via dead transports, causing a 300-second timeout.What's in this PR
This PR contains only the spec (no implementation). It adds:
proposal.md— What & why: 6 lifecycle fixes, no config API changesdesign.md— 8 design decisions (D1-D8) covering session tracking, toolset cache scoping, cleanup wiring, concurrency protection, WebSocket disconnect hookspecs/mcp-session-lifecycle/spec.md— New capability: 7 requirements, 14 scenariosspecs/session-orchestration/spec.md— Modified: close_session cleanup, agent registrationspecs/unified-session-lifecycle/spec.md— Modified: WebSocket disconnect hooktasks.md— 7 task groups, 46 tasks (P1a→P1f + E2E verification)tests/mcp_server/test_stale_mcp_connection.py— 5 reproduction tests (all passing)Design Decisions
_toolset_cachefor pool-level onlyas_capability(session_id)simplified APIReview Status
Reviewed by Momus (PASS) and Oracle (PASS) after 2 revision cycles.
Revision history
Migration Plan
This is Phase 1 of a 2-phase MCP lifecycle redesign:
Related
tests/mcp_server/test_stale_mcp_connection.py(5 tests, all passing)_toolset_cacheremoved in88bdc1758(cross-task fix), re-added ind3b4966c8(regression)Next Steps
Run
/opsx:applyto begin implementation followingtasks.md.