feat: add threading.lock on async stream; expand interrupts to other AgentBase instances

Author: Owen Kaplan

src/strands/agent/_agent_as_tool.py

@@ -6,6 +6,7 @@
 
 import copy
 import logging
+import threading
 from typing import Any
 
 from typing_extensions import override
@@ -81,6 +82,10 @@ def __init__(
         # messages/state attributes.
         self._initial_messages: Messages = []
         self._initial_state: AgentState = AgentState()
+        # Serialize access so _reset_agent_state + stream_async are atomic.
+        # threading.Lock (not asyncio.Lock) because run_async() may create
+        # separate event loops in different threads.
+        self._lock = threading.Lock()
 
         if not preserve_context:
             from .agent import Agent
@@ -157,20 +162,38 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw
 
         tool_use_id = tool_use["toolUseId"]
 
-        # Determine if we are resuming the sub-agent from an interrupt.
-        if self._is_sub_agent_interrupted():
-            prompt = self._build_interrupt_responses()
-            logger.debug(
-                "tool_name=<%s>, tool_use_id=<%s> | resuming sub-agent from interrupt",
+        # Serialize access to the underlying agent. _reset_agent_state() mutates
+        # the agent before stream_async acquires its own lock, so a concurrent
+        # call would corrupt an in-flight invocation.
+        if not self._lock.acquire(blocking=False):
+            logger.warning(
+                "tool_name=<%s>, tool_use_id=<%s> | agent is already processing a request",
                 self._tool_name,
                 tool_use_id,
             )
-        elif not self._preserve_context:
-            self._reset_agent_state(tool_use_id)
-
-        logger.debug("tool_name=<%s>, tool_use_id=<%s> | invoking agent", self._tool_name, tool_use_id)
+            yield ToolResultEvent(
+                {
+                    "toolUseId": tool_use_id,
+                    "status": "error",
+                    "content": [{"text": f"Agent '{self._tool_name}' is already processing a request"}],
+                }
+            )
+            return
 
         try:
+            # Determine if we are resuming the sub-agent from an interrupt.
+            if self._is_sub_agent_interrupted():
+                prompt = self._build_interrupt_responses()
+                logger.debug(
+                    "tool_name=<%s>, tool_use_id=<%s> | resuming sub-agent from interrupt",
+                    self._tool_name,
+                    tool_use_id,
+                )
+            elif not self._preserve_context:
+                self._reset_agent_state(tool_use_id)
+
+            logger.debug("tool_name=<%s>, tool_use_id=<%s> | invoking agent", self._tool_name, tool_use_id)
+
             result = None
             async for event in self._agent.stream_async(prompt):
                 if "result" in event:
@@ -224,6 +247,8 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw
                     "content": [{"text": f"Agent error: {e}"}],
                 }
             )
+        finally:
+            self._lock.release()
 
     def _reset_agent_state(self, tool_use_id: str) -> None:
         """Reset the wrapped agent to its initial state.
@@ -250,11 +275,8 @@ def _reset_agent_state(self, tool_use_id: str) -> None:
 
     def _is_sub_agent_interrupted(self) -> bool:
         """Check whether the wrapped agent is in an activated interrupt state."""
-        from .agent import Agent
-
-        if not isinstance(self._agent, Agent):
-            return False
-        return self._agent._interrupt_state.activated
+        interrupt_state = getattr(self._agent, "_interrupt_state", None)
+        return interrupt_state is not None and interrupt_state.activated
 
     def _build_interrupt_responses(self) -> list[InterruptResponseContent]:
         """Build interrupt response payloads from the sub-agent's interrupt state.
@@ -266,14 +288,13 @@ def _build_interrupt_responses(self) -> list[InterruptResponseContent]:
         Returns:
             List of interrupt response content blocks for resuming the sub-agent.
         """
-        from .agent import Agent
-
-        if not isinstance(self._agent, Agent):
+        interrupt_state = getattr(self._agent, "_interrupt_state", None)
+        if interrupt_state is None:
             return []
 
         return [
             {"interruptResponse": {"interruptId": interrupt.id, "response": interrupt.response}}
-            for interrupt in self._agent._interrupt_state.interrupts.values()
+            for interrupt in interrupt_state.interrupts.values()
             if interrupt.response is not None
         ]
 

tests/strands/agent/test_agent_as_tool.py

@@ -23,6 +23,9 @@ def mock_agent():
     agent = MagicMock()
     agent.name = "test_agent"
     agent.description = "A test agent"
+    # Prevent MagicMock from auto-creating _interrupt_state on access,
+    # so getattr checks in AgentAsTool correctly detect its absence.
+    agent._interrupt_state = None
     return agent
 
 
@@ -615,3 +618,55 @@ async def test_build_interrupt_responses(fake_agent):
     # Only interrupt_a has a response
     assert len(responses) == 1
     assert responses[0] == {"interruptResponse": {"interruptId": "id-a", "response": "yes"}}
+
+
+# --- concurrency ---
+
+
+@pytest.mark.asyncio
+async def test_stream_rejects_concurrent_call(tool, mock_agent, tool_use, agent_result):
+    """A second concurrent call should get an error ToolResultEvent."""
+    mock_agent.stream_async.return_value = _mock_stream_async(agent_result)
+
+    # Simulate the lock already being held by another invocation
+    tool._lock.acquire()
+    try:
+        events = [event async for event in tool.stream(tool_use, {})]
+
+        assert len(events) == 1
+        assert isinstance(events[0], ToolResultEvent)
+        assert events[0]["tool_result"]["status"] == "error"
+        assert "already processing" in events[0]["tool_result"]["content"][0]["text"]
+        mock_agent.stream_async.assert_not_called()
+    finally:
+        tool._lock.release()
+
+
+@pytest.mark.asyncio
+async def test_stream_releases_lock_after_completion(tool, mock_agent, tool_use, agent_result):
+    """Lock should be released after stream completes, allowing subsequent calls."""
+    mock_agent.stream_async.return_value = _mock_stream_async(agent_result)
+
+    async for _ in tool.stream(tool_use, {}):
+        pass
+
+    assert not tool._lock.locked()
+
+    # A second call should succeed
+    mock_agent.stream_async.return_value = _mock_stream_async(agent_result)
+    events = [event async for event in tool.stream(tool_use, {})]
+
+    result_events = [e for e in events if isinstance(e, ToolResultEvent)]
+    assert len(result_events) == 1
+    assert result_events[0]["tool_result"]["status"] == "success"
+
+
+@pytest.mark.asyncio
+async def test_stream_releases_lock_after_error(tool, mock_agent, tool_use):
+    """Lock should be released even when the agent raises an exception."""
+    mock_agent.stream_async.side_effect = RuntimeError("boom")
+
+    async for _ in tool.stream(tool_use, {}):
+        pass
+
+    assert not tool._lock.locked()