diff --git a/agent_core/components/middleware/llm/compaction.py b/agent_core/components/middleware/llm/compaction.py index a4bd55e..73b8328 100644 --- a/agent_core/components/middleware/llm/compaction.py +++ b/agent_core/components/middleware/llm/compaction.py @@ -60,7 +60,7 @@ def _estimate_tokens_heuristic(messages: list[Message]) -> int: """ total = 0 for m in messages: - text = text_of(m.get("content")) + text = text_of(m.get("content")) + str(m.get("reasoning_content") or "") cjk_count = len(_CJK_RE.findall(text)) other_count = len(text) - cjk_count total += cjk_count + (other_count // 4) + _MSG_OVERHEAD @@ -84,7 +84,9 @@ def estimate_tokens(messages: list[Message]) -> int: for m in messages: total += len( encoder.encode( - text_of(m.get("content")), disallowed_special=() + text_of(m.get("content")) + + str(m.get("reasoning_content") or ""), + disallowed_special=(), ) ) + _MSG_OVERHEAD return total diff --git a/agent_core/components/middleware/llm/summarization.py b/agent_core/components/middleware/llm/summarization.py index c3641fa..717594c 100644 --- a/agent_core/components/middleware/llm/summarization.py +++ b/agent_core/components/middleware/llm/summarization.py @@ -70,8 +70,9 @@ def _estimate_tiktoken(self, encoder: Any, messages: list[Message]) -> int: try: total = 0 for m in messages: + text = text_of(m.get("content")) + str(m.get("reasoning_content") or "") total += ( - len(encoder.encode(text_of(m.get("content")), disallowed_special=())) + len(encoder.encode(text, disallowed_special=())) + self._MSG_OVERHEAD ) return total @@ -82,7 +83,7 @@ def _estimate_heuristic(self, messages: list[Message]) -> int: """Regex-based heuristic for mixed CJK/English text.""" total = 0 for m in messages: - text = text_of(m.get("content")) + text = text_of(m.get("content")) + str(m.get("reasoning_content") or "") cjk_count = len(self._CJK_RE.findall(text)) other_count = len(text) - cjk_count total += cjk_count + (other_count // 4) + self._MSG_OVERHEAD diff --git a/agent_core/runtime/loop/compact.py b/agent_core/runtime/loop/compact.py index 3c3d680..31d61a5 100644 --- a/agent_core/runtime/loop/compact.py +++ b/agent_core/runtime/loop/compact.py @@ -388,7 +388,7 @@ def compress_tool_results( continue clone["content"] = _condense(content, budget) compacted.append(clone) - return compacted + return strip_stale_reasoning(compacted, keep_last=1) @runtime_checkable @@ -428,9 +428,35 @@ def estimate_tokens(messages: list[Message]) -> int: total = 0 for msg in messages: total += _est(text_of(msg.get("content"))) + 4 # +4 per message overhead + # ``reasoning_content`` rides the wire on DeepSeek-V4 / o-series turns + # (see ``WIRE_MESSAGE_KEYS``); leaving it out undercounts those prompts + # by an order of magnitude and compaction never fires. + total += _est(str(msg.get("reasoning_content") or "")) return total +def strip_stale_reasoning(messages: list[Message], keep_last: int) -> list[Message]: + """Drop ``reasoning_content`` from all but the last ``keep_last`` assistants. + + Older traces are the bulk of a long thinking run's prompt. The caller + chooses how many recent traces to retain for its provider. ``keep_last < 0`` + keeps everything. Returns the input list when unchanged. + """ + if keep_last < 0: + return messages + idx = [i for i, m in enumerate(messages) + if is_assistant_msg(m) and m.get("reasoning_content")] + stale = set(idx[:-keep_last] if keep_last else idx) + if not stale: + return messages + out: list[Message] = [] + for i, m in enumerate(messages): + if i in stale: + m = cast(Message, {k: v for k, v in m.items() if k != "reasoning_content"}) + out.append(m) + return out + + def partition_for_compaction( messages: list[Message], keep_recent: int, @@ -637,10 +663,10 @@ class KeepLastNToolResultsCompactor: later turn from re-issuing a query whose result it can no longer see. When no spill is configured and the card would not be shorter than the body it replaces, the body is kept verbatim instead. - ``SystemMessage``, ``HumanMessage``, and every ``AIMessage`` (including - its thinking trace) are left intact, so the model retains its full - chain of reasoning and tool-call metadata while dropping the bulk of - old tool-result bodies (the dominant context cost in long ReAct runs). + ``SystemMessage`` and ``HumanMessage`` are left intact. Assistant content + and tool-call metadata stay intact, while old ``reasoning_content`` is + removed. At least the latest assistant reasoning is retained, including + when ``keep_tool_result`` is zero. Idempotent: already-placeheld messages are detected by content match and left alone, so this is safe to invoke on every turn. @@ -695,6 +721,8 @@ def compact( ) -> list[Message]: if self._keep == -1: return messages + # Old thinking traces go with old tool bodies: same window, same reason. + messages = strip_stale_reasoning(messages, keep_last=max(self._keep, 1)) tool_indices = [i for i, m in enumerate(messages) if is_tool_msg(m)] if not tool_indices: diff --git a/changes/79.fix.md b/changes/79.fix.md new file mode 100644 index 0000000..5da435e --- /dev/null +++ b/changes/79.fix.md @@ -0,0 +1 @@ +Compaction now counts `reasoning_content` in loop and middleware token estimates and drops it from older assistant turns (Tier 1 keeps at least the latest turn, or the last `keep_tool_result` turns when greater than one; tool-result compression and the Tier 2 input keep only the latest). Long `thinking_format: reasoning_content` runs previously looked ~19x smaller than their real prompts, so compaction rarely fired; expect compaction to trigger earlier and prompts to shrink on those profiles. diff --git a/tests/test_compaction_reasoning_content.py b/tests/test_compaction_reasoning_content.py new file mode 100644 index 0000000..aaa3748 --- /dev/null +++ b/tests/test_compaction_reasoning_content.py @@ -0,0 +1,115 @@ +"""Compaction must see and shed ``reasoning_content`` (DeepSeek-V4 round-trip).""" + +from __future__ import annotations + +import pytest + +from agent_core.components.middleware.llm import compaction as mw +from agent_core.components.middleware.llm.base import LLMCallContext +from agent_core.components.middleware.llm.summarization import SummarizationMiddleware +from agent_core.messages import assistant_msg_with_reasoning +from agent_core.runtime.loop.compact import ( + KeepLastNToolResultsCompactor, + compress_tool_results, + estimate_tokens, + strip_stale_reasoning, +) + + +def _history(turns: int, reasoning_chars: int = 20_000) -> list[dict]: + msgs: list[dict] = [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}] + for i in range(turns): + msgs.append(assistant_msg_with_reasoning( + "", f"r{i} " * (reasoning_chars // 4), + tool_calls=[{"id": f"c{i}", "type": "function", + "function": {"name": "bash", "arguments": "{}"}}], + thinking_format="reasoning_content", + )) + msgs.append({"role": "tool", "tool_call_id": f"c{i}", "content": "out"}) + return msgs + + +def _with_reasoning(msgs: list[dict]) -> list[int]: + return [i for i, m in enumerate(msgs) if m.get("reasoning_content")] + + +def test_estimators_count_reasoning_content() -> None: + msgs = _history(4) + bare = [{k: v for k, v in m.items() if k != "reasoning_content"} for m in msgs] + assert estimate_tokens(msgs) > 10 * estimate_tokens(bare) + assert mw.estimate_tokens(msgs) > 10 * mw.estimate_tokens(bare) + + +def test_summarization_middleware_estimators_count_reasoning_content() -> None: + msgs = [{"role": "assistant", "content": "answer", "reasoning_content": "think" * 100}] + middleware = SummarizationMiddleware() + + class RecordingEncoder: + def __init__(self) -> None: + self.inputs: list[str] = [] + + def encode(self, text: str, *, disallowed_special: tuple) -> range: + self.inputs.append(text) + return range(len(text)) + + encoder = RecordingEncoder() + assert middleware._estimate_tiktoken(encoder, msgs) == len("answer") + 500 + 4 + assert encoder.inputs == ["answer" + "think" * 100] + assert middleware._estimate_heuristic(msgs) > 100 + + +@pytest.mark.asyncio +async def test_summarization_middleware_triggers_on_reasoning_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from agent_core.runtime.loop import tokenizer + + monkeypatch.setattr(tokenizer, "get_encoding_nonblocking", lambda _name: None) + msgs = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "old question"}, + {"role": "assistant", "content": "", "reasoning_content": "x" * 20_000}, + {"role": "user", "content": "new question"}, + ] + middleware = SummarizationMiddleware(threshold=1_000, keep_recent=1) + result = await middleware.before_llm(LLMCallContext(), msgs) + + assert result is not msgs + assert [m["role"] for m in result] == ["system", "user", "user"] + assert result[-1] is msgs[-1] + assert msgs[2]["reasoning_content"] == "x" * 20_000 + + +def test_strip_keeps_latest_and_does_not_mutate() -> None: + msgs = _history(5) + out = strip_stale_reasoning(msgs, keep_last=2) + assert _with_reasoning(out) == _with_reasoning(msgs)[-2:] + assert len(_with_reasoning(msgs)) == 5 + assert strip_stale_reasoning(msgs, keep_last=-1) is msgs + assert _with_reasoning(strip_stale_reasoning(msgs, keep_last=0)) == [] + + +def test_tier1_sheds_reasoning_outside_its_window() -> None: + msgs = _history(12) + out = KeepLastNToolResultsCompactor(keep_tool_result=3).compact(msgs, 5) + assert len(_with_reasoning(out)) == 3 + assert estimate_tokens(out) < estimate_tokens(msgs) / 3 + + +def test_tier1_disabled_keeps_reasoning() -> None: + msgs = _history(4) + assert KeepLastNToolResultsCompactor(keep_tool_result=-1).compact(msgs, 5) is msgs + + +def test_tier1_zero_keeps_latest_reasoning_without_mutating_source() -> None: + msgs = _history(4) + out = KeepLastNToolResultsCompactor(keep_tool_result=0).compact(msgs, 5) + assert _with_reasoning(out) == _with_reasoning(msgs)[-1:] + assert len(_with_reasoning(msgs)) == 4 + + +def test_compress_tool_results_keeps_only_latest_reasoning() -> None: + msgs = _history(6) + out = compress_tool_results(msgs) + assert _with_reasoning(out) == _with_reasoning(msgs)[-1:] + assert out[-2]["tool_calls"] == msgs[-2]["tool_calls"]