perf(sdk): count tokens once per model call in summarization middleware (#3877)

`SummarizationMiddleware.wrap_model_call` / `awrap_model_call` counted
tokens over the full effective history twice on every model call — once
inside `_truncate_args` and again for the should-summarize check — even
on the common path where nothing is truncated or summarized. The count
is expensive because counting with `tools=` converts every tool schema
on every invocation; in a production deep-agent trace with 23 tools this
middleware added roughly 280ms to every model call, and the duplicate
count is half of that.

The count now runs once per call and is passed into `_truncate_args` via
a new keyword-only `total_tokens` parameter that defaults to computing
itself, so existing direct callers are unaffected. A recount happens
only when truncation actually modified messages (which changes the
total).

Related: langchain-ai/langchain#38073 addresses the underlying
per-conversion tool-schema cost.

---------

Co-authored-by: Mason Daugherty <github@mdrxy.com>
Co-authored-by: Mason Daugherty <mason@langchain.dev>

libs/deepagents/deepagents/middleware/summarization.py

@@ -711,28 +711,47 @@ A condensed summary follows:
             }
         return tool_call
 
-    def _truncate_args(
+    def _count_tokens(
         self,
         messages: list[AnyMessage],
         system_message: SystemMessage | None,
         tools: list[BaseTool | dict[str, Any]] | None,
+    ) -> int:
+        """Count tokens for messages plus optional system message and tools.
+
+        Args:
+            messages: Messages to count.
+            system_message: Optional system message prepended before counting.
+            tools: Optional tools whose schemas contribute to the count.
+
+        Returns:
+            Total token count. Falls back to counting without `tools` when the
+                configured `token_counter` does not accept a `tools` keyword.
+        """
+        counted_messages = [system_message, *messages] if system_message is not None else messages
+        try:
+            return self.token_counter(counted_messages, tools=tools)  # ty: ignore[unknown-argument]
+        except TypeError:
+            return self.token_counter(counted_messages)
+
+    def _truncate_args(
+        self,
+        messages: list[AnyMessage],
+        total_tokens: int,
     ) -> tuple[list[AnyMessage], bool]:
         """Truncate large tool call arguments in old messages.
 
         Args:
             messages: Messages to potentially truncate.
-            system_message: Optional system message for token counting.
-            tools: Optional tools for token counting.
+            total_tokens: Precomputed token count for `messages` (plus system
+                message and tools). Counting tools is expensive (schema
+                conversion per tool), so the caller counts once and shares the
+                result across the truncation and summarization checks.
 
         Returns:
             Tuple of (truncated_messages, modified). If modified is False,
             truncated_messages is the same as input messages.
         """
-        counted_messages = [system_message, *messages] if system_message is not None else messages
-        try:
-            total_tokens = self.token_counter(counted_messages, tools=tools)  # ty: ignore[unknown-argument]
-        except TypeError:
-            total_tokens = self.token_counter(counted_messages)
         if not self._should_truncate_args(messages, total_tokens):
             return messages, False
 
@@ -960,19 +979,19 @@ A condensed summary follows:
         # Get effective messages based on previous summarization events
         effective_messages = self._get_effective_messages(request)
 
+        # Count once; tool-schema conversion makes each count expensive, so the
+        # count is shared between the truncation check and the summarize check.
+        total_tokens = self._count_tokens(effective_messages, request.system_message, request.tools)
+
         # Step 1: Truncate args if configured
-        truncated_messages, _ = self._truncate_args(
+        truncated_messages, truncate_modified = self._truncate_args(
             effective_messages,
-            request.system_message,
-            request.tools,
+            total_tokens,
         )
 
         # Step 2: Check if summarization should happen
-        counted_messages = [request.system_message, *truncated_messages] if request.system_message is not None else truncated_messages
-        try:
-            total_tokens = self.token_counter(counted_messages, tools=request.tools)  # ty: ignore[unknown-argument]
-        except TypeError:
-            total_tokens = self.token_counter(counted_messages)
+        if truncate_modified:
+            total_tokens = self._count_tokens(truncated_messages, request.system_message, request.tools)
         should_summarize = self._should_summarize(truncated_messages, total_tokens)
 
         # If no summarization needed, return with truncated messages
@@ -1081,19 +1100,19 @@ A condensed summary follows:
         # Get effective messages based on previous summarization events
         effective_messages = self._get_effective_messages(request)
 
+        # Count once; tool-schema conversion makes each count expensive, so the
+        # count is shared between the truncation check and the summarize check.
+        total_tokens = self._count_tokens(effective_messages, request.system_message, request.tools)
+
         # Step 1: Truncate args if configured
-        truncated_messages, _ = self._truncate_args(
+        truncated_messages, truncate_modified = self._truncate_args(
             effective_messages,
-            request.system_message,
-            request.tools,
+            total_tokens,
         )
 
         # Step 2: Check if summarization should happen
-        counted_messages = [request.system_message, *truncated_messages] if request.system_message is not None else truncated_messages
-        try:
-            total_tokens = self.token_counter(counted_messages, tools=request.tools)  # ty: ignore[unknown-argument]
-        except TypeError:
-            total_tokens = self.token_counter(counted_messages)
+        if truncate_modified:
+            total_tokens = self._count_tokens(truncated_messages, request.system_message, request.tools)
         should_summarize = self._should_summarize(truncated_messages, total_tokens)
 
         # If no summarization needed, return with truncated messages

libs/deepagents/tests/benchmarks/test_benchmark_summarization_middleware.py

@@ -0,0 +1,169 @@
+"""Wall-time benchmarks for `SummarizationMiddleware` per-model-call overhead.
+
+Run locally:  `make benchmark`
+Run with CodSpeed:  `uv run --group test pytest ./tests -m benchmark --codspeed`
+
+These tests measure the wall time of `wrap_model_call` / `awrap_model_call`
+on the *common* path -- the one taken on every model invocation where nothing
+is truncated and nothing is summarized. That path runs the token counter, and
+counting with `tools=` is expensive because every tool schema is converted on
+each invocation.
+
+The benchmarks pin a deterministic token counter that reproduces that cost
+(it converts each tool to an OpenAI schema per call) and scale the tool count,
+so the per-call counting overhead dominates the measurement. They exist to
+guard the optimization in PR #3877, which collapsed two full token counts per
+model call down to one on this path; with a duplicate count the numbers here
+roughly double.
+
+Regression detection is handled by CodSpeed in CI. Local runs produce
+pytest-benchmark tables (min/max/mean/stddev) for human inspection.
+"""
+
+import asyncio
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any, cast
+
+import pytest
+from langchain.agents.middleware.types import ModelRequest
+from langchain_core.messages import AIMessage, BaseMessage
+from langchain_core.tools import BaseTool, tool
+from langchain_core.utils.function_calling import convert_to_openai_tool
+from pytest_benchmark.fixture import BenchmarkFixture
+
+from deepagents.middleware.summarization import SummarizationMiddleware
+from tests.unit_tests.middleware.test_summarization_middleware import (
+    MockBackend,
+    make_conversation_messages,
+    make_mock_model,
+    make_mock_runtime,
+)
+
+if TYPE_CHECKING:
+    from langchain.agents.middleware.types import AgentState, ModelResponse
+
+# A high threshold both triggers never reach: the conversation stays small, so
+# neither truncation nor summarization fires and the no-op path is exercised.
+_NEVER = 1_000_000_000
+
+# Tool counts to scale the per-call schema-conversion cost. 23 matches the
+# production deep-agent trace cited in PR #3877.
+_TOOL_COUNTS = [1, 5, 23]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_tool(idx: int) -> BaseTool:
+    """Create a named tool with a non-trivial schema for conversion cost."""
+
+    @tool(description=f"Tool number {idx} used for benchmarking schema conversion")
+    def dynamic_tool(query: str, count: int = 1, *, verbose: bool = False) -> str:
+        """Echo the query a number of times."""
+        return f"tool_{idx}({query}, {count}, {verbose})"
+
+    dynamic_tool.name = f"tool_{idx}"
+    return dynamic_tool
+
+
+def _converting_token_counter(messages: list[BaseMessage], **kwargs: Any) -> int:
+    """Token counter that pays the real per-call tool-conversion cost.
+
+    Production token counters convert every tool schema on each call (see
+    langchain-ai/langchain#38073). This reproduces that cost deterministically
+    so the benchmark measures what the optimization actually saves, rather than
+    a trivial counter whose double-invocation is invisible.
+    """
+    tools: list[BaseTool | dict[str, Any]] | None = kwargs.get("tools")
+    total = 0
+    if tools is not None:
+        for t in tools:
+            schema = convert_to_openai_tool(t)
+            total += len(str(schema)) // 4
+    for message in messages:
+        total += len(str(message.content)) // 4
+    return total
+
+
+def _make_middleware() -> SummarizationMiddleware:
+    """Build a middleware whose triggers never fire (common no-op path)."""
+    return SummarizationMiddleware(
+        model=make_mock_model(),
+        backend=MockBackend(),
+        trigger=("tokens", _NEVER),
+        token_counter=_converting_token_counter,
+        truncate_args_settings={
+            "trigger": ("tokens", _NEVER),
+            "keep": ("messages", 2),
+        },
+    )
+
+
+def _make_request(tools: list[BaseTool]) -> ModelRequest:
+    """Build a `ModelRequest` carrying the tools and a realistic conversation."""
+    state = cast("AgentState[Any]", {"messages": make_conversation_messages()})
+    return ModelRequest(
+        model=make_mock_model(),
+        messages=state["messages"],
+        system_message=None,
+        tools=tools,
+        runtime=make_mock_runtime(),
+        state=state,
+    )
+
+
+def _handler(_request: ModelRequest) -> "ModelResponse":
+    """No-op handler standing in for the wrapped model call."""
+    return AIMessage(content="ok")
+
+
+async def _ahandler(_request: ModelRequest) -> "ModelResponse":
+    """Async no-op handler standing in for the wrapped model call."""
+    return AIMessage(content="ok")
+
+
+@pytest.fixture
+def event_loop() -> Iterator[asyncio.AbstractEventLoop]:
+    """A dedicated event loop, closed on teardown to avoid leaking loops."""
+    loop = asyncio.new_event_loop()
+    try:
+        yield loop
+    finally:
+        loop.close()
+
+
+# ---------------------------------------------------------------------------
+# Benchmarks
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.benchmark
+class TestSummarizationWrapModelCall:
+    """Per-call overhead of the no-summarization, no-truncation path."""
+
+    @pytest.mark.parametrize("tool_count", _TOOL_COUNTS, ids=lambda n: f"{n}_tools")
+    def test_wrap_model_call(self, benchmark: BenchmarkFixture, tool_count: int) -> None:
+        """Sync path: token counting must run once, not once per check."""
+        middleware = _make_middleware()
+        tools = [_make_tool(i) for i in range(tool_count)]
+
+        @benchmark  # type: ignore[misc]
+        def _() -> None:
+            middleware.wrap_model_call(_make_request(tools), _handler)
+
+    @pytest.mark.parametrize("tool_count", _TOOL_COUNTS, ids=lambda n: f"{n}_tools")
+    def test_awrap_model_call(
+        self,
+        benchmark: BenchmarkFixture,
+        tool_count: int,
+        event_loop: asyncio.AbstractEventLoop,
+    ) -> None:
+        """Async path: same single-count guarantee as the sync path."""
+        middleware = _make_middleware()
+        tools = [_make_tool(i) for i in range(tool_count)]
+
+        @benchmark  # type: ignore[misc]
+        def _() -> None:
+            event_loop.run_until_complete(middleware.awrap_model_call(_make_request(tools), _ahandler))

libs/deepagents/tests/unit_tests/middleware/test_summarization_middleware.py

@@ -2563,3 +2563,122 @@ async def test_async_offload_and_summary_run_concurrently() -> None:
     # If sequential, elapsed >= 2 * delay (0.2s). If parallel, elapsed ~ delay.
     # Use 2.5x multiplier to allow for CI scheduling jitter.
     assert elapsed < 2.5 * delay, f"Expected parallel execution (<{2.5 * delay}s) but took {elapsed:.2f}s"
+
+
+class TestTokenCountingEfficiency:
+    """The per-call token count is expensive (tool schema conversion), so it must run once."""
+
+    def _make_counting_middleware(self) -> tuple[SummarizationMiddleware, dict[str, int]]:
+        calls = {"count": 0}
+
+        def counting_token_counter(_messages: list[BaseMessage], **_kwargs: Any) -> int:
+            calls["count"] += 1
+            return 10  # Far below every trigger: no truncation, no summarization.
+
+        middleware = SummarizationMiddleware(
+            model=make_mock_model(),
+            backend=MockBackend(),
+            trigger=("tokens", 1_000_000),
+            token_counter=counting_token_counter,
+            truncate_args_settings={
+                "trigger": ("tokens", 1_000_000),
+                "keep": ("messages", 2),
+            },
+        )
+        return middleware, calls
+
+    def test_token_counter_called_once_per_model_call(self) -> None:
+        middleware, calls = self._make_counting_middleware()
+        state = cast("AgentState[Any]", {"messages": make_conversation_messages()})
+
+        _, captured_request = call_wrap_model_call(middleware, state, make_mock_runtime())
+
+        assert captured_request is not None  # Handler ran; nothing was summarized.
+        assert calls["count"] == 1
+
+    async def test_token_counter_called_once_per_model_call_async(self) -> None:
+        middleware, calls = self._make_counting_middleware()
+        state = cast("AgentState[Any]", {"messages": make_conversation_messages()})
+
+        _, captured_request = await call_awrap_model_call(middleware, state, make_mock_runtime())
+
+        assert captured_request is not None  # Handler ran; nothing was summarized.
+        assert calls["count"] == 1
+
+    def _make_truncating_counting_middleware(
+        self,
+    ) -> tuple[SummarizationMiddleware, dict[str, int]]:
+        """Middleware whose truncation trigger fires but summarization never does.
+
+        Truncating tool-call args shrinks the message set, so the count must be
+        refreshed afterward before the summarization check; this middleware
+        exercises that recount branch.
+        """
+        calls = {"count": 0}
+
+        def counting_token_counter(_messages: list[BaseMessage], **_kwargs: Any) -> int:
+            calls["count"] += 1
+            return 10  # Far below the summarization trigger: nothing is summarized.
+
+        middleware = SummarizationMiddleware(
+            model=make_mock_model(),
+            backend=MockBackend(),
+            trigger=("messages", 100),  # High: no summarization.
+            token_counter=counting_token_counter,
+            truncate_args_settings={
+                "trigger": ("messages", 5),  # Low: truncation fires.
+                "keep": ("messages", 2),
+                "max_length": 100,
+            },
+        )
+        return middleware, calls
+
+    def _truncatable_state(self) -> "AgentState[Any]":
+        """A conversation whose old `write_file` arg is large enough to truncate."""
+        messages = [
+            AIMessage(
+                content="",
+                id="a1",
+                tool_calls=[
+                    {
+                        "id": "tc1",
+                        "name": "write_file",
+                        "args": {"file_path": "/test.txt", "content": "x" * 200},
+                    }
+                ],
+            ),
+            ToolMessage(content="File written", tool_call_id="tc1", id="t1"),
+            HumanMessage(content="Request 1", id="h1"),
+            AIMessage(content="Response 1", id="a2"),
+            HumanMessage(content="Request 2", id="h2"),
+            AIMessage(content="Response 2", id="a3"),
+        ]
+        return cast("AgentState[Any]", {"messages": messages})
+
+    def test_token_counter_recounts_when_truncation_modifies_messages(self) -> None:
+        middleware, calls = self._make_truncating_counting_middleware()
+
+        _, captured_request = call_wrap_model_call(middleware, self._truncatable_state(), make_mock_runtime())
+
+        assert captured_request is not None  # Handler ran; nothing was summarized.
+        # Truncation changed the message set, so the count is refreshed:
+        # once before truncation, once after.
+        assert calls["count"] == 2
+        # Confirm the modify path was genuinely taken (not a vacuous recount).
+        first_ai = captured_request.messages[0]
+        assert isinstance(first_ai, AIMessage)
+        assert first_ai.tool_calls[0]["args"]["content"] == "x" * 20 + "...(argument truncated)"
+
+    async def test_token_counter_recounts_when_truncation_modifies_messages_async(self) -> None:
+        middleware, calls = self._make_truncating_counting_middleware()
+
+        _, captured_request = await call_awrap_model_call(middleware, self._truncatable_state(), make_mock_runtime())
+
+        assert captured_request is not None  # Handler ran; nothing was summarized.
+        # Truncation changed the message set, so the count is refreshed:
+        # once before truncation, once after.
+        assert calls["count"] == 2
+        # Confirm the modify path was genuinely taken (not a vacuous recount).
+        first_ai = captured_request.messages[0]
+        assert isinstance(first_ai, AIMessage)
+        assert first_ai.tool_calls[0]["args"]["content"] == "x" * 20 + "...(argument truncated)"