fix(reflect): enforce max_tokens after forced synthesis

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

@@ -18,6 +18,7 @@
 from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMToolChoice
 from ..llm_trace import LLMQueueWait, reset_queue_wait_sink, set_queue_wait_sink
 from ..llm_transport import describe_llm_error
+from ..response_models import TokenUsage
 from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
 from .prompts import (
     _SPLIT_SYNTHESIS_WARN_CHUNKS,
@@ -836,6 +837,26 @@ async def _forced_final_synthesis(iterations_completed: int) -> ReflectAgentResu
                 f"over {len(chunks)} context chunk(s)."
             )
 
+        answer, _, rewrite_usage, rewrite_trace = await _rewrite_final_answer(
+            answer,
+            llm_config,
+            max_tokens,
+        )
+        if rewrite_usage is not None:
+            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
+        if rewrite_trace is not None:
+            llm_trace.append(
+                {
+                    "scope": rewrite_trace.scope,
+                    "duration_ms": rewrite_trace.duration_ms,
+                    "input_tokens": rewrite_trace.input_tokens,
+                    "output_tokens": rewrite_trace.output_tokens,
+                }
+            )
+
         structured_output = None
         # ``answer`` is non-empty past the guard above, so only the schema gates this.
         if response_schema:
@@ -1340,6 +1361,68 @@ def _document_from_rewrite(rewritten: str, previous_answer: str) -> CanonicalDoc
     return CanonicalDocument(markdown=text, structure=split_markdown(text))
 
 
+async def _rewrite_final_answer(
+    answer: str,
+    llm_config: "LLMProvider | None",
+    max_tokens: int | None,
+    document: StructuredDocument | None = None,
+) -> tuple[str, StructuredDocument | None, TokenUsage | None, LLMCall | None]:
+    """Shorten an over-budget final answer, regardless of how reflect completed."""
+    if llm_config is None or max_tokens is None or count_prompt_tokens(answer) <= max_tokens:
+        return answer, document, None, None
+
+    rewrite_start = time.time()
+    # In document mode the trim is asked for as a document too. Asking for
+    # prose here would put the model back in the business of writing the
+    # markdown that gets stored — on the one path where the answer is long
+    # enough that its structure matters most.
+    if document is not None:
+        rewrite_system = (
+            "Shorten the user's document so it fits within the requested token budget. "
+            "Preserve the key facts and the document's structure; drop lower-priority detail. "
+            'Respond ONLY with JSON: {"sections": [{"heading": "...", "level": 2, '
+            '"blocks": ["...", "..."]}]}. A heading carries no "#", and each block is one '
+            "paragraph, list, table or code fence."
+        )
+        rewrite_user = f"Target budget: {max_tokens} tokens.\n\nDocument to shorten:\n{answer}"
+    else:
+        # The token budget is enforced via the prompt, not a hard provider cap:
+        # on thinking models a hard cap is eaten by reasoning tokens and would
+        # truncate the rewrite mid-word (#3365). Cost is bounded by the separate
+        # reflect_max_completion_tokens config (uncapped by default).
+        rewrite_system = (
+            "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."
+        )
+        rewrite_user = f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}"
+
+    call_result = await llm_config.call(
+        messages=[
+            {"role": "system", "content": rewrite_system},
+            {"role": "user", "content": rewrite_user},
+        ],
+        scope="reflect",
+        temperature=get_config().llm_temperature_reflect,
+        max_completion_tokens=get_config().reflect_max_completion_tokens,
+    )
+    rewritten = call_result.content
+    if document is not None:
+        trimmed = _document_from_rewrite(rewritten, answer)
+        document, answer = trimmed.structure, trimmed.markdown
+    else:
+        answer = rewritten.strip()
+
+    usage = call_result.usage
+    trace = LLMCall(
+        scope="final_rewrite",
+        duration_ms=int((time.time() - rewrite_start) * 1000),
+        input_tokens=usage.input_tokens,
+        output_tokens=usage.output_tokens,
+    )
+    return answer, document, usage, trace
+
+
 async def _process_done_tool(
     done_call: "LLMToolCall",
     available_memory_ids: set[str],
@@ -1389,64 +1472,22 @@ async def _process_done_tool(
         )
 
     final_usage = usage
-    if llm_config and max_tokens is not None and count_prompt_tokens(answer) > max_tokens:
-        rewrite_start = time.time()
-        # In document mode the trim is asked for as a document too. Asking for
-        # prose here would put the model back in the business of writing the
-        # markdown that gets stored — on the one path where the answer is long
-        # enough that its structure matters most.
-        if document is not None:
-            rewrite_system = (
-                "Shorten the user's document so it fits within the requested token budget. "
-                "Preserve the key facts and the document's structure; drop lower-priority detail. "
-                'Respond ONLY with JSON: {"sections": [{"heading": "...", "level": 2, '
-                '"blocks": ["...", "..."]}]}. A heading carries no "#", and each block is one '
-                "paragraph, list, table or code fence."
-            )
-            rewrite_user = f"Target budget: {max_tokens} tokens.\n\nDocument to shorten:\n{answer}"
-        else:
-            # The token budget is enforced via the prompt, not a hard provider cap:
-            # on thinking models a hard cap is eaten by reasoning tokens and would
-            # truncate the rewrite mid-word (#3365). Cost is bounded by the separate
-            # reflect_max_completion_tokens config (uncapped by default).
-            rewrite_system = (
-                "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."
-            )
-            rewrite_user = f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}"
-
-        call_result = await llm_config.call(
-            messages=[
-                {"role": "system", "content": rewrite_system},
-                {"role": "user", "content": rewrite_user},
-            ],
-            scope="reflect",
-            temperature=get_config().llm_temperature_reflect,
-            max_completion_tokens=get_config().reflect_max_completion_tokens,
-        )
-        rewritten = call_result.content
-        rewrite_usage = call_result.usage
-        if document is not None:
-            trimmed = _document_from_rewrite(rewritten, answer)
-            document, answer = trimmed.structure, trimmed.markdown
-        else:
-            answer = rewritten.strip()
+    answer, document, rewrite_usage, rewrite_trace = await _rewrite_final_answer(
+        answer,
+        llm_config,
+        max_tokens,
+        document,
+    )
+    if rewrite_usage is not None:
+        assert rewrite_trace is not None
         final_usage = TokenUsageSummary(
             input_tokens=usage.input_tokens + rewrite_usage.input_tokens,
             output_tokens=usage.output_tokens + rewrite_usage.output_tokens,
             total_tokens=usage.total_tokens + rewrite_usage.input_tokens + rewrite_usage.output_tokens,
             cached_tokens=usage.cached_tokens + (getattr(rewrite_usage, "cached_tokens", 0) or 0),
             thoughts_tokens=usage.thoughts_tokens + (getattr(rewrite_usage, "thoughts_tokens", 0) or 0),
         )
-        llm_trace.append(
-            LLMCall(
-                scope="final_rewrite",
-                duration_ms=int((time.time() - rewrite_start) * 1000),
-                input_tokens=rewrite_usage.input_tokens,
-                output_tokens=rewrite_usage.output_tokens,
-            )
-        )
+        llm_trace.append(rewrite_trace)
 
     # 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]

hindsight-api-slim/tests/test_reflect_agent.py

@@ -1012,6 +1012,54 @@ async def test_stop_after_evidence_uses_forced_final_synthesis(self, mock_llm, m
         final_prompt = mock_llm.call.await_args.kwargs["messages"][1]["content"]
         assert f"approximately {cap} tokens" in final_prompt
 
+    @pytest.mark.asyncio
+    async def test_forced_final_synthesis_rewrites_over_budget_answer(self, mock_llm, mock_functions, monkeypatch):
+        config = MagicMock(
+            reflect_prompt_cache_enabled=False,
+            reflect_max_completion_tokens=None,
+            llm_temperature_reflect=0.17,
+        )
+        monkeypatch.setattr("hindsight_api.engine.reflect.agent.get_config", lambda: config)
+        mock_functions["search_mental_models_fn"].return_value = {
+            "mental_models": [{"id": "mm-1", "name": "Prefs", "content": "Fresh content.", "is_stale": False}]
+        }
+        mock_llm.call_with_tools.side_effect = [
+            self._mm_call(),
+            LLMToolCallResult(tool_calls=[], content="I have enough to answer.", finish_reason="stop"),
+        ]
+        mock_llm.call = AsyncMock(
+            side_effect=[
+                LLMCallResult(
+                    content="important detail " * 100,
+                    usage=TokenUsage(input_tokens=40, output_tokens=100, total_tokens=140),
+                ),
+                LLMCallResult(
+                    content="Concise final answer.",
+                    usage=TokenUsage(input_tokens=110, output_tokens=4, total_tokens=114),
+                ),
+            ]
+        )
+
+        result = await run_reflect_agent(
+            llm_config=mock_llm,
+            bank_id="test-bank",
+            query="test query",
+            bank_profile={"name": "Test", "mission": "Testing"},
+            has_mental_models=True,
+            budget="low",
+            max_tokens=8,
+            **mock_functions,
+        )
+
+        assert result.text == "Concise final answer."
+        assert mock_llm.call.await_count == 2
+        rewrite_call = mock_llm.call.await_args_list[1]
+        assert "Target budget: 8 tokens" in rewrite_call.kwargs["messages"][1]["content"]
+        assert rewrite_call.kwargs["max_completion_tokens"] is None
+        assert result.llm_trace[-1].scope == "final_rewrite"
+        assert result.usage.input_tokens == 150
+        assert result.usage.output_tokens == 104
+
     @pytest.mark.asyncio
     async def test_max_iterations_reached(self, mock_llm, mock_functions):
         """Test that agent stops after max iterations even with errors."""