fix: persist MCP connections across turns by eager-entering MCPToolset

[Leoyzen]

Problem

pydantic-ai's agent.iter() creates a new AsyncExitStack per turn that enters/exits MCPToolset. Without a persistent reference, _running_count goes 0→1→0 each turn, causing:

• Performance overhead: MCP re-initialization (server handshake, tool listing, resource listing) every turn
• Reliability: streamable_http_client.__aexit__() can hang when HTTP proxy doesn't promptly close TCP → deadlock chain (StreamCompleteEvent never yielded → wait_for_completion hangs → per-session lock never released → all subsequent messages deadlock)

Closes #175.

Solution

Pre-enter MCPToolset on cache miss. In MCPManager.get_capabilities(), when a new MCPToolset is created and cached, eagerly call await toolset.__aenter__(). This holds one reference open, so pydantic-ai's per-turn enter/exit goes 1→2→1 instead of 0→1→0.

Connection persists until cleanup_session() or disconnect_all() brings the count to 0 and closes the connection.

This leverages pydantic-ai's built-in _running_count reference counting — the same mechanism used by async with agent: (pydantic-ai's intended cross-run persistence pattern, see issue #2983). No wrappers, no new classes, no config changes.

Changes

File	Change
src/agentpool/mcp_server/manager.py	_make_capability() → async; eager __aenter__() on cache miss; disconnect_all() also closes session-scoped caches
tests/conftest.py	Autouse fixture patches MCPToolset.__aenter__/__aexit__ for unit tests; @pytest.mark.real_mcp opts out
tests/mcp_server/test_mcp_persistence.py	6 new TDD tests (red→green)
pyproject.toml	Register real_mcp marker
tests/mcp_server/test_e2e_*.py	Add @pytest.mark.real_mcp to E2E tests needing real MCPToolset

Key design decisions

1. Eager enter, not wrapper: ~15 LOC in 1 file. Uses pydantic-ai's built-in ref counting. No new classes.
2. Failure handling: If __aenter__ fails, toolset is NOT cached — enables retry on next call.
3. Session-scoped cleanup gap fixed: disconnect_all() now also iterates self._session_contexts to close per-session toolset caches (edge case found during analysis).
4. No config changes: Strict improvement, no feature flag needed.

Ref counting trace

Cache miss (first get_capabilities):
  __aenter__ → _running_count: 0→1 (eager enter)
  
Turn 1 (pydantic-ai iter()):
  __aenter__ → _running_count: 1→2
  __aexit__  → _running_count: 2→1  (connection survives!)
  
Turn 2 (pydantic-ai iter()):
  __aenter__ → _running_count: 1→2
  __aexit__  → _running_count: 2→1  (connection survives!)

Session close (cleanup_session):
  __aexit__  → _running_count: 1→0  (connection closed, caches cleared)

Testing

• 6 new tests in test_mcp_persistence.py:
  • test_eager_enter_sets_running_count — verifies _running_count == 1 after get_capabilities()
  • test_eager_enter_persists_across_calls — multiple calls don't re-enter
  • test_disconnect_all_closes_eagerly_entered_toolset — disconnect_all() brings count to 0
  • test_cleanup_session_closes_session_scoped_toolset — session cleanup closes session-scoped toolsets
  • test_eager_enter_failure_does_not_cache — failed enter doesn't leave stale cache entry
  • test_disconnect_all_also_cleans_session_scoped_toolsets — disconnect_all() cleans session caches too
• All 180 existing MCP server tests pass (no regressions)
• TDD workflow: Tests written first (red), then implementation (green)
• Autouse test fixture patches MCPToolset.__aenter__/__aexit__ globally to avoid real MCP connections in unit tests

[gemini-code-assist]

Code Review

This pull request introduces extensive documentation, design specifications, and benchmarks for the AgentPool framework, detailing the OpenSpec workflow, graph architecture, session orchestration, and lifecycle dimensions. It also adds several experimental commands and skills. The review feedback points out a few key issues: a potential FileNotFoundError in the benchmark script when writing to a gitignored directory, a hardcoded absolute path in the diagnostic script, a potential runtime error in the EventBus design due to frozen dataclasses, and a type mismatch in the YAML graph syntax design.

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.

[Leoyzen]

/gemini review

[gemini-code-assist]

Code Review

This pull request resolves issue #175 by ensuring MCP connection persistence across turns. This is achieved by eagerly entering MCPToolset instances via __aenter__ on cache miss within MCPManager.get_capabilities(), holding a persistent reference open for the session/pool lifetime. Additionally, disconnect_all() has been updated to clean up both global and per-session toolset caches to prevent connection leaks. A new test suite has been added to verify this lifecycle behavior, along with a pytest fixture to mock the toolset lifecycle and avoid real MCP connections in unit tests. Feedback from the reviewer suggests improving the robustness of the cleanup process in disconnect_all() by wrapping the __aexit__ calls for per-session toolsets in a try-except block with a timeout and logging unexpected exceptions, preventing potential hangs or aborts from blocking the entire shutdown process.

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.

[gemini-code-assist]

In disconnect_all(), the per-session toolsets are cleaned up by calling await toolset.__aexit__(None, None, None) without a timeout and only suppressing ValueError. As noted in the PR description, __aexit__ calls can hang (e.g., when HTTP proxies do not promptly close TCP connections). If any single toolset hangs or raises an unexpected exception (such as TimeoutError or OSError), it will block or abort the entire disconnect_all() process, preventing other sessions and the global pool from being cleaned up. To prevent deadlocks and resource leaks, wrap each __aexit__ call in a try-except block with a timeout (using _MCP_CLEANUP_TIMEOUT) and catch/log all unexpected exceptions, matching the robust pattern already used in cleanup_session().

Suggested change

	for ctx in self._session_contexts.values():
	for toolset in ctx.toolset_cache.values():
	with contextlib.suppress(ValueError):
	await toolset.__aexit__(None, None, None)
	ctx.toolset_cache.clear()
	for ctx in self._session_contexts.values():
	for toolset in ctx.toolset_cache.values():
	with contextlib.suppress(ValueError):
	try:
	async with asyncio.timeout(_MCP_CLEANUP_TIMEOUT):
	await toolset.__aexit__(None, None, None)
	except TimeoutError:
	logger.warning("MCP session toolset cleanup timed out during disconnect_all")
	except Exception:
	logger.exception("Error cleaning up session toolset during disconnect_all")
	ctx.toolset_cache.clear()

References

1. When closing multiple streams or resources in a loop, use an explicit try-except block to catch expected library-specific exceptions silently, while logging any other unexpected exceptions to prevent them from masking failures or halting the cleanup of remaining resources.