fix(reflect): enforce max_tokens on the done-tool completion path (#2756)

hindsight-api-slim/hindsight_api/engine/reflect/agent.py

@@ -15,7 +15,15 @@
 from typing import TYPE_CHECKING, Any, Awaitable, Callable
 
 from ...config import get_config
-from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
+from .models import (
+    DirectiveInfo,
+    LLMCall,
+    MaxTokensCapResult,
+    ReflectAgentResult,
+    StructuredOutputResult,
+    TokenUsageSummary,
+    ToolCall,
+)
 from .prompts import (
     _extract_directive_rules,
     build_final_prompt,
@@ -358,6 +366,60 @@ def _json_schema_type_to_python(field_schema: dict) -> type:
         return StructuredOutputResult()
 
 
+async def _enforce_answer_token_cap(
+    answer: str,
+    max_tokens: int | None,
+    llm_config: "LLMProvider | None",
+    reflect_id: str,
+) -> MaxTokensCapResult:
+    """Rewrite ``answer`` to fit within ``max_tokens`` if it overshoots.
+
+    The reflect agent's tool-driven completion is intentionally uncapped so the
+    LLM has headroom to emit tool-call JSON and intermediate reasoning. But once
+    the agent's text becomes the user-visible final answer it must respect
+    max_tokens. Every completion path (forced-final, direct-text short-circuit,
+    and the done tool) routes through this helper so the cap is enforced
+    uniformly. See issue #2756.
+
+    When the answer is already within budget (or no cap/LLM is configured), no
+    call is made and the original answer is returned with ``rewritten=False``.
+    """
+    if max_tokens is None or llm_config is None or not answer:
+        return MaxTokensCapResult(answer=answer)
+    if count_cl100k_tokens(answer) <= max_tokens:
+        return MaxTokensCapResult(answer=answer)
+
+    rewrite_start = time.time()
+    rewritten, usage = await llm_config.call(
+        messages=[
+            {
+                "role": "system",
+                "content": (
+                    "Rewrite the user's text so it fits within the requested token "
+                    "budget. Preserve the key facts and structure; drop lower-priority "
+                    "detail. Respond with the rewritten text only, no preamble."
+                ),
+            },
+            {
+                "role": "user",
+                "content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}",
+            },
+        ],
+        scope="reflect",
+        max_completion_tokens=max_tokens,
+        return_usage=True,
+    )
+    return MaxTokensCapResult(
+        answer=_clean_answer_text(rewritten.strip()),
+        rewritten=True,
+        input_tokens=usage.input_tokens,
+        output_tokens=usage.output_tokens,
+        cached_tokens=getattr(usage, "cached_tokens", 0) or 0,
+        thoughts_tokens=getattr(usage, "thoughts_tokens", 0) or 0,
+        duration_ms=int((time.time() - rewrite_start) * 1000),
+    )
+
+
 def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
     """Estimate the token count of the messages list using cl100k_base encoding."""
     total = 0
@@ -870,40 +932,21 @@ def _log_completion(answer: str, iterations: int, forced: bool = False):
                 # must respect max_tokens like the forced-final paths do. If it
                 # overshoots, run one extra capped call to rewrite it within
                 # the cap.
-                if max_tokens is not None and count_cl100k_tokens(answer) > max_tokens:
-                    rewrite_start = time.time()
-                    rewritten, rewrite_usage = await llm_config.call(
-                        messages=[
-                            {
-                                "role": "system",
-                                "content": (
-                                    "Rewrite the user's text so it fits within the requested token "
-                                    "budget. Preserve the key facts and structure; drop lower-priority "
-                                    "detail. Respond with the rewritten text only, no preamble."
-                                ),
-                            },
-                            {
-                                "role": "user",
-                                "content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}",
-                            },
-                        ],
-                        scope="reflect",
-                        max_completion_tokens=max_tokens,
-                        return_usage=True,
-                    )
-                    total_input_tokens += rewrite_usage.input_tokens
-                    total_output_tokens += rewrite_usage.output_tokens
-                    total_cached_tokens += getattr(rewrite_usage, "cached_tokens", 0) or 0
-                    total_thoughts_tokens += getattr(rewrite_usage, "thoughts_tokens", 0) or 0
+                cap = await _enforce_answer_token_cap(answer, max_tokens, llm_config, reflect_id)
+                if cap.rewritten:
+                    answer = cap.answer
+                    total_input_tokens += cap.input_tokens
+                    total_output_tokens += cap.output_tokens
+                    total_cached_tokens += cap.cached_tokens
+                    total_thoughts_tokens += cap.thoughts_tokens
                     llm_trace.append(
                         {
                             "scope": "final_rewrite",
-                            "duration_ms": int((time.time() - rewrite_start) * 1000),
-                            "input_tokens": rewrite_usage.input_tokens,
-                            "output_tokens": rewrite_usage.output_tokens,
+                            "duration_ms": cap.duration_ms,
+                            "input_tokens": cap.input_tokens,
+                            "output_tokens": cap.output_tokens,
                         }
                     )
-                    answer = _clean_answer_text(rewritten.strip())
 
                 # Generate structured output if schema provided
                 structured_output = None
@@ -1035,6 +1078,7 @@ def _log_completion(answer: str, iterations: int, forced: bool = False):
                     directives_applied=directives_applied,
                     llm_config=llm_config,
                     response_schema=response_schema,
+                    max_tokens=max_tokens,
                 )
 
         # Execute other tools in parallel (exclude done tool in all its format variants)
@@ -1244,6 +1288,7 @@ async def _process_done_tool(
     directives_applied: list[DirectiveInfo],
     llm_config: "LLMProvider | None" = None,
     response_schema: dict | None = None,
+    max_tokens: int | None = None,
 ) -> ReflectAgentResult:
     """Process the done tool call and return the result."""
     args = done_call.arguments
@@ -1254,6 +1299,30 @@ async def _process_done_tool(
     if not answer:
         answer = "No answer provided."
 
+    # The done tool is the normal completion path, and its answer comes verbatim
+    # from the (intentionally uncapped) tool-call arguments. Enforce max_tokens
+    # here too, exactly like the forced-final and direct-text short-circuit paths,
+    # so a per-model cap is honoured no matter how the agent finishes (#2756).
+    cap = await _enforce_answer_token_cap(answer, max_tokens, llm_config, reflect_id)
+    if cap.rewritten:
+        answer = cap.answer
+        usage = TokenUsageSummary(
+            input_tokens=usage.input_tokens + cap.input_tokens,
+            output_tokens=usage.output_tokens + cap.output_tokens,
+            total_tokens=usage.total_tokens + cap.input_tokens + cap.output_tokens,
+            cached_tokens=usage.cached_tokens + cap.cached_tokens,
+            thoughts_tokens=usage.thoughts_tokens + cap.thoughts_tokens,
+        )
+        llm_trace = [
+            *llm_trace,
+            LLMCall(
+                scope="final_rewrite",
+                duration_ms=cap.duration_ms,
+                input_tokens=cap.input_tokens,
+                output_tokens=cap.output_tokens,
+            ),
+        ]
+
     # Validate IDs (only include IDs that were actually retrieved)
     used_memory_ids = [mid for mid in (args.get("memory_ids") or []) if mid in available_memory_ids]
     used_mental_model_ids = [mid for mid in (args.get("mental_model_ids") or []) if mid in available_mental_model_ids]

hindsight-api-slim/hindsight_api/engine/reflect/models.py

@@ -106,6 +106,24 @@ class StructuredOutputResult(BaseModel):
     thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider")
 
 
+class MaxTokensCapResult(BaseModel):
+    """Result of enforcing a max_tokens cap on a final answer.
+
+    Carries the (possibly rewritten) answer plus the token usage and timing of
+    the rewrite call, so the caller can fold it into its own accounting. When
+    ``rewritten`` is False the answer was already within budget and no LLM call
+    was made (all token counts are zero).
+    """
+
+    answer: str = Field(description="The answer, rewritten within budget if it overshot")
+    rewritten: bool = Field(default=False, description="Whether a capped rewrite call was actually made")
+    input_tokens: int = Field(default=0, description="Input tokens used by the rewrite call")
+    output_tokens: int = Field(default=0, description="Visible output tokens used by the rewrite call")
+    cached_tokens: int = Field(default=0, description="Cached prefix tokens. Subset of input_tokens.")
+    thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider")
+    duration_ms: int = Field(default=0, description="Wall-clock duration of the rewrite call")
+
+
 class ReflectAgentResult(BaseModel):
     """Result from the reflect agent."""
 

hindsight-api-slim/tests/test_reflect_agent.py

@@ -1263,3 +1263,96 @@ async def test_real_stale_mental_model_forces_and_grounds_in_deeper_evidence(sel
             context="A stale mental model claimed the launch was still pending, but the freshly retrieved raw "
             "fact (deploy log A-1029) shows it shipped on Friday. The agent should correct the stale summary.",
         )
+
+
+class TestDoneToolMaxTokensCap:
+    """Regression tests for #2756: max_tokens must be enforced on the done-tool
+    completion path (the normal path), not only the forced-final and direct-text
+    short-circuit paths."""
+
+    @staticmethod
+    def _done_call(answer: str) -> "LLMToolCall":
+        return LLMToolCall(id="done-1", name="done", arguments={"answer": answer})
+
+    @staticmethod
+    def _process(done_call, llm_config, max_tokens, llm_trace=None):
+        from hindsight_api.engine.reflect.agent import _process_done_tool
+        from hindsight_api.engine.reflect.models import TokenUsageSummary
+
+        return _process_done_tool(
+            done_call,
+            set(),
+            set(),
+            set(),
+            2,  # iterations
+            1,  # total_tools_called
+            [],  # tool_trace
+            list(llm_trace or []),  # llm_trace
+            TokenUsageSummary(input_tokens=100, output_tokens=50, total_tokens=150),
+            lambda *a, **k: None,  # log_completion
+            "test-reflect",  # reflect_id
+            directives_applied=[],
+            llm_config=llm_config,
+            response_schema=None,
+            max_tokens=max_tokens,
+        )
+
+    @pytest.mark.asyncio
+    async def test_done_tool_over_budget_answer_is_rewritten(self):
+        """An over-budget done answer triggers a single capped rewrite call."""
+        long_answer = "This is a sentence with several words. " * 60  # well over 50 tokens
+        llm = MagicMock()
+        llm.call = AsyncMock(
+            return_value=("Short capped answer.", TokenUsage(input_tokens=200, output_tokens=20, total_tokens=220))
+        )
+
+        result = await self._process(self._done_call(long_answer), llm, max_tokens=50)
+
+        # Cap enforced: the rewritten (short) answer is returned, not the raw one.
+        assert result.text == "Short capped answer."
+        # Exactly one rewrite call, with the budget forwarded.
+        llm.call.assert_awaited_once()
+        kwargs = llm.call.await_args.kwargs
+        assert kwargs["max_completion_tokens"] == 50
+        assert kwargs["scope"] == "reflect"
+        # Rewrite usage is folded into the result and traced.
+        assert result.usage.input_tokens == 100 + 200
+        assert result.usage.output_tokens == 50 + 20
+        assert any(c.scope == "final_rewrite" for c in result.llm_trace)
+
+    @pytest.mark.asyncio
+    async def test_done_tool_under_budget_answer_is_untouched(self):
+        """An answer already within budget makes no rewrite call and is returned as-is."""
+        llm = MagicMock()
+        llm.call = AsyncMock()
+
+        result = await self._process(self._done_call("Concise."), llm, max_tokens=50)
+
+        assert result.text == "Concise."
+        llm.call.assert_not_awaited()
+        assert not any(c.scope == "final_rewrite" for c in result.llm_trace)
+
+    @pytest.mark.asyncio
+    async def test_done_tool_no_cap_configured_is_untouched(self):
+        """With max_tokens=None the answer is never rewritten, however long."""
+        long_answer = "This is a sentence with several words. " * 60
+        llm = MagicMock()
+        llm.call = AsyncMock()
+
+        result = await self._process(self._done_call(long_answer), llm, max_tokens=None)
+
+        assert result.text == _clean_done_answer(long_answer.strip())
+        llm.call.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_enforce_cap_helper_skips_when_no_cap(self):
+        """The shared helper is a no-op (no LLM call) when max_tokens is None."""
+        from hindsight_api.engine.reflect.agent import _enforce_answer_token_cap
+
+        llm = MagicMock()
+        llm.call = AsyncMock()
+        cap = await _enforce_answer_token_cap("some long text " * 50, None, llm, "rid")
+
+        assert cap.rewritten is False
+        assert cap.answer == "some long text " * 50
+        llm.call.assert_not_awaited()