diff --git a/CHANGELOG.md b/CHANGELOG.md index 0067c7d..8c795dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- Moved full MCP tool call / callback inputs and outputs, and `execute_typescript`/`execute_bash` code and output, from INFO to DEBUG logs to declutter traces. Added DEBUG logs describing tool result shape (structured content, JSON parse success/fallback). - **Breaking:** `CodeMode::with_callbacks` returns `(Self, CallbackReport)` instead of `Result`, so builder-style callers see which tools failed or degraded. Per-tool isolation means the batch itself cannot fail, so the report is the only outcome. - **Breaking:** `CodeMode::add_callback` returns `Result>` — the reasons that tool's types were degraded to `any`, empty when fully typed. diff --git a/crates/pctx_code_mode/src/code_mode.rs b/crates/pctx_code_mode/src/code_mode.rs index 2597029..fe7e8c2 100644 --- a/crates/pctx_code_mode/src/code_mode.rs +++ b/crates/pctx_code_mode/src/code_mode.rs @@ -513,7 +513,7 @@ impl CodeMode { } /// Execute bash commands directly in the virtual filesystem - #[instrument(skip(self), ret(Display), err)] + #[instrument(skip(self), err)] pub async fn execute_bash(&self, command: &str) -> Result { debug!(command = %command, "Executing bash command"); @@ -578,15 +578,18 @@ export default result;"#, warn!("Bash execution failed with exit code {exit_code}: {stderr}"); } - Ok(ExecuteBashOutput { + let output = ExecuteBashOutput { exit_code, stdout, stderr, - }) + }; + debug!("Bash execution result:\n{output}"); + + Ok(output) } /// Execute TypeScript code with access to registered tools and virtual filesystem - #[instrument(skip(self, registry), ret(Display), err)] + #[instrument(skip(self, registry, code), err)] pub async fn execute_typescript( &self, code: &str, @@ -698,7 +701,7 @@ export default result;"#, } }; - debug!(to_execute = %to_execute, "Executing TypeScript in sandbox"); + debug!("Executing TypeScript in sandbox:\n{to_execute}"); let execution_res = pctx_executor::execute( &to_execute, @@ -712,14 +715,18 @@ export default result;"#, warn!("TypeScript execution failed: {:?}", execution_res.stderr); } - Ok(ExecuteTypescriptOutput { + let output = ExecuteTypescriptOutput { success: execution_res.success, stdout: execution_res.stdout, stderr: execution_res.stderr, output: execution_res.output, registry: execution_res.registry, trace: execution_res.trace, - }) + }; + + debug!("TypeScript execution result:\n{output}"); + + Ok(output) } } diff --git a/crates/pctx_registry/src/registry.rs b/crates/pctx_registry/src/registry.rs index 1e7919b..ec52907 100644 --- a/crates/pctx_registry/src/registry.rs +++ b/crates/pctx_registry/src/registry.rs @@ -15,7 +15,7 @@ use std::{ sync::{Arc, RwLock}, time::SystemTime, }; -use tracing::{debug, info, instrument, warn}; +use tracing::{debug, instrument, warn}; pub type CallbackFn = Arc< dyn Fn( @@ -237,13 +237,7 @@ impl PctxRegistry { /// /// This function will return an error if an action by the provided id doesn't exist /// or if the action itself fails - #[instrument( - name = "invoke_registry_action", - skip_all, - fields(id=id, args = json!(args).to_string()), - ret(Display), - err - )] + #[instrument(name = "invoke_registry_action", skip_all, fields(id = id), err)] pub async fn invoke( &self, id: &str, @@ -333,25 +327,40 @@ impl PctxRegistry { } // Prefer structuredContent if available, otherwise use content array - let has_structured = tool_result.structured_content.is_some(); let val = if let Some(structured) = tool_result.structured_content { + debug!(tool = %mcp_id.id(), "tool result: using structured content"); structured } else if let Some(RawContent::Text(text_content)) = tool_result.content.first().map(|a| &**a) { // Try to parse as JSON, fallback to string value - serde_json::from_str(&text_content.text) - .or_else(|_| Ok(serde_json::Value::String(text_content.text.clone()))) - .map_err(|e: serde_json::Error| { - RegistryError::ToolCall(format!("Failed to parse content: {e}")) - })? + match serde_json::from_str(&text_content.text) { + Ok(json) => { + debug!( + tool = %mcp_id.id(), + "tool result: parsed text content as JSON" + ); + json + } + Err(e) => { + debug!( + tool = %mcp_id.id(), + error = %e, + "tool result: text content is not JSON, using raw string" + ); + serde_json::Value::String(text_content.text.clone()) + } + } } else { // Return the whole content array as JSON + debug!( + tool = %mcp_id.id(), + content_len = tool_result.content.len(), + "tool result: no structured or text content, using raw content array" + ); json!(tool_result.content) }; - info!(structured_content = has_structured, result =? &val, "Tool result"); - Ok(val) })(); diff --git a/pctx-py/CHANGELOG.md b/pctx-py/CHANGELOG.md index 3281618..aa2ae38 100644 --- a/pctx-py/CHANGELOG.md +++ b/pctx-py/CHANGELOG.md @@ -17,6 +17,14 @@ For changes to the underlying Rust crates and CLI, see the ### Fixed +- Tool calls now run concurrently. Each request is handled in its own task + rather than awaited inside the WebSocket read loop, so code fanning out with + `Promise.all` takes the time of its slowest call instead of the sum of all + of them. +- Sync tools run on a worker thread instead of blocking the event loop, so one + slow sync tool no longer stalls the calls beside it. Their bodies now execute + off the main thread, so anything they share must be thread-safe. + ## [v0.4.4] - 2026-07-22 ### Fixed diff --git a/pctx-py/src/pctx_client/_websocket_client.py b/pctx-py/src/pctx_client/_websocket_client.py index 29c350e..6926fc1 100644 --- a/pctx-py/src/pctx_client/_websocket_client.py +++ b/pctx-py/src/pctx_client/_websocket_client.py @@ -68,6 +68,10 @@ def __init__( self._headers = headers or {} self._pending_executions: dict[str | int, asyncio.Future] = {} self._request_counter = 0 + self._message_handler_task: asyncio.Task | None = None + # In-flight tool executions. Held strongly because asyncio only keeps + # weak references to tasks, and cancelled as a group on disconnect. + self._tool_tasks: set[asyncio.Task] = set() async def _connect(self, code_mode_session: str): """ @@ -100,6 +104,10 @@ async def _disconnect(self): if self._message_handler_task: self._message_handler_task.cancel() + for task in self._tool_tasks: + task.cancel() + self._tool_tasks.clear() + if self.ws: await self.ws.close() self.ws = None @@ -192,8 +200,14 @@ async def _handle_messages(self): message: WebSocketMessage = adapter.validate_json(message_data) if isinstance(message, ExecuteToolRequest): - res = await self._handle_execute_tool(message) - await self._send(res) + # Run the tool in its own task so this loop stays free to + # read the next message. Awaiting it here would serialize + # every tool call the server dispatches, so code that fans + # out with `Promise.all` would take the sum of its calls + # rather than the slowest one. + task = asyncio.create_task(self._execute_tool(message)) + self._tool_tasks.add(task) + task.add_done_callback(self._tool_tasks.discard) elif isinstance(message, ExecuteCodeResponse): future = self._pending_executions.get(message.id) if future is not None: @@ -216,6 +230,16 @@ async def _handle_messages(self): except Exception as e: print(f"Message handler error: {e}") + async def _execute_tool(self, req: ExecuteToolRequest): + """Run one tool request and send its response back to the server.""" + try: + res = await self._handle_execute_tool(req) + await self._send(res) + except asyncio.CancelledError: + raise + except Exception as e: + print(f"Error executing tool: {e}") + async def _handle_execute_tool( self, req: ExecuteToolRequest ) -> ExecuteToolResponse | JsonRpcError: @@ -240,10 +264,14 @@ async def _handle_execute_tool( args = req.params.args or {} try: if isinstance(tool, Tool): + # Sync tools go to a worker thread. Calling one inline would + # block the event loop for its whole duration, which stalls + # every other in-flight tool call behind it -- and with it the + # loop reading further requests off the WebSocket. if tool.input_schema is None: - output = tool.invoke() + output = await asyncio.to_thread(tool.invoke) else: - output = tool.invoke(**args) + output = await asyncio.to_thread(tool.invoke, **args) else: if tool.input_schema is None: output = await tool.ainvoke() diff --git a/pctx-py/tests/test_integration.py b/pctx-py/tests/test_integration.py index afa21ff..efbe796 100644 --- a/pctx-py/tests/test_integration.py +++ b/pctx-py/tests/test_integration.py @@ -1,5 +1,8 @@ """Integration tests for pctx code mode against a running server""" +import asyncio +import threading +import time from datetime import datetime import pytest @@ -796,3 +799,150 @@ def format_result(value: int, label: str) -> str: "Please ensure the pctx server is running.\n" "Start the server with: pctx server start" ) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_concurrent_async_tool_calls_run_in_parallel(): + """Tools fanned out with `Promise.all` must execute concurrently. + + The client used to await each tool request inside its WebSocket read loop, + so a batch of N calls took the sum of their durations instead of the + slowest one. Assert on observed overlap rather than wall time alone, so + the test fails on serialization rather than on a slow machine. + """ + sleep_secs = 0.5 + calls = 4 + + in_flight = 0 + max_in_flight = 0 + + @tool + async def slow_echo(value: int) -> int: + """Sleep briefly, then echo the value back""" + nonlocal in_flight, max_in_flight + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + try: + await asyncio.sleep(sleep_secs) + return value + finally: + in_flight -= 1 + + try: + async with Pctx(tools=[slow_echo], execute_timeout=60) as pctx: + code = """ + async function run() { + const values = await Promise.all([ + Tools.slowEcho({ value: 1 }), + Tools.slowEcho({ value: 2 }), + Tools.slowEcho({ value: 3 }), + Tools.slowEcho({ value: 4 }), + ]); + return { values }; + } + """ + + start = time.perf_counter() + output = await pctx.execute_typescript(code) + elapsed = time.perf_counter() - start + + assert output.success, f"Execution should succeed, got: {output.stderr}" + assert output.output is not None, "Execution should return output" + assert output.output.get("values") == [1, 2, 3, 4], ( + f"Expected all four calls to return, got: {output.output}" + ) + + assert max_in_flight == calls, ( + f"All {calls} tool calls should be in flight at once, " + f"peaked at {max_in_flight} -- the client is serializing them" + ) + + # Serialized dispatch takes calls * sleep_secs; concurrent dispatch + # takes ~sleep_secs. Half way between the two is a wide enough + # margin to absorb session setup and round-trip overhead. + serial_secs = calls * sleep_secs + assert elapsed < serial_secs / 2, ( + f"Concurrent calls took {elapsed:.2f}s; serialized dispatch " + f"would take ~{serial_secs:.2f}s" + ) + except ConnectionError: + pytest.fail( + "Failed to connect to pctx server at http://localhost:8080.\n" + "Please ensure the pctx server is running.\n" + "Start the server with: pctx server start" + ) + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_concurrent_sync_tool_calls_run_in_parallel(): + """Sync tools fanned out with `Promise.all` must also run concurrently. + + A sync tool body blocks whatever thread it runs on, so calling it inline + on the event loop would stall every other in-flight call behind it. They + run on worker threads instead -- hence the blocking `time.sleep` here, and + the lock around the counters, which the tool bodies touch off-thread. + """ + sleep_secs = 0.5 + calls = 4 + + lock = threading.Lock() + in_flight = 0 + max_in_flight = 0 + + @tool + def slow_echo_sync(value: int) -> int: + """Block briefly, then echo the value back""" + nonlocal in_flight, max_in_flight + with lock: + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + try: + time.sleep(sleep_secs) + return value + finally: + with lock: + in_flight -= 1 + + try: + async with Pctx(tools=[slow_echo_sync], execute_timeout=60) as pctx: + code = """ + async function run() { + const values = await Promise.all([ + Tools.slowEchoSync({ value: 1 }), + Tools.slowEchoSync({ value: 2 }), + Tools.slowEchoSync({ value: 3 }), + Tools.slowEchoSync({ value: 4 }), + ]); + return { values }; + } + """ + + start = time.perf_counter() + output = await pctx.execute_typescript(code) + elapsed = time.perf_counter() - start + + assert output.success, f"Execution should succeed, got: {output.stderr}" + assert output.output is not None, "Execution should return output" + assert output.output.get("values") == [1, 2, 3, 4], ( + f"Expected all four calls to return, got: {output.output}" + ) + + assert max_in_flight == calls, ( + f"All {calls} tool calls should be in flight at once, " + f"peaked at {max_in_flight} -- sync tools are blocking the " + f"event loop instead of running on worker threads" + ) + + serial_secs = calls * sleep_secs + assert elapsed < serial_secs / 2, ( + f"Concurrent calls took {elapsed:.2f}s; serialized dispatch " + f"would take ~{serial_secs:.2f}s" + ) + except ConnectionError: + pytest.fail( + "Failed to connect to pctx server at http://localhost:8080.\n" + "Please ensure the pctx server is running.\n" + "Start the server with: pctx server start" + ) diff --git a/pctx-py/uv.lock b/pctx-py/uv.lock index bdd086c..efa1ffd 100644 --- a/pctx-py/uv.lock +++ b/pctx-py/uv.lock @@ -3459,7 +3459,7 @@ wheels = [ [[package]] name = "pctx-client" -version = "0.4.3" +version = "0.4.4" source = { editable = "." } dependencies = [ { name = "docstring-parser" }, diff --git a/plans/tool-call-concurrency-controls.md b/plans/tool-call-concurrency-controls.md new file mode 100644 index 0000000..ed97f04 --- /dev/null +++ b/plans/tool-call-concurrency-controls.md @@ -0,0 +1,66 @@ +# Tool call concurrency controls + +## Problem +Tool calls used to be dispatched **one at a time**: the Python client awaited each tool inside its WebSocket read loop, so it never read the next request until the current tool finished. Code fanning out with `Promise.all` took the *sum* of its calls instead of the slowest one. Everything above the client was already concurrent — the Deno op is `#[op2(async)]` ([invoke_ops.rs:13](../crates/pctx_code_execution_runtime/src/invoke_ops.rs#L13)), the registry holds no lock across its await ([registry.rs:241](../crates/pctx_registry/src/registry.rs#L241)), and the server gives each call its own request id and response channel ([ws_manager.rs:146](../crates/pctx_session_server/src/state/ws_manager.rs#L146)). The requests left the server together and queued at the client. + +## Current solution (shipped) +[`_websocket_client.py`](../pctx-py/src/pctx_client/_websocket_client.py): +1. Each `ExecuteToolRequest` runs in its own `asyncio.Task`, so the read loop stays free. Tasks are held in a set (asyncio keeps only weak references) and cancelled as a group on disconnect. +2. Sync tools run via `asyncio.to_thread` — calling one inline blocked the event loop for its whole duration, stalling every other in-flight call behind it. + +Measured, 4 × 2s tool under `Promise.all`: **8.12s → 2.11s**, all four starting at t=0.14. Covered by `test_concurrent_tool_calls_run_in_parallel` and `test_concurrent_sync_tool_calls_run_in_parallel` ([test_integration.py](../pctx-py/tests/test_integration.py)), which assert on observed overlap rather than wall time so they fail on serialization, not on a slow machine. + +**Dispatch is now unbounded for async tools** — N concurrent calls in JS means N concurrent tasks in Python. What follows is about bounding that. + +--- + +## The constraint that shapes every option + +The server's per-tool timeout starts at **dispatch**, not at execution: [`execute_callback`](../crates/pctx_session_server/src/state/ws_manager.rs#L146) sends the request, *then* starts `tokio::time::timeout`. Any client-side queue silently spends the tool's timeout budget waiting. + +Measured — 40 sync calls × 2s work, `tool_timeout_secs=5`: + +``` +starts: 14 @ 0.19s | 14 @ 2.2s | 12 @ 4.2s +result: ok=28, failed=12 + 'Tool `tools__slow_sync` timed out after 5s' +``` + +12 calls died on a 5s deadline while doing 2s of work, purely from queue wait. **A concurrency cap converts a throughput problem into hard failures.** Any cap has to be introduced with this in mind. + +## Where we are today + +Measured on a 10-core machine, 40-way fan-out: + +| | peak in-flight | wall | +|---|---|---| +| async tools | 40 (unbounded) | 0.66s | +| sync tools | **14** | 1.68s | + +The 14 is `min(32, cpu_count + 4)` — asyncio's default `ThreadPoolExecutor`, inherited via `to_thread`. So there is **already a cap on the sync path**, it is just accidental and machine-dependent: 14 here, 6 on a 2-core CI box, moving the timeout knee with the hardware. + +--- + +## Options + +| Option | How | Pros | Cons | +|---|---|---|---| +| **A. Explicit sync executor** | Dedicated `ThreadPoolExecutor` with a configured size instead of asyncio's default | Replaces an invisible machine-dependent cap with a stated one; no behaviour change where `cpu+4 == size` | Doesn't bound async tools | +| **B. Opt-in `max_concurrent_tools`** | Semaphore around tool dispatch, default `None` (unbounded) | Available for rate-limited upstreams; preserves current behaviour by default | Queue wait burns the server timeout (above); one slow tool blocks unrelated fast ones | +| **C. Server sends its deadline** | Include the deadline in `ExecuteToolRequest` | Client can fail fast with a real reason, or drop work it can't service in time; makes caps *safe* rather than merely possible | Protocol + server change; both sides must ship | +| **D. Per-tool limits in the tool body** | Caller puts a semaphore inside their own tool function | Correctly scoped per upstream; no client change | Not discoverable; every caller reimplements it | + +### Why not a global semaphore on its own +A single cap means a slow Excel call blocks an unrelated fast Graph call. A client-level cap is a **resource safety valve** (sockets, memory, threads), not a throttle. Rate limiting is per-upstream and belongs where the upstream is known — **D**, or **B** scoped per tool rather than globally. + +--- + +## Recommendation +Ship the current fix as-is; it is a strict improvement and adds no cap where none existed. + +Then, cheapest-first: +1. **A** — make the sync cap explicit. It's the one cap that already exists and already bites, silently and differently per machine. +2. **B**, defaulting to unbounded. Given the timeout coupling, unbounded is the correct default: everything starts immediately and gets its full budget. Offer the knob to callers protecting an upstream, and document that `tool_timeout_secs` must cover *queue wait plus call duration*, not just call duration. +3. **C** if caps become load-bearing. Without the deadline, a queued client is blind to a clock already running against it, and overload surfaces as mystery timeouts. + +Treat **D** as the documented answer for per-upstream rate limiting regardless of whether B lands.