fix(api): reuse MCP client connections across tool invocations - #42658
LetMeSleep8h wants to merge 4 commits into
Conversation
|
Thanks — the pooling direction is right, and keying by connection fingerprint (URL + effective headers + timeouts) is a sensible first cut. Four things before this can merge, though: (1) the lock ordering can wedge the registry for every caller; (2) the key is built from per-invocation credentials, so it never matches when identity forwarding is on — the Playwright scenario in #42650 stays broken; (3) with forwarding off it matches too broadly — one session shared across users and concurrent runs; (4) nothing closes the pool in production. Details and a runnable repro below, traced against 1. Blocking:
|
|
Thanks for the thorough review — all four points are addressed in 155d75ee:
On §5 — agreed both are pre-existing and pooling amplifies them. Beyond the Tests: |
|
Second pass on In candidate = _PooledClient(client, key)
with self._lock:
existing = self._clients.get(key)
if existing is None:
self._clients[key] = candidate
victims = self._evict_to_capacity() # bound only here
for victim in victims: # reached on both paths
self._close_entry(victim)
if existing is not None:
# Lost a race creating the same connection: use the winner.
self._close_entry(candidate)
return existingWhen two callers race on the same key — exactly the case that loser branch exists for — the loser has The suite doesn't catch it because the concurrency tests use different scopes ( Repro (my transcription of import threading, time
CONNECT_WINDOW = threading.Event()
class Entry:
def __init__(self, key):
self.key, self.last_used = key, time.monotonic()
self.lock, self.doomed, self.closed = threading.Lock(), False, False
self.client = self
def cleanup(self):
self.closed = True
class Manager: # mirrors client_manager.py _acquire
def __init__(self):
self._clients, self._lock, self._max_size = {}, threading.Lock(), 100
def _evict_to_capacity(self):
return [] # no overflow at max_size=100
def _close_entry(self, entry):
if not entry.lock.acquire(blocking=False):
entry.doomed = True
return
try:
if not entry.closed:
entry.closed = True
entry.client.cleanup()
finally:
entry.lock.release()
def acquire(self, key):
with self._lock:
existing = self._clients.get(key)
if existing is not None:
existing.last_used = time.monotonic()
if existing is not None:
return existing
CONNECT_WINDOW.wait() # "connect" outside the registry lock
candidate = Entry(key)
with self._lock:
existing = self._clients.get(key)
if existing is None:
self._clients[key] = candidate
victims = self._evict_to_capacity()
for victim in victims: # <-- UnboundLocalError on the loser
self._close_entry(victim)
if existing is not None:
self._close_entry(candidate)
return existing
return candidate
manager, errors, key = Manager(), [], "same-scope-key"
def caller(n):
try:
manager.acquire(key)
except Exception as exc:
errors.append((n, type(exc).__name__, str(exc)))
threads = [threading.Thread(target=caller, args=(n,), daemon=True) for n in (1, 2)]
for t in threads: t.start()
time.sleep(0.05) # both pass the registry check
CONNECT_WINDOW.set() # then both race the registration
for t in threads: t.join(timeout=5)
print("callers that raised:", errors)
print("pooled keys after the race:", list(manager._clients))Output: Minimal fix — bind duplicate: _PooledClient | None = None
with self._lock:
existing = self._clients.get(key)
if existing is None:
self._clients[key] = candidate
victims = self._evict_to_capacity()
else:
victims = []
duplicate = candidate
for victim in victims:
self._close_entry(victim)
if duplicate is not None:
self._close_entry(duplicate)
return existing
return candidateWorth a test alongside the two you added: two threads on the same key, assert one client created and neither call raises. (The loser's client has already completed |
|
On the race-loser crash — one correction with evidence, though the smell you flagged is real and already fixed: In That said, the shared name across two eviction sites was exactly the trap you identified, and a redundancy pass had already flagged it — Added the test you sketched in Also refreshed the PR description to match the current design (stable-scope key, 14 tests, atexit hook, fail-fast-until-TTL behavior note). |
2 participants
Summary
Each
MCPToolinvocation opened a new MCP connection and closed it as soon as the call returned, so stateful MCP servers lost all session state between calls — e.g. with the official Playwright MCP server,browser_navigatesucceeds but the nextbrowser_snapshotlands on a freshabout:blankcontext.This adds a process-wide
MCPClientManager(api/core/mcp/client_manager.py) that pools successfully initialized clients, and routesMCPTool.invoke_remote_mcp_toolthrough it:initialize()keeps working and the pooled client refreshes its own token on 401. Distinct tenants, end users and provider configs get distinct connections, so a stateful server never shares session state across them;ClientSessiondrives a single-worker executor). Entries are closed outside the registry lock — a busy entry is markeddoomedand closed by its in-flight call on the way out, so a call hung on an unresponsive server (the session has no read timeout) can never wedge other keys;MCP_CLIENT_POOL_IDLE_TTL) and the pool is capped (MCP_CLIENT_POOL_MAX_SIZE, default 100), evicting least-recently-used entries;MCPConnectionError) is evicted and the call retried once on a fresh connection. Auth errors are deliberately not retried (MCPAuthErroralso extendsMCPConnectionError): token refresh already happens insideMCPClientWithAuthRetry, and re-invoking after a final auth failure could execute the tool twice.ValueError(transport-level parsing failures that corrupt the session) evicts the entry without a retry;atexithook (mirroringcore/helper/http_client_pooling.py) closes the pool on process exit;Tool listing is unchanged (still connect-per-call) — happy to follow up if reuse is wanted there too.
Fixes #42650
Tests
api/tests/unit_tests/core/mcp/test_client_manager.py: scope reuse, credential-change reuse, user/tenant/provider isolation, dead-connection eviction + retry, auth-error propagation without eviction,ValueErroreviction without retry, a hung call not blocking other scopes, idle-eviction of an in-flight entry, a same-key registration race (neither caller raises, loser's connection closed, one entry pooled), TTL and LRU eviction, failed-init not pooled,close_all. The two concurrency cases were verified red against the pre-fix implementation during development.MCPClientWithAuthRetryintool.pyto patching the manager.pytest tests/unit_tests/core/mcp tests/unit_tests/core/tools tests/unit_tests/tools→ 867 passed;ruff check/ruff formatclean.