fix(reflect): derive final prompt context budget from max_context_tokens

Replace the hardcoded _FINAL_PROMPT_CONTEXT_BUDGET (60K tokens) with a fraction of max_context_tokens (80%), so the fallback synthesis prompt automatically scales with whatever context window is configured.
fix(reflect): prevent context_length_exceeded on large memory banks (#457 )
2026-03-02 11:38:30 +01:00 · 2026-03-02 10:32:51 +01:00
7 changed files with 342 additions and 12 deletions
@@ -314,6 +314,7 @@ ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLO

 # Reflect agent settings
 ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
+ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
 ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"

 # Disposition settings
@@ -453,6 +454,7 @@ DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2  # Max concurrent consolidation tasks

 # Reflect agent settings
 DEFAULT_REFLECT_MAX_ITERATIONS = 10  # Max tool call iterations before forcing response
+DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000  # Max accumulated context tokens before forcing final prompt

 # Disposition defaults (None = not set, fall back to bank DB value or 3)
 DEFAULT_DISPOSITION_SKEPTICISM = None
@@ -720,6 +722,7 @@ class HindsightConfig:

    # Reflect agent settings
    reflect_max_iterations: int
+    reflect_max_context_tokens: int

    # OpenTelemetry tracing configuration
    otel_traces_enabled: bool
@@ -1134,6 +1137,9 @@ class HindsightConfig:
            ),
            # Reflect agent settings
            reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
+            reflect_max_context_tokens=int(
+                os.getenv(ENV_REFLECT_MAX_CONTEXT_TOKENS, str(DEFAULT_REFLECT_MAX_CONTEXT_TOKENS))
+            ),
            reflect_mission=os.getenv(ENV_REFLECT_MISSION) or None,
            # Disposition settings (None = fall back to DB value)
            disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM))
@@ -4582,6 +4582,7 @@ class MemoryEngine(MemoryEngineInterface):
        budget_multipliers = {Budget.LOW: 0.5, Budget.MID: 1.0, Budget.HIGH: 2.0}
        effective_budget = budget or Budget.LOW
        max_iterations = max(1, int(base_max_iterations * budget_multipliers.get(effective_budget, 1.0)))
+        max_context_tokens = config.reflect_max_context_tokens

        # Run agentic loop - acquire connections only when needed for DB operations
        # (not held during LLM calls which can be slow)
@@ -4691,6 +4692,7 @@ class MemoryEngine(MemoryEngineInterface):
                directives=directives,
                has_mental_models=has_mental_models,
                budget=effective_budget,
+                max_context_tokens=max_context_tokens,
            )

            total_time = time.time() - reflect_start
@@ -14,6 +14,8 @@ import re
 import time
 from typing import TYPE_CHECKING, Any, Awaitable, Callable

+import tiktoken
+
 from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
 from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools
 from .tools_schema import get_reflect_tools
@@ -259,6 +261,46 @@ OUTPUT:"""
        return None, 0, 0


+_TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
+
+
+def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
+    """Estimate the token count of the messages list using cl100k_base encoding."""
+    total = 0
+    for msg in messages:
+        content = msg.get("content") or ""
+        if isinstance(content, str):
+            total += len(_TIKTOKEN_ENCODING.encode(content))
+        elif isinstance(content, list):
+            for part in content:
+                if isinstance(part, dict) and isinstance(part.get("text"), str):
+                    total += len(_TIKTOKEN_ENCODING.encode(part["text"]))
+        # Tool call arguments and results also count
+        for tc in msg.get("tool_calls") or []:
+            if isinstance(tc, dict):
+                func = tc.get("function", {})
+                total += len(_TIKTOKEN_ENCODING.encode(func.get("arguments", "")))
+    return total
+
+
+def _is_context_overflow_error(exc: Exception) -> bool:
+    """Return True if the exception signals the LLM context window was exceeded."""
+    msg = str(exc).lower()
+    return any(
+        phrase in msg
+        for phrase in (
+            "context_length_exceeded",
+            "context length exceeded",
+            "maximum context length",
+            "prompt_too_long",
+            "prompt is too long",
+            "resource_exhausted",
+            "input is too long",
+            "too many tokens",
+        )
+    )
+
+
 async def run_reflect_agent(
    llm_config: "LLMProvider",
    bank_id: str,
@@ -275,6 +317,7 @@ async def run_reflect_agent(
    directives: list[dict[str, Any]] | None = None,
    has_mental_models: bool = False,
    budget: str | None = None,
+    max_context_tokens: int = 100_000,
 ) -> ReflectAgentResult:
    """
    Execute the reflect agent loop using native tool calling.
@@ -388,7 +431,7 @@ async def run_reflect_agent(

        if is_last:
            # Force text response on last iteration - no tools
-            prompt = build_final_prompt(query, context_history, bank_profile, context)
+            prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens)
            llm_start = time.time()
            response, usage = await llm_config.call(
                messages=[
@@ -433,6 +476,60 @@ async def run_reflect_agent(
                directives_applied=directives_applied,
            )

+        # Proactive context-window guard: if accumulated messages would exceed the
+        # configured token budget, bail out early and synthesize from what we have.
+        estimated_tokens = _count_messages_tokens(messages)
+        if estimated_tokens >= max_context_tokens and (
+            bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
+        ):
+            logger.warning(
+                f"[REFLECT {reflect_id}] Context budget exceeded on iteration {iteration + 1}: "
+                f"~{estimated_tokens} tokens >= {max_context_tokens} limit. Forcing final synthesis."
+            )
+            prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens)
+            llm_start = time.time()
+            response, usage = await llm_config.call(
+                messages=[
+                    {"role": "system", "content": FINAL_SYSTEM_PROMPT},
+                    {"role": "user", "content": prompt},
+                ],
+                scope="reflect",
+                max_completion_tokens=max_tokens,
+                return_usage=True,
+            )
+            llm_duration = int((time.time() - llm_start) * 1000)
+            total_input_tokens += usage.input_tokens
+            total_output_tokens += usage.output_tokens
+            llm_trace.append(
+                {
+                    "scope": "final",
+                    "duration_ms": llm_duration,
+                    "input_tokens": usage.input_tokens,
+                    "output_tokens": usage.output_tokens,
+                }
+            )
+            answer = _clean_answer_text(response.strip())
+
+            structured_output = None
+            if response_schema and answer:
+                structured_output, struct_in, struct_out = await _generate_structured_output(
+                    answer, response_schema, llm_config, reflect_id
+                )
+                total_input_tokens += struct_in
+                total_output_tokens += struct_out
+
+            _log_completion(answer, iteration + 1, forced=True)
+            return ReflectAgentResult(
+                text=answer,
+                structured_output=structured_output,
+                iterations=iteration + 1,
+                tools_called=total_tools_called,
+                tool_trace=tool_trace,
+                llm_trace=_get_llm_trace(),
+                usage=_get_usage(),
+                directives_applied=directives_applied,
+            )
+
        # Call LLM with tools
        llm_start = time.time()

@@ -478,13 +575,20 @@ async def run_reflect_agent(
            consecutive_errors += 1
            logger.warning(f"[REFLECT {reflect_id}] LLM error on iteration {iteration + 1}: {e} ({err_duration}ms)")
            llm_trace.append({"scope": f"agent_{iteration + 1}_err", "duration_ms": err_duration})
-            # Guardrail: If no evidence gathered yet, retry (but cap consecutive errors to avoid long hangs)
            has_gathered_evidence = (
                bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
            )
-            if not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2:
+            # Context overflow errors must never be retried — retrying would only make them worse.
+            # Skip straight to final synthesis with whatever evidence we have.
+            if _is_context_overflow_error(e):
+                logger.warning(
+                    f"[REFLECT {reflect_id}] Context window exceeded on iteration {iteration + 1}, "
+                    "forcing final synthesis from gathered evidence."
+                )
+            # For other errors: retry if no evidence yet (but cap consecutive errors to avoid long hangs)
+            elif not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2:
                continue
-            prompt = build_final_prompt(query, context_history, bank_profile, context)
+            prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens)
            llm_start = time.time()
            response, usage = await llm_config.call(
                messages=[
@@ -555,7 +659,7 @@ async def run_reflect_agent(
                    directives_applied=directives_applied,
                )
            # Empty response, force final
-            prompt = build_final_prompt(query, context_history, bank_profile, context)
+            prompt = build_final_prompt(query, context_history, bank_profile, context, max_context_tokens=max_context_tokens)
            llm_start = time.time()
            response, usage = await llm_config.call(
                messages=[
@@ -10,6 +10,14 @@ The reflect agent uses hierarchical retrieval:
 import json
 from typing import Any

+import tiktoken
+
+_TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
+
+# Fraction of max_context_tokens reserved for tool results in the final synthesis prompt.
+# The remainder covers the system prompt, question, bank context, and output tokens.
+_FINAL_PROMPT_CONTEXT_FRACTION = 0.8
+

 def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]:
    """Extract directive rules as a list of strings."""
@@ -394,6 +402,7 @@ def build_final_prompt(
    context_history: list[dict],
    bank_profile: dict,
    additional_context: str | None = None,
+    max_context_tokens: int = 100_000,
 ) -> str:
    """Build the final prompt when forcing a text response (no tools)."""
    parts = []
@@ -423,18 +432,32 @@ def build_final_prompt(
    if additional_context:
        parts.append(f"\n## Additional Context\n{additional_context}")

-    # Tool call history
+    # Tool call history — include as many entries as fit within the token budget,
+    # preferring the most recent calls (they tend to be the most targeted).
    if context_history:
        parts.append("\n## Retrieved Data (synthesize and reason from this data)")
-        for entry in context_history:
+        token_budget = int(max_context_tokens * _FINAL_PROMPT_CONTEXT_FRACTION)
+        # Render entries newest-first, then reverse so the prompt reads chronologically.
+        rendered: list[str] = []
+        truncated = False
+        for entry in reversed(context_history):
            tool = entry["tool"]
            output = entry["output"]
-            # Format as proper JSON for LLM readability
            try:
                output_str = json.dumps(output, indent=2, default=str)
            except (TypeError, ValueError):
                output_str = str(output)
-            parts.append(f"\n### From {tool}:\n```json\n{output_str}\n```")
+            block = f"\n### From {tool}:\n```json\n{output_str}\n```"
+            block_tokens = len(_TIKTOKEN_ENCODING.encode(block))
+            if block_tokens > token_budget:
+                truncated = True
+                break
+            rendered.append(block)
+            token_budget -= block_tokens
+        for block in reversed(rendered):
+            parts.append(block)
+        if truncated:
+            parts.append("\n*Note: Some earlier tool results were omitted to stay within the context window.*")
    else:
        parts.append("\n## Retrieved Data\nNo data was retrieved.")

@@ -292,6 +292,7 @@ def main():
            worker_max_slots=config.worker_max_slots,
            worker_consolidation_max_slots=config.worker_consolidation_max_slots,
            reflect_max_iterations=config.reflect_max_iterations,
+            reflect_max_context_tokens=config.reflect_max_context_tokens,
            reflect_mission=config.reflect_mission,
            disposition_skepticism=config.disposition_skepticism,
            disposition_literalism=config.disposition_literalism,
@@ -7,14 +7,17 @@ These tests verify:
 3. Recovery from tool execution errors
 """

+from unittest.mock import AsyncMock, MagicMock
+
 import pytest
-from unittest.mock import AsyncMock, MagicMock, patch

 from hindsight_api.engine.reflect.agent import (
-    _normalize_tool_name,
-    _is_done_tool,
    _clean_answer_text,
    _clean_done_answer,
+    _count_messages_tokens,
+    _is_context_overflow_error,
+    _is_done_tool,
+    _normalize_tool_name,
    run_reflect_agent,
 )
 from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
@@ -412,3 +415,193 @@ class TestReflectAgentMocked:
        # Should have a result even if no memories found
        assert result is not None
        assert result.iterations == 3
+
+
+class TestContextOverflowHelpers:
+    """Unit tests for context-overflow detection helpers."""
+
+    def test_count_messages_tokens_basic(self):
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What is the capital of France?"},
+        ]
+        count = _count_messages_tokens(messages)
+        assert count > 0
+        # Rough sanity check: ~10 tokens for each message
+        assert count < 100
+
+    def test_count_messages_tokens_with_tool_result(self):
+        """A large tool result should substantially increase the count."""
+        small_messages = [{"role": "user", "content": "hi"}]
+        large_messages = [
+            {"role": "user", "content": "hi"},
+            {
+                "role": "tool",
+                "tool_call_id": "x",
+                "name": "recall",
+                "content": '{"memories": [' + ', '.join([f'{{"id": "m{i}", "content": "A long memory fact about some topic that goes on and on."}}' for i in range(50)]) + ']}',
+            },
+        ]
+        small = _count_messages_tokens(small_messages)
+        large = _count_messages_tokens(large_messages)
+        assert large > small + 200
+
+    def test_is_context_overflow_error_openai(self):
+        assert _is_context_overflow_error(Exception("context_length_exceeded: too many tokens"))
+        assert _is_context_overflow_error(Exception("This model's maximum context length is 128000 tokens. However, your messages resulted in 142164 tokens."))
+
+    def test_is_context_overflow_error_anthropic(self):
+        assert _is_context_overflow_error(Exception("prompt_too_long"))
+        assert _is_context_overflow_error(Exception("prompt is too long for this model"))
+
+    def test_is_context_overflow_error_gemini(self):
+        assert _is_context_overflow_error(Exception("RESOURCE_EXHAUSTED: quota exceeded"))
+
+    def test_is_context_overflow_error_generic(self):
+        assert _is_context_overflow_error(Exception("input is too long to process"))
+        assert _is_context_overflow_error(Exception("too many tokens in the request"))
+
+    def test_is_context_overflow_error_unrelated(self):
+        assert not _is_context_overflow_error(Exception("connection timeout"))
+        assert not _is_context_overflow_error(Exception("rate limit exceeded"))
+        assert not _is_context_overflow_error(ValueError("invalid argument"))
+
+
+class TestContextOverflowBehavior:
+    """Test that the reflect agent handles context overflow gracefully."""
+
+    @pytest.fixture
+    def mock_llm(self):
+        llm = MagicMock()
+        llm.call_with_tools = AsyncMock()
+        llm.call = AsyncMock(
+            return_value=("Synthesized answer from gathered evidence.", TokenUsage(input_tokens=50, output_tokens=20, total_tokens=70))
+        )
+        return llm
+
+    @pytest.fixture
+    def mock_functions_with_large_output(self):
+        """Mock functions that return a large enough payload to exceed a tiny token budget."""
+        large_memories = [
+            {"id": f"mem-{i}", "content": f"Memory fact number {i}: " + "A" * 200}
+            for i in range(20)
+        ]
+        return {
+            "search_mental_models_fn": AsyncMock(return_value={"mental_models": []}),
+            "search_observations_fn": AsyncMock(return_value={"observations": []}),
+            "recall_fn": AsyncMock(return_value={"memories": large_memories}),
+            "expand_fn": AsyncMock(return_value={"memories": []}),
+        }
+
+    @pytest.mark.asyncio
+    async def test_proactive_guard_fires_when_budget_exceeded(self, mock_llm, mock_functions_with_large_output):
+        """When token count exceeds max_context_tokens after a tool call, the agent
+        should immediately synthesize from gathered evidence instead of making
+        another LLM call that would overflow."""
+        # First call: LLM calls recall (forced by iter 0 with no mental models)
+        mock_llm.call_with_tools.return_value = LLMToolCallResult(
+            tool_calls=[LLMToolCall(id="1", name="recall", arguments={"query": "test"})],
+            finish_reason="tool_calls",
+        )
+
+        # Set a tiny token budget — the recall result alone will blow past it
+        result = await run_reflect_agent(
+            llm_config=mock_llm,
+            bank_id="test-bank",
+            query="What do you know?",
+            bank_profile={"name": "Test", "mission": "Testing"},
+            max_context_tokens=100,
+            **mock_functions_with_large_output,
+        )
+
+        assert result.text == "Synthesized answer from gathered evidence."
+        # call_with_tools was called once (for the forced recall), then the guard
+        # kicked in — no further tool-call iterations
+        assert mock_llm.call_with_tools.call_count == 1
+        # llm.call() was invoked to generate the final synthesis
+        mock_llm.call.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_context_overflow_error_skips_retry(self, mock_llm, mock_functions_with_large_output):
+        """A context_length_exceeded error from the LLM should NOT be retried —
+        it should immediately fall back to final synthesis."""
+        mock_llm.call_with_tools.side_effect = Exception(
+            "context_length_exceeded: messages resulted in 150000 tokens."
+        )
+
+        result = await run_reflect_agent(
+            llm_config=mock_llm,
+            bank_id="test-bank",
+            query="What do you know?",
+            bank_profile={"name": "Test", "mission": "Testing"},
+            max_iterations=5,
+            **mock_functions_with_large_output,
+        )
+
+        assert result is not None
+        # Should have attempted only 1 iteration (no retry on overflow error)
+        assert mock_llm.call_with_tools.call_count == 1
+        # Final synthesis was called
+        mock_llm.call.assert_called_once()
+
+
+class TestContextOverflowIntegration:
+    """Integration test: real LLM with a very small max_context_tokens.
+
+    The agent will make one real LLM call (forced tool choice), receive a large
+    tool result that exceeds the tiny budget, then synthesize from it via a second
+    real LLM call — all without raising a context_length_exceeded error.
+    """
+
+    @pytest.mark.asyncio
+    async def test_reflect_completes_with_tiny_context_budget(self, memory, request_context):
+        """End-to-end: reflect on a bank with max_context_tokens=1 (tiny budget).
+
+        Setting max_context_tokens=1 guarantees the proactive guard fires as soon
+        as the first tool result is received and evidence is available.
+        The result must be a non-empty string with no exception raised.
+        """
+        import uuid
+        from unittest.mock import patch
+
+        bank_id = f"test-ctx-overflow-{uuid.uuid4().hex[:8]}"
+        try:
+            # Retain a handful of facts so the recall tool has something to return
+            await memory.retain_async(
+                bank_id=bank_id,
+                content="Alice is a software engineer who enjoys hiking on weekends.",
+                request_context=request_context,
+            )
+            await memory.retain_async(
+                bank_id=bank_id,
+                content="Bob is a designer who loves cooking Italian food.",
+                request_context=request_context,
+            )
+
+            # Patch get_config where memory_engine uses it, injecting a tiny
+            # max_context_tokens.  Everything else delegates to the real config.
+            real_config = memory._get_raw_config() if hasattr(memory, "_get_raw_config") else None
+            from hindsight_api.config import get_config as _real_get_config
+
+            class _TinyContextProxy:
+                """Forwards all attribute access to the real config proxy except
+                reflect_max_context_tokens which is forced to 1."""
+                _real = _real_get_config()
+
+                def __getattr__(self, name: str):
+                    if name == "reflect_max_context_tokens":
+                        return 1
+                    return getattr(self._real, name)
+
+            with patch("hindsight_api.engine.memory_engine.get_config", return_value=_TinyContextProxy()):
+                result = await memory.reflect_async(
+                    bank_id=bank_id,
+                    query="Tell me about the people you know.",
+                    request_context=request_context,
+                )
+
+            assert result.text, "reflect must return a non-empty answer"
+            assert result.usage.total_tokens > 0
+
+        finally:
+            await memory.delete_bank(bank_id, request_context=request_context)
@@ -780,6 +780,7 @@ export HINDSIGHT_API_OBSERVATIONS_MISSION="Observations are recurring patterns i
 | Variable | Description | Default |
 |----------|-------------|---------|
 | `HINDSIGHT_API_REFLECT_MAX_ITERATIONS` | Max tool call iterations before forcing a response | `10` |
+| `HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS` | Max accumulated context tokens in the reflect loop before forcing final synthesis. Prevents `context_length_exceeded` errors on large banks. Lower this if your LLM has a context window smaller than 128K. | `100000` |
 | `HINDSIGHT_API_REFLECT_MISSION` | Global reflect mission (identity and reasoning framing). Overridden per bank via config API. | - |

 #### Disposition