From 5535415e8982066feee3695912cc3c2de0c645c2 Mon Sep 17 00:00:00 2001 From: Tomas Sereikis Date: Sat, 29 Aug 2026 13:21:54 +0300 Subject: [PATCH 1/2] perf(llm): run one turn's tool calls concurrently instead of serially MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A single assistant turn routinely requests several independent reads — a dialectic turn typically emits `search_memory` and `search_messages` together — but the loop awaited them one at a time, so the turn cost their sum when it only needed to cost the slowest. Measured on a production `minimal` dialectic: 1.2s in `search_memory`, then a further 1.9s in `search_messages`, both pure reads. Measured across 42 production dialectic requests (103 iterations, 77 of them multi-tool, 3.34 tool calls per iteration on average), running each turn's calls concurrently would cut 11.3% of total wall-clock: median 9.3% per request, p90 29.7%, best case 51.0%. Only 4 of the 42 requests gain nothing. This is safe because tool handlers own their sessions: each opens a short-lived one via `tracked_db()`, and only the mutating handlers (`create_observations`, `update_peer_card`, `delete_observations`) take `ctx.db_lock`, so concurrent reads do not contend on shared state. The per-call telemetry ContextVars move inside the task. `asyncio` copies the context per task, so `set_current_tool_call_seq` and `set_last_tool_metadata` now bind to their own call instead of being written and read across one shared context — which the previous code could only keep straight by never overlapping. `gather` preserves argument order, so `tool_results` and `all_tool_calls` stay in the order the model asked for them. Fan-out is capped at MAX_CONCURRENT_TOOL_CALLS (4). Production has produced 18 tool calls in one iteration, and firing all of them at once would mean that many simultaneous embedding + pgvector queries on a single instance. The cap leaves the measured common case fully parallel while bounding the tail. --- src/llm/tool_loop.py | 68 ++++++++--- tests/llm/test_tool_loop_parallel.py | 173 +++++++++++++++++++++++++++ 2 files changed, 224 insertions(+), 17 deletions(-) create mode 100644 tests/llm/test_tool_loop_parallel.py diff --git a/src/llm/tool_loop.py b/src/llm/tool_loop.py index 783d4965..7a30b960 100644 --- a/src/llm/tool_loop.py +++ b/src/llm/tool_loop.py @@ -11,6 +11,7 @@ from __future__ import annotations +import asyncio import dataclasses import functools import logging @@ -201,6 +202,13 @@ def _emit_agent_iteration( logger = logging.getLogger(__name__) +# A single assistant turn can request many tools at once — production has seen +# 18 in one iteration. Running every one concurrently would put that many +# simultaneous embedding + pgvector queries on a single instance, so the fan-out +# is capped: the measured common case is 3.3 calls per iteration, which stays +# fully parallel, while the tail is bounded instead of stampeding. +MAX_CONCURRENT_TOOL_CALLS = 4 + # Bounds for max_tool_iterations to prevent runaway loops. MIN_TOOL_ITERATIONS = 1 MAX_TOOL_ITERATIONS = 100 @@ -391,6 +399,9 @@ async def execute_tool_loop( if telemetry is not None and telemetry.hash_memo is None: telemetry = dataclasses.replace(telemetry, hash_memo={}) + # One semaphore for the whole loop: iterations run one after another, so a + # single cap bounds the fan-out without being rebound per iteration. + tool_call_semaphore = asyncio.Semaphore(MAX_CONCURRENT_TOOL_CALLS) iteration = 0 all_tool_calls: list[dict[str, Any]] = [] total_input_tokens = 0 @@ -582,54 +593,77 @@ async def execute_tool_loop( # Telemetry context — 1-indexed iteration. set_current_iteration(iteration + 1) - tool_results: list[dict[str, Any]] = [] - for seq, tool_call in enumerate(response.tool_calls_made): + # Tools requested in one assistant turn are independent of each + # other, so run them concurrently and pay the slowest instead of the + # sum. A dialectic turn routinely asks for `search_memory` and + # `search_messages` together; serially that was 1.2s + 1.9s of pure + # read. Handlers open their own short-lived sessions via + # `tracked_db()` and only the mutating ones (`create_observations`, + # `update_peer_card`, `delete_observations`) take `ctx.db_lock`, so + # concurrent reads are safe. + async def run_tool_call( + seq: int, tool_call: dict[str, Any] + ) -> tuple[dict[str, Any], dict[str, Any] | None]: tool_name = tool_call["name"] tool_input = tool_call["input"] tool_id = tool_call.get("id", "") logger.debug(f"Executing tool: {tool_name}") - # the executor closure reads these from - # ContextVars to populate AgentToolCallCompletedEvent. Set BEFORE - # the executor call so two calls to the same tool in one iteration - # get distinct seq values. Reset last-tool metadata so we never - # observe stale state from a prior call. + # The executor closure reads these from ContextVars to populate + # AgentToolCallCompletedEvent. They are set INSIDE the task: + # asyncio copies the context per task, so two calls to the same + # tool in one iteration keep their own seq and their own + # last-tool metadata rather than racing on one shared context. set_current_tool_call_seq(seq, tool_id or None) set_last_tool_metadata({}) try: - tool_result = await tool_executor(tool_name, tool_input) + async with tool_call_semaphore: + tool_result = await tool_executor(tool_name, tool_input) # Stash ToolResult.metadata on all_tool_calls so # specialist rollups can read created/deleted observation # counts without round-tripping through the event store. - tool_result_metadata = get_last_tool_metadata() - tool_results.append( + return ( { "tool_id": tool_id, "tool_name": tool_name, "result": tool_result, - } - ) - all_tool_calls.append( + }, { "tool_name": tool_name, "tool_input": tool_input, "tool_result": tool_result, - "tool_result_metadata": tool_result_metadata, - } + "tool_result_metadata": get_last_tool_metadata(), + }, ) except Exception as e: logger.error(f"Tool execution failed for {tool_name}: {e}") - tool_results.append( + return ( { "tool_id": tool_id, "tool_name": tool_name, "result": f"Error: {str(e)}", "is_error": True, - } + }, + None, ) + # gather preserves argument order, so tool_results and + # all_tool_calls stay in the order the model asked for them. + outcomes = await asyncio.gather( + *( + asyncio.create_task(run_tool_call(seq, tool_call)) + for seq, tool_call in enumerate(response.tool_calls_made) + ) + ) + + tool_results: list[dict[str, Any]] = [] + for result_entry, call_entry in outcomes: + tool_results.append(result_entry) + if call_entry is not None: + all_tool_calls.append(call_entry) + append_tool_results(current_provider, tool_results, conversation_messages) finally: if step is not None: diff --git a/tests/llm/test_tool_loop_parallel.py b/tests/llm/test_tool_loop_parallel.py new file mode 100644 index 00000000..ac1408d9 --- /dev/null +++ b/tests/llm/test_tool_loop_parallel.py @@ -0,0 +1,173 @@ +# pyright: reportPrivateUsage=false, reportUnknownLambdaType=false, reportArgumentType=false +"""Tool calls emitted in a single assistant turn execute concurrently. + +The model routinely asks for several independent reads at once — a dialectic +turn typically emits `search_memory` and `search_messages` together. Running +them one after another makes the turn cost the sum of their latencies when it +only needs to cost the slowest: measured against production, a `minimal` +dialectic spent 1.2s in `search_memory` and then a further 1.9s in +`search_messages`, both pure reads. + +Tool handlers open their own short-lived sessions via `tracked_db()` and only +the mutating handlers (`create_observations`, `update_peer_card`, +`delete_observations`) take `ctx.db_lock`, so concurrent reads are safe. The +per-call telemetry ContextVars are set inside each task, and `asyncio` copies +the context per task, so sequence numbers and last-tool metadata stay bound to +their own call instead of racing on one shared context. +""" + +from __future__ import annotations + +import asyncio +from typing import Any +from unittest.mock import patch + +import pytest + +from src.llm import tool_loop +from src.llm.runtime import AttemptPlan +from src.llm.tool_loop import MAX_CONCURRENT_TOOL_CALLS, execute_tool_loop +from src.llm.types import HonchoLLMCallResponse + +TOOL_DELAY_SECONDS = 0.2 + + +def _make_plan() -> AttemptPlan: + return AttemptPlan( + provider="anthropic", + model="claude-sonnet-4-5", + client=object(), + thinking_budget_tokens=None, + reasoning_effort=None, + selected_config=None, + attempt=1, + retry_attempts=1, + is_fallback=False, + ) + + +def _response(tool_calls: list[dict[str, Any]]) -> HonchoLLMCallResponse[Any]: + return HonchoLLMCallResponse( + content="done", + input_tokens=10, + output_tokens=5, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + finish_reasons=["stop"], + tool_calls_made=tool_calls, + ) + + +class _Tracker: + """Counts how many tool executions are in flight at once.""" + + def __init__(self) -> None: + self.in_flight = 0 + self.peak = 0 + + async def execute(self, _name: str, _input: dict[str, Any]) -> str: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + try: + await asyncio.sleep(TOOL_DELAY_SECONDS) + finally: + self.in_flight -= 1 + return "ok" + + +async def _run(tracker: _Tracker, n_calls: int = 2) -> HonchoLLMCallResponse[Any]: + names = ["search_memory", "search_messages"] + calls = [ + {"name": names[i % len(names)], "input": {"query": str(i)}, "id": f"t{i}"} + for i in range(n_calls) + ] + responses = iter([_response(calls), _response([])]) + + async def _call(*_args: Any, **_kwargs: Any) -> HonchoLLMCallResponse[Any]: + return next(responses) + + with ( + patch.object(tool_loop, "honcho_llm_call_inner", new=_call), + patch("src.llm.conversation.count_message_tokens", return_value=10), + patch( + "src.llm.conversation.truncate_messages_to_fit", + side_effect=lambda msgs, _cap: msgs, + ), + ): + return await execute_tool_loop( + prompt="hi", + max_tokens=64, + messages=[{"role": "user", "content": "q"}], + tools=[ + { + "name": "search_memory", + "description": "", + "input_schema": {"type": "object"}, + }, + { + "name": "search_messages", + "description": "", + "input_schema": {"type": "object"}, + }, + ], + tool_choice="auto", + tool_executor=tracker.execute, + max_tool_iterations=5, + response_model=None, + json_mode=False, + temperature=None, + stop_seqs=None, + verbosity=None, + enable_retry=False, + retry_attempts=1, + max_input_tokens=1000, + get_attempt_plan=_make_plan, + before_retry_callback=lambda _r: None, + stream_final=False, + telemetry=None, + ) + + +@pytest.mark.asyncio +async def test_tool_calls_in_one_iteration_run_concurrently(): + """Two tools requested in one turn overlap rather than queueing.""" + tracker = _Tracker() + + started = asyncio.get_running_loop().time() + await _run(tracker) + elapsed = asyncio.get_running_loop().time() - started + + assert tracker.peak == 2, f"tools ran sequentially (peak in-flight {tracker.peak})" + # Sequential would be >= 2 * delay; concurrent stays near one delay. + assert elapsed < TOOL_DELAY_SECONDS * 1.8, ( + f"elapsed {elapsed:.3f}s suggests serial execution" + ) + + +@pytest.mark.asyncio +async def test_results_keep_request_order(): + """Concurrency must not reorder results relative to the calls.""" + tracker = _Tracker() + + result = await _run(tracker) + + names = [call["tool_name"] for call in result.tool_calls_made] + assert names == ["search_memory", "search_messages"] + + +@pytest.mark.asyncio +async def test_fan_out_is_capped(): + """A turn asking for many tools does not stampede the database. + + Production has produced 18 tool calls in a single iteration; firing all of + them at once would mean that many concurrent embedding + pgvector queries + on one instance. + """ + tracker = _Tracker() + + await _run(tracker, n_calls=18) + + assert tracker.peak <= MAX_CONCURRENT_TOOL_CALLS, ( + f"fan-out reached {tracker.peak}, above the {MAX_CONCURRENT_TOOL_CALLS} cap" + ) + assert tracker.peak > 1, "cap must not serialise execution entirely" From 083ef01aab28a6032bec74277219a23e6c319b3d Mon Sep 17 00:00:00 2001 From: Tomas Sereikis Date: Sat, 29 Aug 2026 13:38:03 +0300 Subject: [PATCH 2/2] fix(llm): keep state-changing tool calls in the order the model asked for MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review on #1100 caught that the previous commit parallelised every tool call in a turn, mutations included. `ctx.db_lock` makes the mutating handlers take turns but does not decide whose turn comes first, so two writes from one assistant turn could apply in an order the model did not request — `tool_results` stayed ordered while the database did not. Reads may still overlap freely, because they change nothing. So the loop now splits the turn: anything named in `mutating_tools` runs one at a time in request order, everything else runs concurrently alongside it. `mutating_tools=None` is the default and means "assume every tool mutates", which is the fully sequential behaviour that existed before this branch. A caller opts in by naming its mutating set, so no existing caller changes behaviour by being left alone. `MUTATING_TOOL_NAMES` lives in agent_tools next to the dispatch table and lists exactly the five names routing to a handler that takes `ctx.db_lock`: create_observations, create_observations_deductive, create_observations_inductive, update_peer_card, delete_observations. `extract_preferences` and `finish_consolidation` are deliberately absent — both only return text telling the model what to call next. Dialectic passes that set rather than a "these are all reads" boolean. Its loadout is reads today, but DIALECTIC_TOOLS already carries a commented-out create_observations_deductive, so a flag would have gone quietly wrong the day someone uncommented it. The parameter is threaded through honcho_llm_call instead of imported inside src/llm, which keeps the llm layer free of tool semantics and avoids the existing agent_tools <- dreamer.specialists cycle. Two tests added: mutating calls are not reordered even when a later call finishes first, and the default schedule is fully sequential. --- src/dialectic/core.py | 11 ++++++ src/llm/api.py | 5 +++ src/llm/tool_loop.py | 39 +++++++++++++++++----- src/utils/agent_tools.py | 17 ++++++++++ tests/llm/test_tool_loop_parallel.py | 50 ++++++++++++++++++++++++++-- 5 files changed, 111 insertions(+), 11 deletions(-) diff --git a/src/dialectic/core.py b/src/dialectic/core.py index 57964c87..6842fbe1 100644 --- a/src/dialectic/core.py +++ b/src/dialectic/core.py @@ -35,6 +35,7 @@ from src.telemetry.prometheus.metrics import DialecticComponents, TokenTypes from src.utils.agent_tools import ( DIALECTIC_TOOLS, DIALECTIC_TOOLS_MINIMAL, + MUTATING_TOOL_NAMES, create_tool_executor, search_memory, ) @@ -508,6 +509,11 @@ class DialecticAgent: tool_choice=level_settings.TOOL_CHOICE, tool_executor=tool_executor, max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, + # Dialectic's loadout is reads today, but say so explicitly rather + # than relying on that: the loop keeps any state-changing call + # in the order the model asked for, so adding a write tool here + # later stays correct without anyone remembering this. + mutating_tools=MUTATING_TOOL_NAMES, messages=self.messages, max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, trace_name="dialectic_chat", @@ -584,6 +590,11 @@ class DialecticAgent: tool_choice=level_settings.TOOL_CHOICE, tool_executor=tool_executor, max_tool_iterations=level_settings.MAX_TOOL_ITERATIONS, + # Dialectic's loadout is reads today, but say so explicitly rather + # than relying on that: the loop keeps any state-changing call + # in the order the model asked for, so adding a write tool here + # later stays correct without anyone remembering this. + mutating_tools=MUTATING_TOOL_NAMES, messages=self.messages, max_input_tokens=settings.DIALECTIC.MAX_INPUT_TOKENS, trace_name="dialectic_chat", diff --git a/src/llm/api.py b/src/llm/api.py index 13c3a5e5..ca9795d6 100644 --- a/src/llm/api.py +++ b/src/llm/api.py @@ -68,6 +68,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -97,6 +98,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -126,6 +128,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -154,6 +157,7 @@ async def honcho_llm_call( tool_choice: str | dict[str, Any] | None = None, tool_executor: Callable[[str, dict[str, Any]], Any] | None = None, max_tool_iterations: int = 10, + mutating_tools: frozenset[str] | None = None, messages: list[dict[str, Any]] | None = None, max_input_tokens: int | None = None, trace_name: str | None = None, @@ -450,6 +454,7 @@ async def honcho_llm_call( tool_choice=tool_choice, tool_executor=tool_executor, max_tool_iterations=max_tool_iterations, + mutating_tools=mutating_tools, response_model=response_model, json_mode=json_mode, temperature=temperature, diff --git a/src/llm/tool_loop.py b/src/llm/tool_loop.py index 7a30b960..8a70291d 100644 --- a/src/llm/tool_loop.py +++ b/src/llm/tool_loop.py @@ -353,6 +353,7 @@ async def execute_tool_loop( tool_choice: str | dict[str, Any] | None, tool_executor: Callable[[str, dict[str, Any]], Any], max_tool_iterations: int, + mutating_tools: frozenset[str] | None = None, response_model: type[BaseModel] | None, json_mode: bool, temperature: float | None, @@ -649,17 +650,37 @@ async def execute_tool_loop( None, ) - # gather preserves argument order, so tool_results and - # all_tool_calls stay in the order the model asked for them. - outcomes = await asyncio.gather( - *( - asyncio.create_task(run_tool_call(seq, tool_call)) - for seq, tool_call in enumerate(response.tool_calls_made) - ) - ) + # Reads may overlap freely; state-changing calls may not, because + # `ctx.db_lock` decides that they take turns but not which turn each + # one gets. Anything named in `mutating_tools` therefore runs in the + # order the model asked for, while the rest run concurrently + # alongside it. `mutating_tools=None` means "assume every tool + # mutates", so a caller that has not opted in keeps the fully + # sequential behaviour it had before. + calls = list(enumerate(response.tool_calls_made)) + + def _must_stay_ordered(tool_call: dict[str, Any]) -> bool: + return mutating_tools is None or tool_call["name"] in mutating_tools + + # create_task even for the ordered ones: each still needs its own + # context copy for the telemetry ContextVars, and awaiting them one + # at a time is what preserves the order. + concurrent = { + seq: asyncio.create_task(run_tool_call(seq, tool_call)) + for seq, tool_call in calls + if not _must_stay_ordered(tool_call) + } + + outcomes: dict[int, tuple[dict[str, Any], dict[str, Any] | None]] = {} + for seq, tool_call in calls: + if _must_stay_ordered(tool_call): + outcomes[seq] = await run_tool_call(seq, tool_call) + for seq, task in concurrent.items(): + outcomes[seq] = await task tool_results: list[dict[str, Any]] = [] - for result_entry, call_entry in outcomes: + for seq, _ in calls: + result_entry, call_entry = outcomes[seq] tool_results.append(result_entry) if call_entry is not None: all_tool_calls.append(call_entry) diff --git a/src/utils/agent_tools.py b/src/utils/agent_tools.py index b753c462..37d90b8f 100644 --- a/src/utils/agent_tools.py +++ b/src/utils/agent_tools.py @@ -2604,6 +2604,23 @@ _TOOL_HANDLERS: dict[str, Callable[[ToolContext, dict[str, Any]], Any]] = { "get_reasoning_chain": _handle_get_reasoning_chain, } +# Tools that change state, and so must not be reordered relative to one another +# within a single assistant turn. Each of these routes to a handler that takes +# `ctx.db_lock` (`_handle_create_observations_impl`, `_handle_update_peer_card`, +# `_handle_delete_observations`); the lock stops them interleaving, but it does +# not decide who gets it first, so the tool loop keeps them in the order the +# model asked for. `extract_preferences` and `finish_consolidation` are not +# here: both only return text telling the model what to call next. +MUTATING_TOOL_NAMES: frozenset[str] = frozenset( + { + "create_observations", + "create_observations_deductive", + "create_observations_inductive", + "update_peer_card", + "delete_observations", + } +) + async def create_tool_executor( workspace_name: str, diff --git a/tests/llm/test_tool_loop_parallel.py b/tests/llm/test_tool_loop_parallel.py index ac1408d9..f49e66f2 100644 --- a/tests/llm/test_tool_loop_parallel.py +++ b/tests/llm/test_tool_loop_parallel.py @@ -75,8 +75,13 @@ class _Tracker: return "ok" -async def _run(tracker: _Tracker, n_calls: int = 2) -> HonchoLLMCallResponse[Any]: - names = ["search_memory", "search_messages"] +async def _run( + tracker: _Tracker, + n_calls: int = 2, + mutating_tools: frozenset[str] | None = frozenset(), + names: list[str] | None = None, +) -> HonchoLLMCallResponse[Any]: + names = names or ["search_memory", "search_messages"] calls = [ {"name": names[i % len(names)], "input": {"query": str(i)}, "id": f"t{i}"} for i in range(n_calls) @@ -113,6 +118,7 @@ async def _run(tracker: _Tracker, n_calls: int = 2) -> HonchoLLMCallResponse[Any tool_choice="auto", tool_executor=tracker.execute, max_tool_iterations=5, + mutating_tools=mutating_tools, response_model=None, json_mode=False, temperature=None, @@ -171,3 +177,43 @@ async def test_fan_out_is_capped(): f"fan-out reached {tracker.peak}, above the {MAX_CONCURRENT_TOOL_CALLS} cap" ) assert tracker.peak > 1, "cap must not serialise execution entirely" + + +@pytest.mark.asyncio +async def test_mutating_tools_are_not_reordered(): + """State-changing calls keep the order the model asked for. + + `ctx.db_lock` makes the mutating handlers take turns but does not decide + whose turn comes first, so running them concurrently could apply a turn's + writes in an order the model did not request. + """ + order: list[str] = [] + + class _Recorder(_Tracker): + async def execute(self, name: str, _input: dict[str, Any]) -> str: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + # Later calls sleep less, so anything concurrent finishes reversed. + await asyncio.sleep(TOOL_DELAY_SECONDS / (len(order) + 1)) + order.append(name) + self.in_flight -= 1 + return "ok" + + await _run( + _Recorder(), + n_calls=3, + mutating_tools=frozenset({"update_peer_card"}), + names=["update_peer_card"], + ) + + assert order == ["update_peer_card"] * 3 + + +@pytest.mark.asyncio +async def test_default_is_fully_sequential(): + """A caller that does not opt in keeps the old one-at-a-time behaviour.""" + tracker = _Tracker() + + await _run(tracker, n_calls=4, mutating_tools=None) + + assert tracker.peak == 1