Skip to content

Commit 337c43e

Browse files
authored
fix: make mcp_instrumentation idempotent to prevent recursion errors (#892)
- Add module-level flag _instrumentation_applied to track patch state - Return early from mcp_instrumentation() if already applied - Prevents wrapper accumulation that causes RecursionError with multiple MCPClient instances - Add integration tests for multiple client creation and thread safety Fixes #869 🤖 Assisted by Amazon Q Developer
1 parent 68103f6 commit 337c43e

File tree

5 files changed

+51
-5
lines changed

5 files changed

+51
-5
lines changed

src/strands/tools/mcp/mcp_client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ def stop(
149149
- _background_thread_event_loop: AsyncIO event loop in background thread
150150
- _close_event: AsyncIO event to signal thread shutdown
151151
- _init_future: Future for initialization synchronization
152-
152+
153153
Cleanup order:
154154
1. Signal close event to background thread (if session initialized)
155155
2. Wait for background thread to complete

src/strands/tools/mcp/mcp_instrumentation.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@
1818
from opentelemetry import context, propagate
1919
from wrapt import ObjectProxy, register_post_import_hook, wrap_function_wrapper
2020

21+
# Module-level flag to ensure instrumentation is applied only once
22+
_instrumentation_applied = False
23+
2124

2225
@dataclass(slots=True, frozen=True)
2326
class ItemWithContext:
@@ -48,7 +51,14 @@ def mcp_instrumentation() -> None:
4851
- Adding OpenTelemetry context to the _meta field of MCP requests
4952
- Extracting and activating context on the server side
5053
- Preserving context across async message processing boundaries
54+
55+
This function is idempotent - multiple calls will not accumulate wrappers.
5156
"""
57+
global _instrumentation_applied
58+
59+
# Return early if instrumentation has already been applied
60+
if _instrumentation_applied:
61+
return
5262

5363
def patch_mcp_client(wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any) -> Any:
5464
"""Patch MCP client to inject OpenTelemetry context into tool calls.
@@ -167,6 +177,9 @@ def traced_method(
167177
"mcp.server.session",
168178
)
169179

180+
# Mark instrumentation as applied
181+
_instrumentation_applied = True
182+
170183

171184
class TransportContextExtractingReader(ObjectProxy):
172185
"""A proxy reader that extracts OpenTelemetry context from MCP messages.

tests/strands/tools/mcp/test_mcp_client.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -522,15 +522,16 @@ def test_stop_with_background_thread_but_no_event_loop():
522522
# Verify cleanup occurred
523523
assert client._background_thread is None
524524

525-
525+
526526
def test_mcp_client_state_reset_after_timeout():
527527
"""Test that all client state is properly reset after timeout."""
528+
528529
def slow_transport():
529530
time.sleep(4) # Longer than timeout
530531
return MagicMock()
531532

532533
client = MCPClient(slow_transport, startup_timeout=2)
533-
534+
534535
# First attempt should timeout
535536
with pytest.raises(MCPClientInitializationError, match="background thread did not start in 2 seconds"):
536537
client.start()
@@ -539,4 +540,4 @@ def slow_transport():
539540
assert client._background_thread is None
540541
assert client._background_thread_session is None
541542
assert client._background_thread_event_loop is None
542-
assert not client._init_future.done() # New future created
543+
assert not client._init_future.done() # New future created

tests/strands/tools/mcp/test_mcp_instrumentation.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from mcp.types import JSONRPCMessage, JSONRPCRequest
66
from opentelemetry import context, propagate
77

8+
from strands.tools.mcp.mcp_client import MCPClient
89
from strands.tools.mcp.mcp_instrumentation import (
910
ItemWithContext,
1011
SessionContextAttachingReader,
@@ -14,6 +15,17 @@
1415
)
1516

1617

18+
@pytest.fixture(autouse=True)
19+
def reset_mcp_instrumentation():
20+
"""Reset MCP instrumentation state before each test."""
21+
import strands.tools.mcp.mcp_instrumentation as mcp_inst
22+
23+
mcp_inst._instrumentation_applied = False
24+
yield
25+
# Reset after test too
26+
mcp_inst._instrumentation_applied = False
27+
28+
1729
class TestItemWithContext:
1830
def test_item_with_context_creation(self):
1931
"""Test that ItemWithContext correctly stores item and context."""
@@ -328,6 +340,27 @@ def __getattr__(self, name):
328340

329341

330342
class TestMCPInstrumentation:
343+
def test_mcp_instrumentation_idempotent_with_multiple_clients(self):
344+
"""Test that mcp_instrumentation is only called once even with multiple MCPClient instances."""
345+
346+
# Mock the wrap_function_wrapper to count calls
347+
with patch("strands.tools.mcp.mcp_instrumentation.wrap_function_wrapper") as mock_wrap:
348+
# Mock transport
349+
def mock_transport():
350+
read_stream = AsyncMock()
351+
write_stream = AsyncMock()
352+
return read_stream, write_stream
353+
354+
# Create first MCPClient instance - should apply instrumentation
355+
MCPClient(mock_transport)
356+
first_call_count = mock_wrap.call_count
357+
358+
# Create second MCPClient instance - should NOT apply instrumentation again
359+
MCPClient(mock_transport)
360+
361+
# wrap_function_wrapper should not be called again for the second client
362+
assert mock_wrap.call_count == first_call_count
363+
331364
def test_mcp_instrumentation_calls_wrap_function_wrapper(self):
332365
"""Test that mcp_instrumentation calls the expected wrapper functions."""
333366
with (

tests_integ/test_mcp_client.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -297,4 +297,3 @@ def slow_transport():
297297
with client:
298298
tools = client.list_tools_sync()
299299
assert len(tools) >= 0 # Should work now
300-

0 commit comments

Comments
 (0)