fix(summarization): count replayed vLLM reasoning (#8077)

* fix(summarization): count replayed vLLM reasoning

* fix(summarization): unwrap nested runnable bindings

* test(summarization): cover replayed reasoning in content blocks

---------

Co-authored-by: lovesky514514-cmd <273698517+lovesky514514-cmd@users.noreply.github.com>
Co-authored-by: Totoro <msy626836554@gmail.com>

CHANGELOG.md

@@ -47,6 +47,10 @@ This section accumulates work toward the **2.3.0** milestone
 
 ### Fixed
 
+- **summarization:** vLLM compaction counts reasoning tokens replayed to the
+  model when checking the trigger and estimating how much history to keep.
+  ([#8077](https://github.com/bytedance/deer-flow/pull/8077))
+
 - **gateway:** Posting feedback twice for the same run now returns a 409
   instead of a 500. `POST /api/threads/{id}/runs/{run_id}/feedback` allows one
   feedback row per user and run; a retry, or feedback on another message of the

README.md

@@ -2737,7 +2737,8 @@ A pass is limited to 64 chunks and leaves history unchanged on failure or limit
 exhaustion. A limit breach starts a five-minute per-process cooldown for the
 same thread/history window and policy, including rebuilt middleware. A larger
 input budget or smaller source window permits an immediate retry. Model fallback
-and SDK retries can add calls. See the
+and SDK retries can add calls. For vLLM run models, history and retention
+estimates include assistant reasoning replayed to the provider. See the
 [summarization guide](backend/docs/summarization.md).
 
 With `task_continuity.enabled`, `history_search` searches the current task's active

backend/docs/summarization.md

@@ -282,6 +282,11 @@ The default LangChain prompt instructs the model to:
 - Uses approximate token counting based on character count
 - For Anthropic models: ~3.3 characters per token
 - For other models: Uses LangChain's default estimation
+- For `VllmChatModel` using Chat Completions: History and token-based retention
+  estimates include the assistant reasoning replayed to vLLM. A non-null `reasoning` field takes
+  precedence over `reasoning_content`, including when it is empty. This applies
+  to both automatic and manual compaction. The estimate is approximate and
+  leaves the summary prompt and its trimming budget unchanged.
 - Can be customized with a custom `token_counter` function
 
 ### Message Preservation

backend/packages/harness/deerflow/agents/middlewares/AGENTS.md

@@ -25,6 +25,8 @@ PII before splitting, budgets before escaping; hooks/history/removal on success.
 Chunk/carry/cooldown invariants and tests:
 [contract](../../../../../docs/summarization.md).
 
+vLLM compaction counts replayed reasoning. Keep it out of summary input.
+
 Todo compaction: [contract](../../../../../docs/summarization.md#todo-reminders).
 
 Delegation verdicts are untrusted: revalidate persisted values, ignore malformed

backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py

@@ -10,6 +10,7 @@ import threading
 from collections import OrderedDict
 from collections.abc import Mapping
 from dataclasses import dataclass
+from functools import partial
 from time import monotonic
 from typing import Any, Protocol, override, runtime_checkable
 
@@ -17,6 +18,7 @@ from deerflow_extension_api import CompactionEvent, ExtensionData, canonical_has
 from langchain.agents import AgentState
 from langchain.agents.middleware import SummarizationMiddleware
 from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, SystemMessage, get_buffer_string
+from langchain_core.runnables import RunnableBinding
 from langgraph.config import get_config
 from langgraph.constants import TAG_NOSTREAM
 from langgraph.graph.message import REMOVE_ALL_MESSAGES
@@ -32,6 +34,7 @@ from deerflow.config.task_continuity_config import TaskContinuityConfig
 from deerflow.extensions.notify import notify_context_compacted
 from deerflow.models import create_chat_model
 from deerflow.models.reasoning import resolve_reasoning_contract
+from deerflow.models.vllm_provider import VllmChatModel, count_tokens_with_replayed_reasoning
 from deerflow.runtime.events.catalog import MIDDLEWARE_SUMMARIZE_TAG
 from deerflow.utils.assembly_io import run_assembly
 
@@ -177,6 +180,15 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
         **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
+        run_model = self.model
+        seen_bindings: set[int] = set()
+        while isinstance(run_model, RunnableBinding) and id(run_model) not in seen_bindings:
+            seen_bindings.add(id(run_model))
+            run_model = run_model.bound
+        if "token_counter" not in kwargs and isinstance(run_model, VllmChatModel) and run_model.replays_reasoning():
+            self.token_counter = partial(count_tokens_with_replayed_reasoning, token_counter=self.token_counter)
+            self._partial_token_counter = partial(count_tokens_with_replayed_reasoning, token_counter=self._partial_token_counter)
+            # Preserve LangChain's unscaled suffix estimate.
         self._task_continuity_config = task_continuity_config if task_continuity_config is not None and task_continuity_config.enabled is True else None
         self._before_summarization_hooks = before_summarization or []
         # Model-ownership state. The model that actually executes the run is selected
@@ -1350,8 +1362,6 @@ def create_summarization_middleware(
 
     hooks: list[BeforeSummarizationHook] = []
     if resolved_app_config.memory.enabled and not skip_memory_flush:
-        from functools import partial
-
         from deerflow.agents.memory.summarization_hook import memory_flush_hook
 
         hooks.append(partial(memory_flush_hook, pii_redaction_config=getattr(resolved_app_config, "pii_redaction", None)))

backend/packages/harness/deerflow/models/vllm_provider.py

@@ -23,7 +23,7 @@ import json
 import threading
 import time
 from collections import OrderedDict
-from collections.abc import Mapping
+from collections.abc import Callable, Iterable, Mapping
 from typing import Any, cast
 
 import openai
@@ -31,12 +31,15 @@ from langchain_core.language_models import LanguageModelInput
 from langchain_core.messages import (
     AIMessage,
     AIMessageChunk,
+    BaseMessage,
     BaseMessageChunk,
     ChatMessageChunk,
     FunctionMessageChunk,
     HumanMessageChunk,
+    MessageLikeRepresentation,
     SystemMessageChunk,
     ToolMessageChunk,
+    convert_to_messages,
 )
 from langchain_core.messages.ai import UsageMetadata, subtract_usage
 from langchain_core.messages.tool import tool_call_chunk
@@ -112,6 +115,24 @@ def _pick_reasoning(fields: Mapping[str, Any]) -> Any:
     return fields.get("reasoning_content") if reasoning is None else reasoning
 
 
+def count_tokens_with_replayed_reasoning(
+    messages: Iterable[MessageLikeRepresentation],
+    *,
+    token_counter: Callable[[Iterable[MessageLikeRepresentation]], int],
+) -> int:
+    counted: list[BaseMessage] = []
+    for message in convert_to_messages(messages):
+        if isinstance(message, AIMessage):
+            reasoning = _pick_reasoning(message.additional_kwargs)
+            text = _reasoning_to_text(reasoning) if reasoning is not None else ""
+            if text:
+                content = message.content + text if isinstance(message.content, str) else [*message.content, text]
+                message = message.model_copy(update={"content": content})
+                # Count replayed reasoning on a copy.
+        counted.append(message)
+    return token_counter(counted)
+
+
 def _convert_delta_to_message_chunk_with_reasoning(_dict: Mapping[str, Any], default_class: type[BaseMessageChunk]) -> BaseMessageChunk:
     """Convert a streaming delta to a LangChain message chunk while preserving reasoning."""
     id_ = _dict.get("id")
@@ -204,6 +225,9 @@ class VllmChatModel(GuardedChatOpenAI):
     def _llm_type(self) -> str:
         return "vllm-openai-compatible"
 
+    def replays_reasoning(self) -> bool:
+        return not self._use_responses_api(self._default_params)  # Chat Completions only.
+
     def _usage_delta(self, completion_id: str, usage: UsageMetadata, *, terminal: bool) -> UsageMetadata:
         """Convert a completion's cumulative usage snapshot into a delta."""
         with self._cumulative_usage_lock:

backend/tests/test_summarization_middleware.py

@@ -10,7 +10,9 @@ import pytest
 from langchain.agents import create_agent
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage
+from langchain_core.messages.utils import count_tokens_approximately
 from langchain_core.outputs import ChatGeneration, ChatResult
+from langchain_core.runnables import RunnableBinding
 from langgraph.constants import TAG_NOSTREAM
 from pydantic import ValidationError
 
@@ -23,6 +25,7 @@ from deerflow.config.memory_config import MemoryConfig
 from deerflow.config.model_config import ModelConfig, ReasoningCapabilities
 from deerflow.config.sandbox_config import SandboxConfig
 from deerflow.config.summarization_config import ContextSize, SummarizationConfig
+from deerflow.models.vllm_provider import VllmChatModel
 
 
 def _messages() -> list:
@@ -1207,6 +1210,107 @@ def _factory_app_config(model_names, *, summary_model_name=None, summarization_k
     )
 
 
+@pytest.fixture
+def vllm_compaction(monkeypatch):
+    models = {
+        "vllm": VllmChatModel(model="test-model", api_key="dummy", base_url="http://localhost:8000/v1"),
+        "summary": _StaticChatModel(text="compressed summary"),
+    }
+    monkeypatch.setattr(
+        "deerflow.agents.middlewares.summarization_middleware.create_chat_model",
+        lambda name, **_: models[name],
+    )
+    config = _factory_app_config(
+        ("vllm", "summary"),
+        summary_model_name="summary",
+        summarization_kwargs={"trigger": ContextSize(type="tokens", value=150)},
+    )
+    middleware = create_summarization_middleware(app_config=config, run_model_name="vllm", keep=("tokens", 100))
+    assert middleware is not None
+    return middleware
+
+
+_REPLAYED_REASONING = "Review the tool results. " * 100
+
+
+@pytest.mark.parametrize("nested", [False, True], ids=["single-binding", "double-binding"])
+def test_vllm_compaction_counts_reasoning_through_bindings(vllm_compaction, nested):
+    model = vllm_compaction.model
+    if nested:
+        model = RunnableBinding(bound=model)
+    middleware = DeerFlowSummarizationMiddleware(model=model)
+    message = AIMessage(content="Checked.", additional_kwargs={"reasoning": _REPLAYED_REASONING})
+    counted = AIMessage(content="Checked." + _REPLAYED_REASONING)
+
+    assert middleware.token_counter([message]) == count_tokens_approximately([counted])
+
+
+@pytest.mark.parametrize(
+    ("fields", "replayed"),
+    [
+        pytest.param({}, None, id="ordinary-message"),
+        pytest.param({"reasoning": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="reasoning-only"),
+        pytest.param({"reasoning_content": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="legacy-only"),
+        pytest.param({"reasoning": None, "reasoning_content": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="null-canonical"),
+        pytest.param({"reasoning": _REPLAYED_REASONING, "reasoning_content": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="duplicate-alias"),
+        pytest.param({"reasoning": _REPLAYED_REASONING, "reasoning_content": "different text" * 500}, _REPLAYED_REASONING, id="canonical-wins"),
+        pytest.param({"reasoning": "", "reasoning_content": _REPLAYED_REASONING}, "", id="empty-canonical"),
+    ],
+)
+def test_vllm_compaction_counts_replayed_reasoning_once(vllm_compaction, fields, replayed):
+    message = AIMessage(
+        content="Checked.",
+        additional_kwargs=fields,
+        tool_calls=[{"name": "check", "args": {}, "id": "call-1"}],
+    )
+    payload = vllm_compaction.model._get_request_payload([message])
+    assert payload["messages"][0].get("reasoning") == replayed
+    assert "reasoning_content" not in payload["messages"][0]
+
+    counted = message.model_copy(update={"content": message.content + (replayed or "")})
+    assert vllm_compaction.token_counter([message]) == count_tokens_approximately([counted])
+    assert message.content == "Checked."
+    assert message.additional_kwargs == fields
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("manual", [False, True], ids=["auto", "manual"])
+async def test_vllm_compaction_uses_replayed_reasoning_for_trigger_and_retention(vllm_compaction, manual):
+    messages = _messages()
+    plain_state = {"messages": messages}
+    assert await vllm_compaction.abefore_model(plain_state, _runtime()) is None
+    messages[1] = messages[1].model_copy(update={"additional_kwargs": {"reasoning": _REPLAYED_REASONING}})
+    state = {"messages": messages}
+
+    if manual:
+        result = await vllm_compaction.acompact_state(state, _runtime(), force=True)
+        assert result is not None
+        assert result.total_tokens >= 150
+        assert list(result.preserved_messages) == messages[2:]
+        assert result.summary_text == "compressed summary"
+    else:
+        update = await vllm_compaction.abefore_model(state, _runtime())
+        assert update is not None
+        assert update["messages"][1:] == messages[2:]
+        assert update["summary_text"] == "compressed summary"
+
+
+def test_non_vllm_compaction_keeps_existing_token_estimate():
+    middleware = DeerFlowSummarizationMiddleware(model=_StaticChatModel())
+    message = AIMessage(content="Checked.", additional_kwargs={"reasoning": _REPLAYED_REASONING})
+    assert middleware.token_counter([message]) == count_tokens_approximately([message])
+
+
+def test_vllm_compaction_ignores_reasoning_when_using_responses_api(vllm_compaction):
+    model = vllm_compaction.model.bound.model_copy(update={"use_responses_api": True})
+    middleware = DeerFlowSummarizationMiddleware(model=model)
+    message = AIMessage(content="Checked.", additional_kwargs={"reasoning": _REPLAYED_REASONING})
+    payload = model._get_request_payload([message])
+    assert "input" in payload
+    assert _REPLAYED_REASONING not in repr(payload["input"])
+    assert middleware.token_counter([message]) == count_tokens_approximately([message])
+
+
 @pytest.mark.parametrize("trim_limit", [None, 80, 4000])
 def test_factory_preserves_explicit_summary_input_limit(monkeypatch, trim_limit):
     monkeypatch.setattr("deerflow.agents.middlewares.summarization_middleware.create_chat_model", _tracking_create_chat_model([]))

backend/tests/test_vllm_provider.py

@@ -6,8 +6,9 @@ import json
 import httpx
 import pytest
 from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage
+from langchain_core.messages.utils import count_tokens_approximately
 
-from deerflow.models.vllm_provider import VllmChatModel
+from deerflow.models.vllm_provider import VllmChatModel, count_tokens_with_replayed_reasoning
 
 
 def _make_model(*, cumulative_stream_usage: bool = False) -> VllmChatModel:
@@ -62,6 +63,23 @@ def _assert_usage(message, *, input_tokens: int, output_tokens: int, total_token
     assert message.usage_metadata["total_tokens"] == total_tokens
 
 
+def test_count_tokens_with_replayed_reasoning_preserves_content_blocks():
+    blocks = [
+        {"type": "text", "text": "Checked the source."},
+        {"type": "text", "text": "The output matches."},
+    ]
+    reasoning = "Compare each output field with its source."
+    message = AIMessage(content=blocks, additional_kwargs={"reasoning": reasoning})
+    original = message.model_copy(deep=True)
+
+    count = count_tokens_with_replayed_reasoning([message], token_counter=count_tokens_approximately)
+
+    assert count == count_tokens_approximately([AIMessage(content=[*blocks, reasoning])])
+    assert count > count_tokens_approximately([original])
+    assert count > count_tokens_approximately([AIMessage(content=[reasoning])])
+    assert message == original
+
+
 def test_vllm_provider_restores_reasoning_in_request_payload():
     model = _make_model()
     payload = model._get_request_payload(