diff --git a/src/praisonai-agents/praisonaiagents/agent/agent.py b/src/praisonai-agents/praisonaiagents/agent/agent.py index 0e13659ab..a116b6c22 100644 --- a/src/praisonai-agents/praisonaiagents/agent/agent.py +++ b/src/praisonai-agents/praisonaiagents/agent/agent.py @@ -3287,6 +3287,12 @@ Your Goal: {self.goal} if k in ("_db_run_context", "_db_session_init_lock"): # A clone starts outside the source agent's active DB turn. continue + elif k == "_tool_executor": + # Thread pools can't be copied; the clone lazily builds its own. + continue + elif k == "_Agent__hook_runner": + # Holds locks; rebuilt lazily on first use. + object.__setattr__(result, k, None) elif k in ("_Agent__cache_lock",): object.__setattr__(result, k, threading.RLock()) elif k in ("_cost_lock", "_auto_save_lock"): diff --git a/src/praisonai-agents/praisonaiagents/agent/tool_execution.py b/src/praisonai-agents/praisonaiagents/agent/tool_execution.py index 64311c736..4da112a27 100644 --- a/src/praisonai-agents/praisonaiagents/agent/tool_execution.py +++ b/src/praisonai-agents/praisonaiagents/agent/tool_execution.py @@ -49,8 +49,9 @@ _CALLER_ERROR_TYPES = ( # Worker count for the per-agent tool ThreadPoolExecutor. Used both when the # pool is created and to decide (with the per-executor in-flight counter) -# whether a fresh call would have to queue behind hung workers. -_TOOL_POOL_MAX_WORKERS = 2 +# whether a fresh call would have to queue behind hung workers. Matches the +# default ParallelToolCallExecutor width so a parallel batch never queues. +_TOOL_POOL_MAX_WORKERS = 5 # Upper bound on overflow ("spill") threads a single agent may run at once for # timed async tool calls that find every pool worker busy. A hung spill thread @@ -1461,8 +1462,10 @@ class ToolExecutionMixin: with tool_progress_channel(_attempt_sink): # Use copy_context to preserve injection context in executor thread ctx = contextvars.copy_context() + started = threading.Event() def execute_with_context(): + started.set() with with_injection_context(state): return self._execute_tool_with_circuit_breaker(function_name, arguments) @@ -1483,6 +1486,13 @@ class ToolExecutionMixin: } break try: + # Charge tool_timeout from when the body starts, + # not from submit, so time spent queued for a + # worker never times out a healthy tool. The + # queue wait is itself bounded by tool_timeout + # so a pool stuck behind hung workers still fails. + if not started.wait(timeout=tool_timeout): + raise concurrent.futures.TimeoutError() result = future.result(timeout=tool_timeout) except concurrent.futures.TimeoutError: if _progress_active is not None: diff --git a/src/praisonai-agents/praisonaiagents/hooks/runner.py b/src/praisonai-agents/praisonaiagents/hooks/runner.py index 7b6c5aa64..faf6b6ddf 100644 --- a/src/praisonai-agents/praisonaiagents/hooks/runner.py +++ b/src/praisonai-agents/praisonaiagents/hooks/runner.py @@ -11,6 +11,7 @@ import time import asyncio from praisonaiagents._logging import get_logger import subprocess +import threading from concurrent.futures import ThreadPoolExecutor from typing import List, Optional, Dict, Any @@ -78,25 +79,14 @@ class HookRunner: # misbehaving hook can only exhaust hook capacity, never the loop's # shared default executor that guardrails/tools/locks also rely on. self._hook_pool: Optional[ThreadPoolExecutor] = None - # hook.id -> in-flight Future for a sync hook. The slot is reserved - # before submission and only cleared when the worker finishes, so a - # still-running (timed-out or cancelled) hook caps new launches at one - # thread per hook. Guarded by ``_stuck_lock`` so concurrent invocations - # of the same hook cannot both reserve the slot and race past the cap. + # hook.id -> worker Future of a sync hook that ran past its timeout. + # It is cleared only when that worker finishes, so a stuck hook caps + # new launches at one leaked thread per hook. Healthy concurrent calls + # (e.g. parallel tool calls) are not tracked and never denied. A + # threading lock, because execute_sync() drives each call on its own + # event loop in the caller's thread. self._stuck: Dict[str, Any] = {} - self._stuck_lock_obj: Optional[asyncio.Lock] = None - - @property - def _stuck_lock(self) -> asyncio.Lock: - """Lazily create the in-flight-slot lock on first use. - - Created lazily (not in ``__init__``) because a ``HookRunner`` may be - constructed before any event loop exists; the lock is only needed once - we are already inside ``async def execute`` / ``_execute_function_hook``. - """ - if self._stuck_lock_obj is None: - self._stuck_lock_obj = asyncio.Lock() - return self._stuck_lock_obj + self._stuck_lock = threading.Lock() def _release_stuck(self, hook_id: str, fut: Any) -> None: """Clear the in-flight slot once the worker Future completes. @@ -104,8 +94,9 @@ class HookRunner: Only removes the entry if it still points at *this* Future, so a newer reserved invocation is never clobbered by an older worker draining. """ - if self._stuck.get(hook_id) is fut: - self._stuck.pop(hook_id, None) + with self._stuck_lock: + if self._stuck.get(hook_id) is fut: + self._stuck.pop(hook_id, None) def _get_hook_pool(self) -> ThreadPoolExecutor: """Lazily create the dedicated sync-hook executor.""" @@ -439,14 +430,9 @@ class HookRunner: timeout=timeout ) else: - # Reserve the single in-flight slot for this hook *under a lock* - # and *before* submitting, so two concurrent invocations can't - # both see an empty ``_stuck`` and both launch a worker (which - # would exceed the intended one-worker-per-hook cap and double - # the side effects). The slot is only cleared once the worker - # finishes; a still-running previous worker fails this one - # closed immediately instead of leaking another thread. - async with self._stuck_lock: + # A previous invocation that timed out and is still running + # fails this one closed instead of leaking another thread. + with self._stuck_lock: prev = self._stuck.get(hook.id) if prev is not None and not prev.done(): duration = (time.time() - start_time) * 1000 @@ -463,29 +449,31 @@ class HookRunner: ), ) - # Execute sync function in the dedicated hook pool with - # timeout. ``copy_context_to_callable`` preserves - # trace/session contextvars across the thread, matching - # other executor sites. - from ..trace.context_events import copy_context_to_callable - loop = asyncio.get_event_loop() - func = hook.func - data = input_data - fut = loop.run_in_executor( - self._get_hook_pool(), - copy_context_to_callable(lambda: func(data)), - ) - # Reserve the slot now; a completion callback clears it only - # when the worker actually finishes, so a timeout *or* a - # cancellation leaves the slot held until the thread drains. - self._stuck[hook.id] = fut - fut.add_done_callback( - lambda f, hid=hook.id: self._release_stuck(hid, f) - ) - - result = await asyncio.wait_for( - asyncio.shield(fut), timeout=timeout + # Execute sync function in the dedicated hook pool with + # timeout. ``copy_context_to_callable`` preserves + # trace/session contextvars across the thread, matching + # other executor sites. A concurrent Future (not a loop-bound + # one) so the stuck-slot release works after the loop closes. + from ..trace.context_events import copy_context_to_callable + func = hook.func + data = input_data + fut = self._get_hook_pool().submit( + copy_context_to_callable(lambda: func(data)) ) + try: + result = await asyncio.wait_for( + asyncio.shield(asyncio.wrap_future(fut)), timeout=timeout + ) + except (asyncio.TimeoutError, asyncio.CancelledError): + # The worker cannot be cancelled; hold the slot until it + # actually finishes. + if not fut.done(): + with self._stuck_lock: + self._stuck[hook.id] = fut + fut.add_done_callback( + lambda f, hid=hook.id: self._release_stuck(hid, f) + ) + raise duration = (time.time() - start_time) * 1000 diff --git a/src/praisonai-agents/praisonaiagents/mcp/mcp.py b/src/praisonai-agents/praisonaiagents/mcp/mcp.py index 7c3b80e2e..6394a2311 100644 --- a/src/praisonai-agents/praisonaiagents/mcp/mcp.py +++ b/src/praisonai-agents/praisonaiagents/mcp/mcp.py @@ -127,43 +127,42 @@ class MCPToolRunner(threading.Thread): # Signal that initialization is complete self.initialized.set() - # Process requests - while True: + # Process requests. Each request runs as its own task so + # one slow call never serializes the others behind it (the + # session multiplexes requests by id), and the queue is read + # with a short blocking get instead of a busy 10ms poll. The + # bounded get keeps the reader thread from outliving the loop. + loop = asyncio.get_running_loop() + tasks = set() + + def _next_item(): try: - # Check for new requests - try: - item = self.queue.get(block=False) - if item is None: # Shutdown signal - break - - response_queue, kind, name, arguments, cancelled = item - if cancelled.is_set(): - # Caller already gave up (timed out); never - # execute a call it was told did not happen. - continue - try: - if kind == "resource": - result = await asyncio.wait_for(session.read_resource(name), timeout=self.timeout) - elif kind == "prompt": - result = await asyncio.wait_for(session.get_prompt(name, arguments or None), timeout=self.timeout) - else: - result = await asyncio.wait_for(session.call_tool(name, arguments), timeout=self.timeout) - response_queue.put((True, result)) - except asyncio.TimeoutError: - response_queue.put((False, f"MCP {kind} call timed out after {self.timeout} seconds (server side)")) - except Exception as e: - if _is_transport_closed(e): - self._dead = f"MCP server connection lost: {str(e) or type(e).__name__}" - response_queue.put((False, self._dead)) - break # leave stdio_client() so the child is reaped - response_queue.put((False, str(e) or type(e).__name__)) - except queue.Empty: - pass - - # Give other tasks a chance to run - await asyncio.sleep(0.01) - except asyncio.CancelledError: - break + return self.queue.get(timeout=0.25) + except queue.Empty: + return queue.Empty + + try: + while True: + item = await loop.run_in_executor(None, _next_item) + if item is queue.Empty: + if self._dead: + break # leave stdio_client() so the child is reaped + continue + if item is None: # Shutdown signal + break + if self._dead: + self.queue.put(item) # answered by _fail_pending() + break + task = asyncio.create_task(self._handle_request(session, item)) + tasks.add(task) + task.add_done_callback(tasks.discard) + except asyncio.CancelledError: + pass + finally: + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) except asyncio.CancelledError: # stop() cancelled a startup that never completed; leaving the # contexts above has already terminated the stdio child. @@ -197,6 +196,35 @@ class MCPToolRunner(threading.Thread): except Exception: pass + async def _handle_request(self, session, item): + """Run one queued request on the session and post its result.""" + response_queue, kind, name, arguments, cancelled = item + if cancelled.is_set(): + # Caller already gave up (timed out); never execute a call it was + # told did not happen. + return + try: + if kind == "resource": + coro = session.read_resource(name) + elif kind == "prompt": + coro = session.get_prompt(name, arguments or None) + else: + coro = session.call_tool(name, arguments) + result = await asyncio.wait_for(coro, timeout=self.timeout) + response_queue.put((True, result)) + except asyncio.TimeoutError: + response_queue.put((False, f"MCP {kind} call timed out after {self.timeout} seconds (server side)")) + except asyncio.CancelledError: + # Report the real cause when a sibling call saw the transport die. + response_queue.put((False, self._dead or f"MCP {kind} call cancelled: runner shut down")) + raise + except Exception as e: + if _is_transport_closed(e): + self._dead = f"MCP server connection lost: {str(e) or type(e).__name__}" + response_queue.put((False, self._dead)) + return + response_queue.put((False, str(e) or type(e).__name__)) + async def _discover_resources_and_prompts(self, session): """Enumerate server resources, resource templates and prompts. @@ -391,8 +419,8 @@ class MCPToolRunner(threading.Thread): Queues the shutdown sentinel and joins the thread so the request loop exits and its stdio subprocess is terminated instead of leaking a background thread (and its child) across the process lifetime. The - request loop polls the queue between async operations, so once any - in-flight call unwinds the sentinel is consumed promptly. + request loop reads the sentinel promptly and cancels any in-flight + calls, which are answered with an error instead of left waiting. Args: timeout: Seconds to wait for the thread to exit. Defaults to the diff --git a/src/praisonai-agents/tests/unit/agent/test_issue_5795_tool_concurrency.py b/src/praisonai-agents/tests/unit/agent/test_issue_5795_tool_concurrency.py new file mode 100644 index 000000000..2a77f99a0 --- /dev/null +++ b/src/praisonai-agents/tests/unit/agent/test_issue_5795_tool_concurrency.py @@ -0,0 +1,162 @@ +"""Regression tests for #5795: tool/MCP calls must not be timed out by queue wait.""" +import asyncio +import concurrent.futures as cf +import queue +import threading +import time +from contextlib import asynccontextmanager +from types import SimpleNamespace + +from praisonaiagents import Agent +from praisonaiagents.hooks import HookEvent, HookRegistry + + +def _allow_all_hooks(): + """Own registry (no leaked global hooks) with a sync BEFORE_TOOL hook, + so parallel tool calls also run the hook concurrently.""" + registry = HookRegistry() + registry.register_function(HookEvent.BEFORE_TOOL, lambda data: None) + return registry + + +def test_parallel_sync_tools_not_timed_out_by_queue_wait(): + def slow(n: int) -> str: + """slow tool""" + time.sleep(0.5) + return f"ok{n}" + + agent = Agent(name="t", instructions="x", tools=[slow], llm="gpt-4o-mini", hooks=_allow_all_hooks()) + agent._tool_timeout = 0.8 + with cf.ThreadPoolExecutor(5) as ex: + results = list(ex.map(lambda i: agent.execute_tool("slow", {"n": i}), range(5))) + for i, result in enumerate(results): + assert isinstance(result, str) and result.startswith(f"ok{i}"), result + + +def test_parallel_async_sync_tools_not_timed_out_by_queue_wait(): + def slow(n: int) -> str: + """slow tool""" + time.sleep(0.5) + return f"ok{n}" + + agent = Agent(name="a", instructions="x", tools=[slow], llm="gpt-4o-mini", hooks=_allow_all_hooks()) + agent._tool_timeout = 0.8 + + async def main(): + # Twice the pool size: the second half queues for a worker first. + return await asyncio.gather( + *(agent.execute_tool_async("slow", {"n": i}) for i in range(10)) + ) + + for i, result in enumerate(asyncio.run(main())): + assert isinstance(result, str) and result.startswith(f"ok{i}"), result + + +def test_hung_async_sync_tools_do_not_starve_default_executor(): + release = threading.Event() + tools = [] + for i in range(10): + def hang(x: int = 0) -> str: + """hang""" + release.wait(10) + return "x" + hang.__name__ = f"hang{i}" + tools.append(hang) + + agent = Agent(name="u", instructions="x", tools=tools, llm="gpt-4o-mini", hooks=HookRegistry()) + agent._tool_timeout = 0.2 + + async def main(): + for i in range(10): + result = await asyncio.wait_for(agent.execute_tool_async(f"hang{i}", {"x": 1}), 5) + # Times out, or is refused fast once hung calls hit the pool cap. + assert isinstance(result, dict) and (result.get("timeout") or result.get("tool_pool_exhausted")), result + # The loop's default executor is still free for other to_thread users. + assert await asyncio.wait_for(asyncio.to_thread(lambda: 42), 2) == 42 + + try: + asyncio.run(main()) + finally: + release.set() + + +def test_transport_wrapper_cancels_call_on_timeout(): + from praisonaiagents.mcp.mcp_sse import SSEMCPTool + + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_forever, daemon=True).start() + finished = threading.Event() + + class _Session: + async def call_tool(self, name, kwargs): + await asyncio.sleep(1.0) + finished.set() + + try: + tool = SSEMCPTool("t", "d", _Session(), {}, timeout=0.1) + tool._loop = loop + result = tool() + assert isinstance(result, dict) and result.get("timeout"), result + time.sleep(1.3) + assert not finished.is_set() + finally: + loop.call_soon_threadsafe(loop.stop) + + +def test_stdio_runner_runs_requests_concurrently_and_cancels_on_shutdown(monkeypatch): + import praisonaiagents.mcp.mcp as mcp_mod + + class _Session: + async def initialize(self): + pass + + async def list_tools(self): + return SimpleNamespace(tools=[]) + + async def call_tool(self, name, arguments): + await asyncio.sleep(10 if name == "slow" else 0) + return name + + @asynccontextmanager + async def _stdio_client(params): + yield None, None + + @asynccontextmanager + async def _client_session(read, write): + yield _Session() + + async def _no_discovery(self, session): + pass + + monkeypatch.setattr(mcp_mod, "stdio_client", _stdio_client) + monkeypatch.setattr(mcp_mod, "ClientSession", _client_session) + monkeypatch.setattr(mcp_mod.MCPToolRunner, "_discover_resources_and_prompts", _no_discovery) + + runner = mcp_mod.MCPToolRunner(None, timeout=30) + assert runner.initialized.wait(5) and runner._init_error is None + + slow_q, fast_q = queue.Queue(), queue.Queue() + runner.queue.put((slow_q, "tool", "slow", {}, threading.Event())) + runner.queue.put((fast_q, "tool", "fast", {}, threading.Event())) + # The fast call is not serialized behind the in-flight slow one. + assert fast_q.get(timeout=2) == (True, "fast") + + runner.queue.put(None) + ok, message = slow_q.get(timeout=2) + assert not ok and "shut down" in message + runner.join(2) + assert not runner.is_alive() + + +def test_agent_deepcopy_after_async_tool_use(): + import copy + + def echo(x: int) -> int: + """echo""" + return x + + agent = Agent(name="d", instructions="x", tools=[echo], llm="gpt-4o-mini") + assert asyncio.run(agent.execute_tool_async("echo", {"x": 1})) == 1 + clone = copy.deepcopy(agent) + assert "_tool_executor" not in clone.__dict__ + assert asyncio.run(clone.execute_tool_async("echo", {"x": 2})) == 2