fix(mcp): keep connection cleanup in owner tasks by hertznsk · Pull Request #358 · microsoft/conductor

hertznsk · 2026-08-01T10:07:14Z

Summary

move each MCP stdio/session context lifecycle into a persistent owner task
signal owner tasks during manager shutdown so AnyIO cancel scopes exit in the task that entered them
preserve caller cancellation while consuming and logging cleanup failures
protect owner-task teardown when manager shutdown is cancelled
reject duplicate server names without orphaning the existing connection

Testing

PYTHONASYNCIODEBUG=1 uv run pytest tests/test_mcp -q (51 passed)
make check
direct stdio MCP lifecycle driver covering connect, tool call, duplicate names, failure, cancellation, concurrency, and idempotent close

Fixes #353

Jason Robert (jrob5756)

Good fix overall. The owner-task-per-connection model is the right way to solve the anyio cancel-scope task affinity problem, and the new task-affinity tests are a solid way to prove it.

Left a few small comments on double-cancellation edge cases and a couple of small cleanups, but nothing blocking here.

Jason Robert (jrob5756) · 2026-08-03T19:13:36Z

-        self._initialized = False
+            for result in results:
+                if isinstance(result, Exception):
+                    logger.warning(f"Error closing MCP connections: {result}")


If close() gets cancelled a second time while it's sitting at the unshielded await cleanup on line 573, that cancellation propagates into the still-running owner tasks and can interrupt teardown mid-flight (e.g. while a subprocess is being terminated). But cleanup.done() is true even for a cancelled future, so this finally block clears all bookkeeping as if everything shut down cleanly, and there's no log line for this path at all. Worth tracking whether results was actually obtained before treating state as cleared, and logging when a repeated cancellation prevents confirming cleanup completed.

Fixed in ca8551d. The unshielded await cleanup fallback is replaced with a re-shielding loop: repeated cancellations now land on a throwaway asyncio.shield(cleanup) while the gather keeps running, and bookkeeping is cleared only after results is actually obtained (a cancelled gather result is never read). The behavior goes slightly further than suggested: close() now deliberately absorbs cancellation until owner-task teardown completes (documented in the docstring), so teardown is never interrupted mid-flight and state is always consistent. Covered by test_double_cancelled_close_still_finishes_cleanup.

Jason Robert (jrob5756) · 2026-08-03T19:13:36Z

+            cleanup = asyncio.gather(task, return_exceptions=True)
+            try:
+                results = await asyncio.shield(cleanup)
+            except asyncio.CancelledError:


Same gap here as in close(): if a second cancellation lands while awaiting the unshielded await cleanup on this line, it propagates straight out and skips the pop()/_discard_server_state() calls below. That leaves _connection_tasks[name] populated with a task that may not have finished cleanup, so subsequent connect_server(name=...) calls for that name will incorrectly raise "already connected or connecting" until the whole manager is torn down, with no log explaining why.

Fixed in ca8551d with the same re-shielding loop as in close(): a repeated cancellation lands on the shield instead of the gather, so cleanup completes and the pop()/_discard_server_state() calls always run before the cancellation is re-raised. _connection_tasks[name] can no longer be left populated, so a later connect_server(name=...) will not be spuriously rejected. Covered by test_double_cancelled_connection_clears_bookkeeping.

Jason Robert (jrob5756) · 2026-08-03T19:13:36Z

+            self._connection_tasks.pop(name, None)
+            self._connection_stops.pop(name, None)
+            self._discard_server_state(name)
+            logger.error(f"Failed to connect to MCP server '{name}': {exc}")


{exc} only calls str(). If the AsyncExitStack teardown raises a different exception while unwinding after the real connect failure (e.g. a broken pipe during stdio_client/ClientSession exit), that secondary exception replaces the original as what's caught and logged here, masking the actual root cause.

Suggested change

logger.error(f"Failed to connect to MCP server '{name}': {exc}")

logger.error(f"Failed to connect to MCP server '{name}': {exc}", exc_info=exc)

Applied in ca8551d - the log call now passes exc_info=exc, so the original connect failure traceback is preserved even when AsyncExitStack teardown raises a secondary exception while unwinding.

Jason Robert (jrob5756) · 2026-08-03T19:13:36Z

+            except asyncio.CancelledError:
+                results = await cleanup
+            for result in results:
+                if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError):


asyncio.CancelledError derives from BaseException, not Exception, since Python 3.8, so isinstance(result, Exception) already excludes it. The trailing check is dead code.

Suggested change

if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError):

if isinstance(result, Exception):

Applied in ca8551d - the dead CancelledError check is removed; the condition is now plain isinstance(result, Exception).

Jason Robert (jrob5756) · 2026-08-03T19:13:36Z

-        self._exit_stack = AsyncExitStack()
+        self._connection_tasks: dict[str, asyncio.Task[None]] = {}
+        self._connection_stops: dict[str, asyncio.Event] = {}
        self._initialized = False


_initialized is set here and at two other spots but nothing reads it anymore — close()'s guard now checks self._connection_tasks instead. Worth removing it so a future reader doesn't assume it reflects live connection state.

Removed in ca8551d - _initialized is gone from all three spots; close() keeps guarding on _connection_tasks.

Jason Robert (jrob5756) · 2026-08-03T19:13:36Z

-        if not self._initialized:
+        if not self._connection_tasks:
            return



This counts len(self.sessions), which is only populated once a connection's ready future resolves. If close() runs while a connect_server() call is still mid-handshake, this understates how many owner tasks are actually being torn down. len(self._connection_tasks) would be accurate.

Fixed in ca8551d - the debug log now counts len(self._connection_tasks), which includes connections still mid-handshake.

Address review feedback on double-cancellation edge cases: - Replace the unshielded `await cleanup` fallback in both connect_server() and close() with a re-shielding loop. A repeated cancellation previously landed on the gather itself, discarding its result and skipping (close()) or racing past (connect_server()) the bookkeeping cleanup, which could leave a stale _connection_tasks entry blocking reconnects. Repeated cancellations now land on a throwaway shield while the gather completes. - Log the connect failure with exc_info so a secondary AsyncExitStack teardown exception cannot mask the original root cause. - Drop the dead `not isinstance(result, asyncio.CancelledError)` check (CancelledError derives from BaseException since Python 3.8). - Remove the unused _initialized flag; close() already guards on _connection_tasks. - Count _connection_tasks (not sessions) when logging shutdown, so mid-handshake connections are included. close() now deliberately absorbs cancellation until owner-task teardown finishes, documented in its docstring. Regression tests cover double cancellation of both connect_server() and close().

hertznsk

All six review comments are addressed in ca8551d. Per-thread replies via the review-comment API return 404 for my token (fork-based PR), so summarizing here:

close() double-cancellation — the unshielded await cleanup fallback is replaced with a re-shielding loop: repeated cancellations land on a throwaway asyncio.shield(cleanup) while the gather keeps running, and bookkeeping is cleared only after results is actually obtained. close() now deliberately absorbs cancellation until owner-task teardown completes (documented in the docstring); covered by test_double_cancelled_close_still_finishes_cleanup.
connect_server() double-cancellation — same re-shielding loop: cleanup completes and the pop()/_discard_server_state() calls always run before the cancellation is re-raised, so _connection_tasks[name] can no longer be left populated and block reconnects; covered by test_double_cancelled_connection_clears_bookkeeping.
Connect-failure logging — now passes exc_info=exc, so a secondary AsyncExitStack teardown exception cannot mask the original root cause.
Dead not isinstance(result, asyncio.CancelledError) check removed.
Unused _initialized flag removed from all three spots.
Shutdown log now counts len(self._connection_tasks), including mid-handshake connections.

Full test suite passes (4580 tests), ruff and ty clean.

Jason Robert (jrob5756)

LGTM. Approved!

Codecov Comments Bot (codecov-commenter) · 2026-08-03T20:27:45Z

Codecov Report

❌ Patch coverage is 94.84536% with 5 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@f39bd63). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/conductor/mcp/manager.py	94.84%	5 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #358   +/-   ##
=======================================
  Coverage        ?   90.64%           
=======================================
  Files           ?       85           
  Lines           ?    14464           
  Branches        ?        0           
=======================================
  Hits            ?    13111           
  Misses          ?     1353           
  Partials        ?        0

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

fix(mcp): keep connection cleanup in owner tasks

14d7e0c

Jason Robert (jrob5756) reviewed Aug 3, 2026

View reviewed changes

hertznsk commented Aug 3, 2026

View reviewed changes

Jason Robert (jrob5756) approved these changes Aug 3, 2026

View reviewed changes

Jason Robert (jrob5756) merged commit 8f95a8f into microsoft:main Aug 3, 2026
10 checks passed

	logger.error(f"Failed to connect to MCP server '{name}': {exc}")
	logger.error(f"Failed to connect to MCP server '{name}': {exc}", exc_info=exc)

	if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError):
	if isinstance(result, Exception):

Uh oh!

Conversation

hertznsk commented Aug 1, 2026

Summary

Testing

Uh oh!

Jason Robert (jrob5756) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hertznsk left a comment

Choose a reason for hiding this comment

Uh oh!

Jason Robert (jrob5756) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Codecov Comments Bot (codecov-commenter) commented Aug 3, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants