Merge 083ef01aab into ced1514200

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",

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,

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
@@ -345,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,
@@ -391,6 +400,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 +594,97 @@ 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,
                     )
 
+            # 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 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)
+
             append_tool_results(current_provider, tool_results, conversation_messages)
         finally:
             if step is not None:

src/utils/agent_tools.py

@@ -2619,6 +2619,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,

tests/llm/test_tool_loop_parallel.py

@@ -0,0 +1,219 @@
+# 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,
+    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)
+    ]
+    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,
+            mutating_tools=mutating_tools,
+            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"
+
+
+@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