From c7ea8ab546bd412df1d17f77f3312910ff4313e5 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Fri, 28 Aug 2026 17:23:42 +0200
Subject: [PATCH 1/9] fix(agent): stream LLM completions to avoid gateway
 idle-timeout

Corporate LLM gateways (e.g. AI.proxy on AWS) enforce an idle timeout on
buffered (non-streaming) requests, so a long-running compile step can hit
a Gateway Timeout even though the provider would have eventually finished.

Switch _llm_call() and _llm_call_async() in openkb/agent/compiler.py to
litellm.completion()/acompletion() with stream=True: streaming keeps
bytes flowing over the connection, so idle-timeout gateways never see a
silent connection. Chunks are merged back into the existing response
shape via a new _merge_stream_chunks() helper, using LiteLLM's own
litellm.stream_chunk_builder() for genuine multi-chunk streams. An
exception raised mid-stream propagates as a complete failure (list()
never returns a partial buffer), matching prior all-or-nothing behavior.

Adapts the compiler test mocks (_mock_completion/_mock_acompletion and a
handful of inline mocks) to return a single-chunk fake stream, plus the
litellm.completion/acompletion mocks in test_llm_timeout.py.

Resolves #235.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 openkb/agent/compiler.py  |  49 ++++++++++++++++--
 tests/test_compiler.py    | 106 ++++++++++++--------------------------
 tests/test_llm_timeout.py |  19 ++++---
 3 files changed, 91 insertions(+), 83 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index d0c9f878d..2f9df71e1 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -397,6 +397,23 @@ class TruncatedResponseError(Exception):
     treat truncation as a failure (so a partial page is skipped, not written)."""
 
 
+def _merge_stream_chunks(chunks: list, messages: list[dict]):
+    """Merge streamed LLM chunks back into a single, non-streaming response.
+
+    Genuine LiteLLM stream chunks only ever carry a ``.delta`` (never a
+    ``.message``), so a real multi-chunk stream is merged via LiteLLM's own
+    :func:`litellm.stream_chunk_builder`. A single chunk that already looks
+    like a complete, non-streaming ``ModelResponse`` (exposing ``.message``)
+    is used as-is — there's nothing left to merge, and it lets test doubles
+    fake a one-shot response without simulating LiteLLM's internal delta
+    format.
+    """
+    choices = getattr(chunks[0], "choices", None) or []
+    if len(chunks) == 1 and choices and hasattr(choices[0], "message"):
+        return chunks[0]
+    return litellm.stream_chunk_builder(chunks, messages=messages)
+
+
 def _llm_call(
     model: str,
     messages: list[dict],
@@ -406,7 +423,15 @@ def _llm_call(
     bundle=None,
     **kwargs,
 ) -> str:
-    """Single LLM call with animated progress and debug logging."""
+    """Single LLM call with animated progress and debug logging.
+
+    Uses ``stream=True``: some corporate LLM gateways enforce an idle
+    timeout on buffered (non-streaming) requests, which a long-running
+    completion can hit before the response is ever sent. Streaming keeps
+    bytes flowing over the connection so that timeout never fires; the
+    chunks are merged back into a single response via
+    :func:`_merge_stream_chunks` so callers see the same shape as before.
+    """
     messages = _prepare_messages(model, messages)
     extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers()
     if extra_headers:
@@ -417,6 +442,7 @@ def _llm_call(
     if bundle is not None:
         kwargs.setdefault("api_key", bundle.api_key)
         kwargs.setdefault("base_url", bundle.base_url)
+    kwargs.setdefault("stream_options", {"include_usage": True})
     logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages))
     if kwargs:
         logger.debug("LLM kwargs [%s]: %s", step_name, kwargs)
@@ -425,7 +451,11 @@ def _llm_call(
     spinner.start()
     t0 = time.time()
 
-    response = litellm.completion(model=model, messages=messages, **kwargs)
+    stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs)
+    chunks = list(stream)
+    if not chunks:
+        raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
+    response = _merge_stream_chunks(chunks, messages)
     content = response.choices[0].message.content or ""
     truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
 
@@ -449,7 +479,10 @@ async def _llm_call_async(
     bundle=None,
     **kwargs,
 ) -> str:
-    """Async LLM call with timing output and debug logging."""
+    """Async LLM call with timing output and debug logging.
+
+    See ``_llm_call`` for why ``stream=True`` is used.
+    """
     messages = _prepare_messages(model, messages)
     extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers()
     if extra_headers:
@@ -460,13 +493,21 @@ async def _llm_call_async(
     if bundle is not None:
         kwargs.setdefault("api_key", bundle.api_key)
         kwargs.setdefault("base_url", bundle.base_url)
+    kwargs.setdefault("stream_options", {"include_usage": True})
     logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages))
     if kwargs:
         logger.debug("LLM kwargs [%s]: %s", step_name, kwargs)
 
     t0 = time.time()
 
-    response = await litellm.acompletion(model=model, messages=messages, **kwargs)
+    stream = await litellm.acompletion(model=model, messages=messages, stream=True, **kwargs)
+    if hasattr(stream, "__aiter__"):
+        chunks = [chunk async for chunk in stream]
+    else:
+        chunks = list(stream)
+    if not chunks:
+        raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
+    response = _merge_stream_chunks(chunks, messages)
     content = response.choices[0].message.content or ""
     truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
 
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 95a57cc4c..e17ad47d7 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -1112,36 +1112,45 @@ def test_frontmatter_without_sources_line_gets_one_inserted(self, tmp_path):
         assert "[[summaries/new-doc]]" in text
 
 
+def _mock_response(content, finish_reason: str = "stop") -> MagicMock:
+    """Build a fake, already-complete LLM response (single-chunk stream).
+
+    ``_llm_call``/``_llm_call_async`` now call ``litellm.completion``/
+    ``acompletion`` with ``stream=True`` and merge the resulting chunks back
+    into one response (see ``_merge_stream_chunks``). Exposing ``.message``
+    (rather than the ``.delta`` a genuine stream chunk carries) tells
+    ``_merge_stream_chunks`` this single chunk *is* the final response, so it
+    is used as-is without needing to fake LiteLLM's internal delta format.
+    """
+    mock_resp = MagicMock()
+    mock_resp.choices = [MagicMock()]
+    mock_resp.choices[0].message.content = content
+    mock_resp.choices[0].finish_reason = finish_reason
+    mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
+    mock_resp.usage.prompt_tokens_details = None
+    return mock_resp
+
+
 def _mock_completion(responses: list[str]):
-    """Create a mock for litellm.completion that returns responses in order."""
+    """Create a mock for litellm.completion returning a single-chunk stream."""
     call_count = {"n": 0}
 
     def side_effect(*args, **kwargs):
         idx = min(call_count["n"], len(responses) - 1)
         call_count["n"] += 1
-        mock_resp = MagicMock()
-        mock_resp.choices = [MagicMock()]
-        mock_resp.choices[0].message.content = responses[idx]
-        mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
-        mock_resp.usage.prompt_tokens_details = None
-        return mock_resp
+        return [_mock_response(responses[idx])]
 
     return side_effect
 
 
 def _mock_acompletion(responses: list[str]):
-    """Create an async mock for litellm.acompletion."""
+    """Create an async mock for litellm.acompletion returning a single-chunk stream."""
     call_count = {"n": 0}
 
     async def side_effect(*args, **kwargs):
         idx = min(call_count["n"], len(responses) - 1)
         call_count["n"] += 1
-        mock_resp = MagicMock()
-        mock_resp.choices = [MagicMock()]
-        mock_resp.choices[0].message.content = responses[idx]
-        mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
-        mock_resp.usage.prompt_tokens_details = None
-        return mock_resp
+        return [_mock_response(responses[idx])]
 
     return side_effect
 
@@ -1342,15 +1351,7 @@ def sync_side_effect(*args, **kwargs):
             sync_call_count["n"] += 1
             if idx == 2:  # the summary-rewrite call
                 raise RuntimeError("simulated API failure")
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
-            mock_resp.choices[0].message.content = [
-                summary_response,
-                plan_response,
-            ][idx]
-            mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            return [_mock_response([summary_response, plan_response][idx])]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=sync_side_effect)
@@ -1507,21 +1508,11 @@ async def test_short_doc_marks_doc_and_summary(self, tmp_path):
         def sync_side_effect(*args, **kwargs):
             captured_sync_calls.append(kwargs["messages"])
             idx = min(len(captured_sync_calls) - 1, len(sync_responses) - 1)
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
-            mock_resp.choices[0].message.content = sync_responses[idx]
-            mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            return [_mock_response(sync_responses[idx])]
 
         async def async_side_effect(*args, **kwargs):
             captured_async_calls.append(kwargs["messages"])
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
-            mock_resp.choices[0].message.content = concept_response
-            mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            return [_mock_response(concept_response)]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=sync_side_effect)
@@ -1586,15 +1577,9 @@ async def test_long_doc_marks_doc_message(self, tmp_path):
 
         def sync_side_effect(*args, **kwargs):
             captured.append(kwargs["messages"])
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
             # First call: overview (plain text); second: plan (JSON).
-            mock_resp.choices[0].message.content = (
-                "Overview text" if len(captured) == 1 else plan_response
-            )
-            mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            content = "Overview text" if len(captured) == 1 else plan_response
+            return [_mock_response(content)]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=sync_side_effect)
@@ -1726,16 +1711,9 @@ async def test_create_and_update_flow(self, tmp_path):
         async def ordered_acompletion(*args, **kwargs):
             idx = call_order["n"]
             call_order["n"] += 1
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
             # create tasks come first, then update tasks
-            if idx == 0:
-                mock_resp.choices[0].message.content = create_page_response
-            else:
-                mock_resp.choices[0].message.content = update_page_response
-            mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            content = create_page_response if idx == 0 else update_page_response
+            return [_mock_response(content)]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response]))
@@ -1823,13 +1801,7 @@ async def test_truncated_update_preserves_existing_page(self, tmp_path):
         )
 
         async def truncated_acompletion(*args, **kwargs):
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
-            mock_resp.choices[0].message.content = truncated_page
-            mock_resp.choices[0].finish_reason = "length"
-            mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            return [_mock_response(truncated_page, finish_reason="length")]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response]))
@@ -1859,13 +1831,7 @@ async def test_truncated_create_skips_partial_page(self, tmp_path):
         truncated_page = json.dumps({"brief": "x", "content": "# Ghost\n\nPartial"})
 
         async def truncated_acompletion(*args, **kwargs):
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
-            mock_resp.choices[0].message.content = truncated_page
-            mock_resp.choices[0].finish_reason = "length"
-            mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            return [_mock_response(truncated_page, finish_reason="length")]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response]))
@@ -1928,13 +1894,7 @@ async def test_truncated_entity_update_preserves_existing_page(self, tmp_path):
         )
 
         async def truncated_acompletion(*args, **kwargs):
-            mock_resp = MagicMock()
-            mock_resp.choices = [MagicMock()]
-            mock_resp.choices[0].message.content = truncated_page
-            mock_resp.choices[0].finish_reason = "length"
-            mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
-            mock_resp.usage.prompt_tokens_details = None
-            return mock_resp
+            return [_mock_response(truncated_page, finish_reason="length")]
 
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
             mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response]))
diff --git a/tests/test_llm_timeout.py b/tests/test_llm_timeout.py
index ca7d80e68..db119df3c 100644
--- a/tests/test_llm_timeout.py
+++ b/tests/test_llm_timeout.py
@@ -17,6 +17,13 @@
 
 
 def _fake_response():
+    """A fake, already-complete LLM response (single-chunk stream).
+
+    See ``openkb.agent.compiler._merge_stream_chunks``: a chunk exposing
+    ``.message`` (as this one does) is treated as already-complete and used
+    as-is, so callers of ``litellm.completion``/``acompletion`` with
+    ``stream=True`` can be mocked to just return a one-item list.
+    """
     choice = MagicMock()
     choice.message.content = "ok"
     choice.finish_reason = "stop"
@@ -28,7 +35,7 @@ def _fake_response():
 def test_llm_call_forwards_configured_timeout():
     set_timeout(1200.0)
     with patch(
-        "openkb.agent.compiler.litellm.completion", return_value=_fake_response()
+        "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()]
     ) as completion:
         _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step")
     assert completion.call_args.kwargs["timeout"] == 1200.0
@@ -37,7 +44,7 @@ def test_llm_call_forwards_configured_timeout():
 def test_llm_call_omits_timeout_when_unset():
     set_timeout(None)
     with patch(
-        "openkb.agent.compiler.litellm.completion", return_value=_fake_response()
+        "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()]
     ) as completion:
         _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step")
     assert "timeout" not in completion.call_args.kwargs
@@ -47,7 +54,7 @@ def test_llm_call_does_not_override_explicit_timeout():
     # An explicit per-call timeout kwarg wins over the configured default.
     set_timeout(1200.0)
     with patch(
-        "openkb.agent.compiler.litellm.completion", return_value=_fake_response()
+        "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()]
     ) as completion:
         _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30)
     assert completion.call_args.kwargs["timeout"] == 30
@@ -58,7 +65,7 @@ def test_llm_call_async_forwards_configured_timeout():
     with patch(
         "openkb.agent.compiler.litellm.acompletion",
         new_callable=AsyncMock,
-        return_value=_fake_response(),
+        return_value=[_fake_response()],
     ) as acompletion:
         asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step"))
     assert acompletion.call_args.kwargs["timeout"] == 900.0
@@ -69,7 +76,7 @@ def test_llm_call_async_omits_timeout_when_unset():
     with patch(
         "openkb.agent.compiler.litellm.acompletion",
         new_callable=AsyncMock,
-        return_value=_fake_response(),
+        return_value=[_fake_response()],
     ) as acompletion:
         asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step"))
     assert "timeout" not in acompletion.call_args.kwargs
@@ -80,7 +87,7 @@ def test_llm_call_async_does_not_override_explicit_timeout():
     with patch(
         "openkb.agent.compiler.litellm.acompletion",
         new_callable=AsyncMock,
-        return_value=_fake_response(),
+        return_value=[_fake_response()],
     ) as acompletion:
         asyncio.run(
             _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30)

From 76905b8c90c4c006ff356e4b1e7a74b86a826dfc Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Sat, 29 Aug 2026 16:07:41 +0200
Subject: [PATCH 2/9] fix(agent): add debug timing for streamed completions

Add per-chunk debug logging around streamed LiteLLM completion consumption in `openkb.agent.compiler`.

This diagnostic instrumentation is enabled via the existing `openkb -v` flag and helps narrow the still-observed ~60s production cutoff to either a no-first-byte case (for example slow time-to-first-token) or a proxy/gateway path that buffers or drops streamed bytes.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 openkb/agent/compiler.py |  86 ++++++++++++++++++++-
 tests/test_compiler.py   | 162 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 245 insertions(+), 3 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 2f9df71e1..311ce4f16 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -414,6 +414,86 @@ def _merge_stream_chunks(chunks: list, messages: list[dict]):
     return litellm.stream_chunk_builder(chunks, messages=messages)
 
 
+def _log_chunk_timing(
+    step_name: str, chunk_number: int, t0: float, last_t: float, now: float
+) -> None:
+    """Debug-log arrival timing for one streamed chunk.
+
+    Run with ``openkb -v`` to see, per LLM call, how long the first chunk took
+    (time-to-first-token) and the gap to each subsequent chunk. If chunks do
+    arrive before a timeout, that shows the proxy/gateway is forwarding the
+    stream; if no chunk is ever logged before a timeout, the instrumentation
+    narrows the problem to a no-first-byte case (e.g. slow TTFT or an
+    intermediary buffering the response).
+    """
+    logger.debug(
+        "LLM stream chunk [%s] #%d after %.2fs total (+%.2fs since previous)",
+        step_name,
+        chunk_number,
+        now - t0,
+        now - last_t,
+    )
+
+
+def _consume_stream(stream, step_name: str, t0: float) -> list:
+    """Collect a sync LiteLLM stream into a list, logging per-chunk timing.
+
+    A mid-stream exception (e.g. the gateway idle-timeout firing) propagates
+    after logging how many chunks arrived and when, so callers still see a
+    complete failure — no partial buffer is ever returned.
+    """
+    if not logger.isEnabledFor(logging.DEBUG):
+        return list(stream)
+
+    chunks: list = []
+    last_t = t0
+    try:
+        for chunk in stream:
+            now = time.time()
+            _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now)
+            chunks.append(chunk)
+            last_t = now
+    except Exception:
+        logger.debug(
+            "LLM stream [%s] failed after %.2fs with %d chunk(s) received",
+            step_name,
+            time.time() - t0,
+            len(chunks),
+            exc_info=True,
+        )
+        raise
+    return chunks
+
+
+async def _consume_stream_async(stream, step_name: str, t0: float) -> list:
+    """Collect an async LiteLLM stream into a list, logging per-chunk timing.
+
+    Mirrors :func:`_consume_stream`, including the no-partial-buffer invariant
+    on failure.
+    """
+    if not logger.isEnabledFor(logging.DEBUG):
+        return [chunk async for chunk in stream]
+
+    chunks: list = []
+    last_t = t0
+    try:
+        async for chunk in stream:
+            now = time.time()
+            _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now)
+            chunks.append(chunk)
+            last_t = now
+    except Exception:
+        logger.debug(
+            "LLM stream [%s] failed after %.2fs with %d chunk(s) received",
+            step_name,
+            time.time() - t0,
+            len(chunks),
+            exc_info=True,
+        )
+        raise
+    return chunks
+
+
 def _llm_call(
     model: str,
     messages: list[dict],
@@ -452,7 +532,7 @@ def _llm_call(
     t0 = time.time()
 
     stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs)
-    chunks = list(stream)
+    chunks = _consume_stream(stream, step_name, t0)
     if not chunks:
         raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
     response = _merge_stream_chunks(chunks, messages)
@@ -502,9 +582,9 @@ async def _llm_call_async(
 
     stream = await litellm.acompletion(model=model, messages=messages, stream=True, **kwargs)
     if hasattr(stream, "__aiter__"):
-        chunks = [chunk async for chunk in stream]
+        chunks = await _consume_stream_async(stream, step_name, t0)
     else:
-        chunks = list(stream)
+        chunks = _consume_stream(stream, step_name, t0)
     if not chunks:
         raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
     response = _merge_stream_chunks(chunks, messages)
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index e17ad47d7..a5119ff2a 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import logging
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -1155,6 +1156,53 @@ async def side_effect(*args, **kwargs):
     return side_effect
 
 
+class _NoOpSpinner:
+    """Test double that disables spinner side effects."""
+
+    def __init__(self, *_args, **_kwargs):
+        pass
+
+    def start(self) -> None:
+        pass
+
+    def stop(self, _suffix: str = "") -> None:
+        pass
+
+
+class _AsyncStream:
+    """Simple async iterator for exercising streamed LiteLLM responses in tests."""
+
+    def __init__(
+        self,
+        chunks: list[object],
+        *,
+        error: Exception | None = None,
+        raise_after: int | None = None,
+    ) -> None:
+        self._chunks = chunks
+        self._error = error
+        self._raise_after = raise_after
+        self._index = 0
+
+    def __aiter__(self) -> _AsyncStream:
+        return self
+
+    async def __anext__(self) -> object:
+        if self._raise_after is not None and self._index == self._raise_after:
+            raise self._error or RuntimeError("stream exploded")
+        if self._index >= len(self._chunks):
+            raise StopAsyncIteration
+        chunk = self._chunks[self._index]
+        self._index += 1
+        return chunk
+
+
+def _stream_then_raise(chunks: list[object], error: Exception):
+    """Yield all chunks, then raise ``error`` on the next iteration."""
+    yield from chunks
+    raise error
+
+
 class TestCompileShortDoc:
     @pytest.mark.asyncio
     async def test_full_pipeline(self, tmp_path):
@@ -2661,6 +2709,120 @@ async def test_llm_call_async_injects_extra_headers(self):
         assert kwargs["extra_headers"] == {"Copilot-Integration-Id": "vscode-chat"}
 
 
+class TestLLMStreamTimingDebugLogging:
+    """Per-chunk debug timing should be visible when verbose logging is enabled."""
+
+    def test_llm_call_logs_each_stream_chunk_at_debug(self, caplog):
+        from openkb.agent.compiler import _llm_call
+
+        caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
+        with (
+            patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.completion = MagicMock(
+                return_value=iter(["chunk-1", "chunk-2", "chunk-3"])
+            )
+            mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok"))
+
+            out = _llm_call("m", [{"role": "user", "content": "hi"}], "sync-step")
+
+        assert out == "ok"
+        chunk_logs = [
+            record.getMessage()
+            for record in caplog.records
+            if "LLM stream chunk [sync-step]" in record.getMessage()
+        ]
+        assert len(chunk_logs) == 3
+        assert "#1" in chunk_logs[0]
+        assert "#3" in chunk_logs[-1]
+
+    @pytest.mark.asyncio
+    async def test_llm_call_async_logs_each_stream_chunk_at_debug(self, caplog):
+        from openkb.agent.compiler import _llm_call_async
+
+        caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.acompletion = AsyncMock(
+                return_value=_AsyncStream(["chunk-1", "chunk-2", "chunk-3"])
+            )
+            mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok"))
+
+            out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-step")
+
+        assert out == "ok"
+        chunk_logs = [
+            record.getMessage()
+            for record in caplog.records
+            if "LLM stream chunk [async-step]" in record.getMessage()
+        ]
+        assert len(chunk_logs) == 3
+        assert "#1" in chunk_logs[0]
+        assert "#3" in chunk_logs[-1]
+
+    def test_llm_call_logs_stream_failure_before_reraising(self, caplog):
+        from openkb.agent.compiler import _llm_call
+
+        caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
+        error = RuntimeError("stream exploded")
+        with (
+            patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.completion = MagicMock(
+                return_value=_stream_then_raise(["chunk-1", "chunk-2"], error)
+            )
+
+            with pytest.raises(RuntimeError, match="stream exploded"):
+                _llm_call("m", [{"role": "user", "content": "hi"}], "sync-fail-step")
+
+        chunk_logs = [
+            record.getMessage()
+            for record in caplog.records
+            if "LLM stream chunk [sync-fail-step]" in record.getMessage()
+        ]
+        failure_logs = [
+            record.getMessage()
+            for record in caplog.records
+            if "LLM stream [sync-fail-step] failed after" in record.getMessage()
+        ]
+        assert len(chunk_logs) == 2
+        assert len(failure_logs) == 1
+        assert "2 chunk(s) received" in failure_logs[0]
+
+    @pytest.mark.asyncio
+    async def test_llm_call_async_logs_stream_failure_before_reraising(self, caplog):
+        from openkb.agent.compiler import _llm_call_async
+
+        caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
+        error = RuntimeError("stream exploded")
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.acompletion = AsyncMock(
+                return_value=_AsyncStream(
+                    ["chunk-1", "chunk-2"],
+                    error=error,
+                    raise_after=2,
+                )
+            )
+
+            with pytest.raises(RuntimeError, match="stream exploded"):
+                await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-fail-step")
+
+        chunk_logs = [
+            record.getMessage()
+            for record in caplog.records
+            if "LLM stream chunk [async-fail-step]" in record.getMessage()
+        ]
+        failure_logs = [
+            record.getMessage()
+            for record in caplog.records
+            if "LLM stream [async-fail-step] failed after" in record.getMessage()
+        ]
+        assert len(chunk_logs) == 2
+        assert len(failure_logs) == 1
+        assert "2 chunk(s) received" in failure_logs[0]
+
+
 class TestCacheControlStripping:
     """cache_control markers must only reach providers that honour them.
 

From 6c613b0d1baddae4f1dacb3de6b18f2caaf98369 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Mon, 31 Aug 2026 08:07:33 +0200
Subject: [PATCH 3/9] fix(agent): log stream chunk phase as start/end, not one
 line per chunk
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per-chunk DEBUG logging (_log_chunk_timing) drowned out the rest of
a log on a long response — hundreds of LLM stream chunk [...] #N
lines for a single LLM call, e.g. one concepts-plan request logging
205 individual chunk lines.

Replaces it with exactly two log lines per LLM call:
- _log_stream_start: logged once, when the first chunk arrives
  (time-to-first-token).
- _log_stream_end: logged once, when the stream finishes cleanly
  (total chunk count + elapsed time for the last chunk).
- _log_stream_interrupted: logged once instead of _log_stream_end
  if the stream raises mid-iteration — reports how many chunks were
  successfully received and when, right before the exception is
  re-raised (still a complete failure, no partial buffer). Special-cases
  zero chunks (failure before any byte arrived) with dedicated wording
  instead of an inapplicable chunk number.

Updated tests/test_compiler.py::TestLLMStreamTimingDebugLogging to
match: asserts exactly one start + one end/interrupted line, and the
explicit absence of the old per-chunk lines.
---
 openkb/agent/compiler.py | 105 +++++++++++++++++++++++-------------
 tests/test_compiler.py   | 112 +++++++++++++++++++++++----------------
 2 files changed, 135 insertions(+), 82 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 311ce4f16..756f66ec7 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -414,33 +414,72 @@ def _merge_stream_chunks(chunks: list, messages: list[dict]):
     return litellm.stream_chunk_builder(chunks, messages=messages)
 
 
-def _log_chunk_timing(
-    step_name: str, chunk_number: int, t0: float, last_t: float, now: float
+def _log_stream_start(step_name: str, t0: float, first_chunk_t: float) -> None:
+    """Debug-log the time-to-first-chunk (TTFT) once a stream's first chunk arrives.
+
+    Marks the start of a "chunk phase" in the log. The counterpart is
+    :func:`_log_stream_end` (clean finish) or :func:`_log_stream_interrupted`
+    (mid-stream failure) — together these replace a debug line per chunk
+    (which used to drown out the rest of the log on a long response, e.g.
+    hundreds of lines for one LLM call) with exactly one line at the start
+    and exactly one more at the end/interruption.
+    """
+    logger.debug(
+        "LLM stream started [%s]: first chunk after %.2fs",
+        step_name,
+        first_chunk_t - t0,
+    )
+
+
+def _log_stream_end(step_name: str, chunk_count: int, t0: float, last_chunk_t: float) -> None:
+    """Debug-log a stream's clean completion: total chunk count and elapsed time."""
+    logger.debug(
+        "LLM stream finished [%s]: %d chunk(s), last chunk after %.2fs total",
+        step_name,
+        chunk_count,
+        last_chunk_t - t0,
+    )
+
+
+def _log_stream_interrupted(
+    step_name: str, chunk_count: int, t0: float, last_chunk_t: float
 ) -> None:
-    """Debug-log arrival timing for one streamed chunk.
-
-    Run with ``openkb -v`` to see, per LLM call, how long the first chunk took
-    (time-to-first-token) and the gap to each subsequent chunk. If chunks do
-    arrive before a timeout, that shows the proxy/gateway is forwarding the
-    stream; if no chunk is ever logged before a timeout, the instrumentation
-    narrows the problem to a no-first-byte case (e.g. slow TTFT or an
-    intermediary buffering the response).
+    """Debug-log a stream that raised mid-iteration, right before it is re-raised.
+
+    ``chunk_count`` is how many chunks were successfully received before the
+    failure (0 if the very first chunk never arrived). The exception itself
+    (with traceback) is attached via ``exc_info=True`` so the failure and the
+    chunk-phase summary land in a single log record.
     """
+    now = time.time()
+    if chunk_count == 0:
+        logger.debug(
+            "LLM stream [%s] interrupted unexpectedly before any chunk arrived (%.2fs total)",
+            step_name,
+            now - t0,
+            exc_info=True,
+        )
+        return
     logger.debug(
-        "LLM stream chunk [%s] #%d after %.2fs total (+%.2fs since previous)",
+        "LLM stream [%s] interrupted unexpectedly after chunk %d "
+        "(last chunk after %.2fs, failure after %.2fs total)",
         step_name,
-        chunk_number,
+        chunk_count,
+        last_chunk_t - t0,
         now - t0,
-        now - last_t,
+        exc_info=True,
     )
 
 
 def _consume_stream(stream, step_name: str, t0: float) -> list:
-    """Collect a sync LiteLLM stream into a list, logging per-chunk timing.
-
-    A mid-stream exception (e.g. the gateway idle-timeout firing) propagates
-    after logging how many chunks arrived and when, so callers still see a
-    complete failure — no partial buffer is ever returned.
+    """Collect a sync LiteLLM stream into a list, debug-logging the chunk phase.
+
+    Logs exactly one line when the first chunk arrives (time-to-first-token)
+    and exactly one more line when the stream ends — either
+    :func:`_log_stream_end` on a clean finish or :func:`_log_stream_interrupted`
+    if it raises mid-iteration. A mid-stream exception (e.g. the gateway
+    idle-timeout firing) propagates after being logged, so callers still see
+    a complete failure — no partial buffer is ever returned.
     """
     if not logger.isEnabledFor(logging.DEBUG):
         return list(stream)
@@ -450,26 +489,22 @@ def _consume_stream(stream, step_name: str, t0: float) -> list:
     try:
         for chunk in stream:
             now = time.time()
-            _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now)
+            if not chunks:
+                _log_stream_start(step_name, t0, now)
             chunks.append(chunk)
             last_t = now
     except Exception:
-        logger.debug(
-            "LLM stream [%s] failed after %.2fs with %d chunk(s) received",
-            step_name,
-            time.time() - t0,
-            len(chunks),
-            exc_info=True,
-        )
+        _log_stream_interrupted(step_name, len(chunks), t0, last_t)
         raise
+    _log_stream_end(step_name, len(chunks), t0, last_t)
     return chunks
 
 
 async def _consume_stream_async(stream, step_name: str, t0: float) -> list:
-    """Collect an async LiteLLM stream into a list, logging per-chunk timing.
+    """Collect an async LiteLLM stream into a list, debug-logging the chunk phase.
 
-    Mirrors :func:`_consume_stream`, including the no-partial-buffer invariant
-    on failure.
+    Mirrors :func:`_consume_stream`, including the start/end-or-interrupted
+    logging and the no-partial-buffer invariant on failure.
     """
     if not logger.isEnabledFor(logging.DEBUG):
         return [chunk async for chunk in stream]
@@ -479,18 +514,14 @@ async def _consume_stream_async(stream, step_name: str, t0: float) -> list:
     try:
         async for chunk in stream:
             now = time.time()
-            _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now)
+            if not chunks:
+                _log_stream_start(step_name, t0, now)
             chunks.append(chunk)
             last_t = now
     except Exception:
-        logger.debug(
-            "LLM stream [%s] failed after %.2fs with %d chunk(s) received",
-            step_name,
-            time.time() - t0,
-            len(chunks),
-            exc_info=True,
-        )
+        _log_stream_interrupted(step_name, len(chunks), t0, last_t)
         raise
+    _log_stream_end(step_name, len(chunks), t0, last_t)
     return chunks
 
 
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index a5119ff2a..34a94ac6c 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -2710,9 +2710,11 @@ async def test_llm_call_async_injects_extra_headers(self):
 
 
 class TestLLMStreamTimingDebugLogging:
-    """Per-chunk debug timing should be visible when verbose logging is enabled."""
+    """Chunk-phase debug logging should be visible when verbose logging is
+    enabled: one line when the first chunk arrives, one more when the stream
+    ends cleanly or is interrupted — never one line per chunk (see #<TODO>)."""
 
-    def test_llm_call_logs_each_stream_chunk_at_debug(self, caplog):
+    def test_llm_call_logs_stream_start_and_end_at_debug(self, caplog):
         from openkb.agent.compiler import _llm_call
 
         caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
@@ -2728,17 +2730,17 @@ def test_llm_call_logs_each_stream_chunk_at_debug(self, caplog):
             out = _llm_call("m", [{"role": "user", "content": "hi"}], "sync-step")
 
         assert out == "ok"
-        chunk_logs = [
-            record.getMessage()
-            for record in caplog.records
-            if "LLM stream chunk [sync-step]" in record.getMessage()
-        ]
-        assert len(chunk_logs) == 3
-        assert "#1" in chunk_logs[0]
-        assert "#3" in chunk_logs[-1]
+        messages = [record.getMessage() for record in caplog.records]
+        start_logs = [m for m in messages if "LLM stream started [sync-step]" in m]
+        end_logs = [m for m in messages if "LLM stream finished [sync-step]" in m]
+        # Exactly one start line and one end line — never a line per chunk.
+        assert len(start_logs) == 1
+        assert len(end_logs) == 1
+        assert "3 chunk(s)" in end_logs[0]
+        assert not any("LLM stream chunk [sync-step]" in m for m in messages)
 
     @pytest.mark.asyncio
-    async def test_llm_call_async_logs_each_stream_chunk_at_debug(self, caplog):
+    async def test_llm_call_async_logs_stream_start_and_end_at_debug(self, caplog):
         from openkb.agent.compiler import _llm_call_async
 
         caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
@@ -2751,16 +2753,15 @@ async def test_llm_call_async_logs_each_stream_chunk_at_debug(self, caplog):
             out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-step")
 
         assert out == "ok"
-        chunk_logs = [
-            record.getMessage()
-            for record in caplog.records
-            if "LLM stream chunk [async-step]" in record.getMessage()
-        ]
-        assert len(chunk_logs) == 3
-        assert "#1" in chunk_logs[0]
-        assert "#3" in chunk_logs[-1]
-
-    def test_llm_call_logs_stream_failure_before_reraising(self, caplog):
+        messages = [record.getMessage() for record in caplog.records]
+        start_logs = [m for m in messages if "LLM stream started [async-step]" in m]
+        end_logs = [m for m in messages if "LLM stream finished [async-step]" in m]
+        assert len(start_logs) == 1
+        assert len(end_logs) == 1
+        assert "3 chunk(s)" in end_logs[0]
+        assert not any("LLM stream chunk [async-step]" in m for m in messages)
+
+    def test_llm_call_logs_interruption_before_reraising(self, caplog):
         from openkb.agent.compiler import _llm_call
 
         caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
@@ -2776,22 +2777,21 @@ def test_llm_call_logs_stream_failure_before_reraising(self, caplog):
             with pytest.raises(RuntimeError, match="stream exploded"):
                 _llm_call("m", [{"role": "user", "content": "hi"}], "sync-fail-step")
 
-        chunk_logs = [
-            record.getMessage()
-            for record in caplog.records
-            if "LLM stream chunk [sync-fail-step]" in record.getMessage()
-        ]
-        failure_logs = [
-            record.getMessage()
-            for record in caplog.records
-            if "LLM stream [sync-fail-step] failed after" in record.getMessage()
+        messages = [record.getMessage() for record in caplog.records]
+        start_logs = [m for m in messages if "LLM stream started [sync-fail-step]" in m]
+        interrupted_logs = [
+            m for m in messages if "LLM stream [sync-fail-step] interrupted unexpectedly" in m
         ]
-        assert len(chunk_logs) == 2
-        assert len(failure_logs) == 1
-        assert "2 chunk(s) received" in failure_logs[0]
+        # Exactly one start line and one interruption line, no end-of-stream line,
+        # and no per-chunk lines in between.
+        assert len(start_logs) == 1
+        assert len(interrupted_logs) == 1
+        assert "after chunk 2" in interrupted_logs[0]
+        assert not any("LLM stream finished [sync-fail-step]" in m for m in messages)
+        assert not any("LLM stream chunk [sync-fail-step]" in m for m in messages)
 
     @pytest.mark.asyncio
-    async def test_llm_call_async_logs_stream_failure_before_reraising(self, caplog):
+    async def test_llm_call_async_logs_interruption_before_reraising(self, caplog):
         from openkb.agent.compiler import _llm_call_async
 
         caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
@@ -2808,19 +2808,41 @@ async def test_llm_call_async_logs_stream_failure_before_reraising(self, caplog)
             with pytest.raises(RuntimeError, match="stream exploded"):
                 await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-fail-step")
 
-        chunk_logs = [
-            record.getMessage()
-            for record in caplog.records
-            if "LLM stream chunk [async-fail-step]" in record.getMessage()
+        messages = [record.getMessage() for record in caplog.records]
+        start_logs = [m for m in messages if "LLM stream started [async-fail-step]" in m]
+        interrupted_logs = [
+            m for m in messages if "LLM stream [async-fail-step] interrupted unexpectedly" in m
         ]
-        failure_logs = [
-            record.getMessage()
-            for record in caplog.records
-            if "LLM stream [async-fail-step] failed after" in record.getMessage()
+        assert len(start_logs) == 1
+        assert len(interrupted_logs) == 1
+        assert "after chunk 2" in interrupted_logs[0]
+        assert not any("LLM stream finished [async-fail-step]" in m for m in messages)
+        assert not any("LLM stream chunk [async-fail-step]" in m for m in messages)
+
+    def test_llm_call_logs_interruption_before_any_chunk(self, caplog):
+        """No chunk ever arrives (e.g. a proxy silently buffering despite
+        stream=True): no start line, and the interruption line says so
+        instead of an inapplicable chunk number."""
+        from openkb.agent.compiler import _llm_call
+
+        caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
+        error = RuntimeError("connect timeout")
+        with (
+            patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.completion = MagicMock(return_value=_stream_then_raise([], error))
+
+            with pytest.raises(RuntimeError, match="connect timeout"):
+                _llm_call("m", [{"role": "user", "content": "hi"}], "sync-no-chunk-step")
+
+        messages = [record.getMessage() for record in caplog.records]
+        interrupted_logs = [
+            m for m in messages if "LLM stream [sync-no-chunk-step] interrupted unexpectedly" in m
         ]
-        assert len(chunk_logs) == 2
-        assert len(failure_logs) == 1
-        assert "2 chunk(s) received" in failure_logs[0]
+        assert len(interrupted_logs) == 1
+        assert "before any chunk arrived" in interrupted_logs[0]
+        assert not any("LLM stream started [sync-no-chunk-step]" in m for m in messages)
 
 
 class TestCacheControlStripping:

From 5cd479915ddcb1b1b87782e8188159c61a701b31 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Fri, 4 Sep 2026 11:06:49 +0200
Subject: [PATCH 4/9] fix(agent): retry transient stream errors up to 2 extra
 times

_llm_call/_llm_call_async now wrap the stream fetch/consume/merge step in a
fixed 3-attempt loop. A dropped connection or other transient error during
the now-streamed completion (#236) previously failed the whole concept/entity
generation immediately; it now gets 2 automatic retries before giving up.

Not a tunable knob, just a resilience floor sitting below the end-of-batch
sweep retry added for insert_mode (#241).
---
 openkb/agent/compiler.py | 63 +++++++++++++++++++++++++++++++---------
 tests/test_compiler.py   | 36 ++++++++++++++---------
 2 files changed, 73 insertions(+), 26 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 756f66ec7..420f0bcc8 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -562,11 +562,29 @@ def _llm_call(
     spinner.start()
     t0 = time.time()
 
-    stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs)
-    chunks = _consume_stream(stream, step_name, t0)
-    if not chunks:
-        raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
-    response = _merge_stream_chunks(chunks, messages)
+    # Fixed 2 extra attempts for transient stream/LLM errors — not a tunable
+    # knob, just a resilience floor. The concept/entity sweep in
+    # _compile_concepts is the next retry tier above this one.
+    attempts = 3
+    for attempt in range(attempts):
+        try:
+            stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs)
+            chunks = _consume_stream(stream, step_name, t0)
+            if not chunks:
+                raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
+            response = _merge_stream_chunks(chunks, messages)
+            break
+        except Exception as exc:
+            if attempt == attempts - 1:
+                spinner.stop("failed")
+                raise
+            logger.warning(
+                "LLM [%s] attempt %d/%d failed: %s; retrying...",
+                step_name,
+                attempt + 1,
+                attempts,
+                exc,
+            )
     content = response.choices[0].message.content or ""
     truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
 
@@ -611,14 +629,33 @@ async def _llm_call_async(
 
     t0 = time.time()
 
-    stream = await litellm.acompletion(model=model, messages=messages, stream=True, **kwargs)
-    if hasattr(stream, "__aiter__"):
-        chunks = await _consume_stream_async(stream, step_name, t0)
-    else:
-        chunks = _consume_stream(stream, step_name, t0)
-    if not chunks:
-        raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
-    response = _merge_stream_chunks(chunks, messages)
+    # Fixed 2 extra attempts for transient stream/LLM errors — not a tunable
+    # knob, just a resilience floor. The concept/entity sweep in
+    # _compile_concepts is the next retry tier above this one.
+    attempts = 3
+    for attempt in range(attempts):
+        try:
+            stream = await litellm.acompletion(
+                model=model, messages=messages, stream=True, **kwargs
+            )
+            if hasattr(stream, "__aiter__"):
+                chunks = await _consume_stream_async(stream, step_name, t0)
+            else:
+                chunks = _consume_stream(stream, step_name, t0)
+            if not chunks:
+                raise RuntimeError(f"LLM [{step_name}] stream produced no chunks")
+            response = _merge_stream_chunks(chunks, messages)
+            break
+        except Exception as exc:
+            if attempt == attempts - 1:
+                raise
+            logger.warning(
+                "LLM [%s] attempt %d/%d failed: %s; retrying...",
+                step_name,
+                attempt + 1,
+                attempts,
+                exc,
+            )
     content = response.choices[0].message.content or ""
     truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens"))
 
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 34a94ac6c..665827295 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -2770,8 +2770,12 @@ def test_llm_call_logs_interruption_before_reraising(self, caplog):
             patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
             patch("openkb.agent.compiler.litellm") as mock_litellm,
         ):
+            # Fresh generator per call: _llm_call now retries 3 times total
+            # (fixed resilience floor), and a real litellm.completion() call
+            # returns an independent stream every time, unlike reusing one
+            # already-exhausted mock generator across attempts.
             mock_litellm.completion = MagicMock(
-                return_value=_stream_then_raise(["chunk-1", "chunk-2"], error)
+                side_effect=lambda *a, **k: _stream_then_raise(["chunk-1", "chunk-2"], error)
             )
 
             with pytest.raises(RuntimeError, match="stream exploded"):
@@ -2782,11 +2786,11 @@ def test_llm_call_logs_interruption_before_reraising(self, caplog):
         interrupted_logs = [
             m for m in messages if "LLM stream [sync-fail-step] interrupted unexpectedly" in m
         ]
-        # Exactly one start line and one interruption line, no end-of-stream line,
-        # and no per-chunk lines in between.
-        assert len(start_logs) == 1
-        assert len(interrupted_logs) == 1
-        assert "after chunk 2" in interrupted_logs[0]
+        # One start + one interruption line per attempt (3 attempts total),
+        # no end-of-stream line, and no per-chunk lines in between.
+        assert len(start_logs) == 3
+        assert len(interrupted_logs) == 3
+        assert all("after chunk 2" in m for m in interrupted_logs)
         assert not any("LLM stream finished [sync-fail-step]" in m for m in messages)
         assert not any("LLM stream chunk [sync-fail-step]" in m for m in messages)
 
@@ -2797,8 +2801,10 @@ async def test_llm_call_async_logs_interruption_before_reraising(self, caplog):
         caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler")
         error = RuntimeError("stream exploded")
         with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            # Fresh _AsyncStream per call (see the sync test above for why):
+            # _llm_call_async now retries 3 times total.
             mock_litellm.acompletion = AsyncMock(
-                return_value=_AsyncStream(
+                side_effect=lambda *a, **k: _AsyncStream(
                     ["chunk-1", "chunk-2"],
                     error=error,
                     raise_after=2,
@@ -2813,9 +2819,9 @@ async def test_llm_call_async_logs_interruption_before_reraising(self, caplog):
         interrupted_logs = [
             m for m in messages if "LLM stream [async-fail-step] interrupted unexpectedly" in m
         ]
-        assert len(start_logs) == 1
-        assert len(interrupted_logs) == 1
-        assert "after chunk 2" in interrupted_logs[0]
+        assert len(start_logs) == 3
+        assert len(interrupted_logs) == 3
+        assert all("after chunk 2" in m for m in interrupted_logs)
         assert not any("LLM stream finished [async-fail-step]" in m for m in messages)
         assert not any("LLM stream chunk [async-fail-step]" in m for m in messages)
 
@@ -2831,7 +2837,11 @@ def test_llm_call_logs_interruption_before_any_chunk(self, caplog):
             patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
             patch("openkb.agent.compiler.litellm") as mock_litellm,
         ):
-            mock_litellm.completion = MagicMock(return_value=_stream_then_raise([], error))
+            # Fresh generator per call (see test_llm_call_logs_interruption_before_reraising):
+            # _llm_call now retries 3 times total.
+            mock_litellm.completion = MagicMock(
+                side_effect=lambda *a, **k: _stream_then_raise([], error)
+            )
 
             with pytest.raises(RuntimeError, match="connect timeout"):
                 _llm_call("m", [{"role": "user", "content": "hi"}], "sync-no-chunk-step")
@@ -2840,8 +2850,8 @@ def test_llm_call_logs_interruption_before_any_chunk(self, caplog):
         interrupted_logs = [
             m for m in messages if "LLM stream [sync-no-chunk-step] interrupted unexpectedly" in m
         ]
-        assert len(interrupted_logs) == 1
-        assert "before any chunk arrived" in interrupted_logs[0]
+        assert len(interrupted_logs) == 3
+        assert all("before any chunk arrived" in m for m in interrupted_logs)
         assert not any("LLM stream started [sync-no-chunk-step]" in m for m in messages)
 
 

From 65de48e7272dbbefc117f1ce5bed0890b7d77178 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Tue, 8 Sep 2026 08:36:15 +0200
Subject: [PATCH 5/9] fix(agent): skip non-retryable context-window errors and
 oversized docs

_llm_call/_llm_call_async no longer waste their 3-attempt retry floor on litellm.ContextWindowExceededError, since an identical retry is guaranteed to fail again. compile_short_doc additionally preflights the summary prompt's token count (when the model is recognized by litellm) and, on either that check or a ContextWindowExceededError from the summary call itself, skips LLM ingestion for the document entirely: the raw source stays in the KB as a plain reference (like an unreadable image would), with a stub summary noting why, instead of aborting the add or discarding the file.
---
 openkb/agent/compiler.py | 100 +++++++++++++++++++++++++++++++++----
 tests/test_compiler.py   | 105 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 196 insertions(+), 9 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 420f0bcc8..0475de56e 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -397,6 +397,33 @@ class TruncatedResponseError(Exception):
     treat truncation as a failure (so a partial page is skipped, not written)."""
 
 
+# Exceptions that are guaranteed to fail again on an identical retry (e.g. a
+# prompt that already exceeds the model's context window) — retrying wastes
+# whole attempts (and, for a stream, the connection time to discover the
+# failure again) for zero chance of success.
+_NON_RETRYABLE_LLM_ERRORS: tuple[type[Exception], ...] = (litellm.ContextWindowExceededError,)
+
+
+def _max_input_tokens(model: str) -> int | None:
+    """Best-effort context-window lookup for ``model``; ``None`` if unknown.
+
+    Used only to skip a call that's already known to be doomed before it's
+    even sent — never to second-guess a model litellm/the provider doesn't
+    also recognize, so an unmapped model just disables the preflight check.
+    """
+    try:
+        max_input_tokens = litellm.get_model_info(model).get("max_input_tokens")
+    except Exception:
+        return None
+    return int(max_input_tokens) if isinstance(max_input_tokens, int | float) else None
+
+
+# Reserved headroom (completion + rough token-counting slack) subtracted from
+# a model's context window before comparing it to the prompt's token count —
+# a prompt that just barely fits leaves no room for the model to respond.
+_CONTEXT_WINDOW_HEADROOM_TOKENS = 4096
+
+
 def _merge_stream_chunks(chunks: list, messages: list[dict]):
     """Merge streamed LLM chunks back into a single, non-streaming response.
 
@@ -575,7 +602,7 @@ def _llm_call(
             response = _merge_stream_chunks(chunks, messages)
             break
         except Exception as exc:
-            if attempt == attempts - 1:
+            if attempt == attempts - 1 or isinstance(exc, _NON_RETRYABLE_LLM_ERRORS):
                 spinner.stop("failed")
                 raise
             logger.warning(
@@ -647,7 +674,7 @@ async def _llm_call_async(
             response = _merge_stream_chunks(chunks, messages)
             break
         except Exception as exc:
-            if attempt == attempts - 1:
+            if attempt == attempts - 1 or isinstance(exc, _NON_RETRYABLE_LLM_ERRORS):
                 raise
             logger.warning(
                 "LLM [%s] attempt %d/%d failed: %s; retrying...",
@@ -1146,6 +1173,30 @@ def _write_summary(
     atomic_write_text(summaries_dir / f"{doc_name}.md", fm_block + summary)
 
 
+def _write_unprocessable_stub(wiki_dir: Path, doc_name: str, reason: str) -> None:
+    """Write a placeholder summary for a doc whose content can't be fed to the LLM.
+
+    Mirrors how an unreadable/undecodable image is handled: the source stays
+    in the knowledge base as a plain reference instead of aborting the whole
+    ``add`` or discarding it, but it's skipped for LLM ingestion entirely (no
+    summary/concept/entity generation), since a request this size is either
+    already known to exceed the model's context window or has just failed
+    with ``litellm.ContextWindowExceededError``.
+    """
+    body = (
+        "This document was not processed by the LLM: its content is too "
+        f"large for the model's context window ({reason}). The raw source "
+        "is still kept in the knowledge base for reference, but no summary "
+        "or concept/entity extraction was generated for it."
+    )
+    _write_summary(
+        wiki_dir,
+        doc_name,
+        body,
+        description="Not processed by the LLM \u2014 content too large for the context window.",
+    )
+
+
 _SAFE_NAME_RE = re.compile(r"[^\w\-]")
 
 
@@ -2434,18 +2485,49 @@ async def compile_short_doc(
         ),
     }
 
+    # Preflight: skip the LLM entirely for a doc that's already known to
+    # exceed the model's context window (only when the model is recognized —
+    # see _max_input_tokens) instead of sending a request that's certain to
+    # fail. The doc stays in the KB as a plain reference, like an unreadable
+    # image would.
+    max_input_tokens = _max_input_tokens(model)
+    if max_input_tokens is not None:
+        prompt_tokens = litellm.token_counter(model=model, messages=[system_msg, doc_msg])
+        if prompt_tokens > max_input_tokens - _CONTEXT_WINDOW_HEADROOM_TOKENS:
+            logger.warning(
+                "Skipping LLM ingestion for %s: %d prompt tokens > %s's %d-token context window",
+                doc_name,
+                prompt_tokens,
+                model,
+                max_input_tokens,
+            )
+            _write_unprocessable_stub(
+                wiki_dir,
+                doc_name,
+                f"{prompt_tokens} tokens > {model}'s {max_input_tokens}-token context window",
+            )
+            return
+
     # --- Step 1: Generate summary (v1, held in memory) ---
     # The summary is NOT written to disk yet — it's used as cache context
     # for the plan + concept-generation calls, then rewritten into a final
     # v2 (with a whitelist of known wikilink targets) inside
     # _compile_concepts before being written to disk.
-    summary_raw = _llm_call(
-        model,
-        [system_msg, doc_msg],
-        "summary",
-        response_format=_JSON_RESPONSE_FORMAT,
-        bundle=bundle,
-    )
+    try:
+        summary_raw = _llm_call(
+            model,
+            [system_msg, doc_msg],
+            "summary",
+            response_format=_JSON_RESPONSE_FORMAT,
+            bundle=bundle,
+        )
+    except _NON_RETRYABLE_LLM_ERRORS as exc:
+        # The preflight check above is best-effort (unmapped model, or
+        # litellm's token_counter estimate came in under the real one) — this
+        # is the safety net for when it still slips through.
+        logger.warning("Skipping LLM ingestion for %s: %s", doc_name, exc)
+        _write_unprocessable_stub(wiki_dir, doc_name, str(exc))
+        return
     try:
         summary_parsed = _parse_json(summary_raw)
         doc_brief = summary_parsed.get("description", "")
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 665827295..4c7e274b8 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -6,6 +6,7 @@
 import logging
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import litellm
 import pytest
 
 from openkb.agent.compiler import (
@@ -1503,6 +1504,110 @@ async def test_scalar_plan_handled_gracefully(self, tmp_path):
         assert not list((wiki / "concepts").glob("*.md"))
 
 
+class TestOversizedDocumentSkip:
+    """A document whose content can't fit an LLM call is treated like an
+    unreadable image: kept in the KB as a plain reference, but skipped for
+    LLM ingestion entirely rather than aborting the whole ``add`` or wasting
+    retries on a request that's guaranteed to fail again (#<TODO>)."""
+
+    @pytest.mark.asyncio
+    async def test_skips_llm_call_when_preflight_detects_oversized_prompt(self, tmp_path):
+        wiki, source_path = TestCompileShortDocFallbacks._setup_kb(tmp_path)
+
+        with (
+            patch("openkb.agent.compiler._max_input_tokens", return_value=1000),
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.token_counter = MagicMock(return_value=999_999)
+            await compile_short_doc("doc", source_path, tmp_path, "gpt-4o-mini")
+            mock_litellm.completion.assert_not_called()
+
+        summary_path = wiki / "summaries" / "doc.md"
+        assert summary_path.exists()
+        text = summary_path.read_text()
+        assert "too large" in text.lower()
+        assert not list((wiki / "concepts").glob("*.md"))
+
+    @pytest.mark.asyncio
+    async def test_writes_stub_when_summary_call_raises_context_window_exceeded(self, tmp_path):
+        wiki, source_path = TestCompileShortDocFallbacks._setup_kb(tmp_path)
+
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 569887 tokens > 200000 maximum",
+            model="claude-sonnet-4-5",
+            llm_provider="anthropic",
+        )
+        with (
+            patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.completion = MagicMock(side_effect=error)
+            # Must not raise out of compile_short_doc.
+            await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5")
+            # Only the summary call was attempted — no retries wasted on a
+            # deterministically doomed request, no concept-plan call either.
+            assert mock_litellm.completion.call_count == 1
+
+        summary_path = wiki / "summaries" / "doc.md"
+        assert summary_path.exists()
+        text = summary_path.read_text()
+        assert "too large" in text.lower()
+        assert not list((wiki / "concepts").glob("*.md"))
+
+
+class TestMaxInputTokens:
+    """``_max_input_tokens`` must only ever be a best-effort hint: an
+    unrecognized model disables the preflight check instead of guessing."""
+
+    def test_known_model_returns_context_window(self):
+        from openkb.agent.compiler import _max_input_tokens
+
+        assert _max_input_tokens("claude-sonnet-4-5") == 200_000
+
+    def test_unknown_model_returns_none(self):
+        from openkb.agent.compiler import _max_input_tokens
+
+        assert _max_input_tokens("totally-unknown-model-xyz") is None
+
+
+class TestNonRetryableLLMErrors:
+    """litellm.ContextWindowExceededError means the exact same request will
+    fail again — retrying it just burns the fixed 3-attempt resilience floor
+    on a request that can never succeed."""
+
+    def test_llm_call_does_not_retry_context_window_exceeded(self):
+        from openkb.agent.compiler import _llm_call
+
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 569887 tokens > 200000 maximum",
+            model="m",
+            llm_provider="anthropic",
+        )
+        with (
+            patch("openkb.agent.compiler._Spinner", _NoOpSpinner),
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.completion = MagicMock(side_effect=error)
+            with pytest.raises(litellm.ContextWindowExceededError):
+                _llm_call("m", [{"role": "user", "content": "hi"}], "step")
+            assert mock_litellm.completion.call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_llm_call_async_does_not_retry_context_window_exceeded(self):
+        from openkb.agent.compiler import _llm_call_async
+
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 569887 tokens > 200000 maximum",
+            model="m",
+            llm_provider="anthropic",
+        )
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.acompletion = AsyncMock(side_effect=error)
+            with pytest.raises(litellm.ContextWindowExceededError):
+                await _llm_call_async("m", [{"role": "user", "content": "hi"}], "step")
+            assert mock_litellm.acompletion.call_count == 1
+
+
 class TestCacheControl:
     """Verify cache_control breakpoints are emitted on the right messages
     so Anthropic prompt caching can hit on every reuse of the base context.

From 69c79b0a6ff4ff14edf1a650c453d897aa394938 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Tue, 8 Sep 2026 09:24:50 +0200
Subject: [PATCH 6/9] fix(agent): add index.md entry for the oversized-document
 stub summary

_write_unprocessable_stub() previously only wrote the stub summary page, never calling _update_index() (that only happened inside _compile_concepts(), which the oversized-doc path deliberately skips). The stub page was therefore an undiscoverable orphan: missing from index.md's ## Documents section, which openkb lint's check_index_sync() flags as an index-sync error. Now calls _update_index(wiki_dir, doc_name, [], doc_brief=description) too, same as every other short-doc path.
---
 openkb/agent/compiler.py | 14 +++++++-------
 tests/test_compiler.py   |  8 ++++++++
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 0475de56e..b41628be2 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -1181,20 +1181,20 @@ def _write_unprocessable_stub(wiki_dir: Path, doc_name: str, reason: str) -> Non
     ``add`` or discarding it, but it's skipped for LLM ingestion entirely (no
     summary/concept/entity generation), since a request this size is either
     already known to exceed the model's context window or has just failed
-    with ``litellm.ContextWindowExceededError``.
+    with ``litellm.ContextWindowExceededError``. Also adds the usual
+    ``## Documents`` index.md entry (via ``_update_index``, with no concepts)
+    — without it the summary page would be an undiscoverable orphan, which
+    ``openkb lint`` flags as an index-sync error.
     """
+    description = "Not processed by the LLM \u2014 content too large for the context window."
     body = (
         "This document was not processed by the LLM: its content is too "
         f"large for the model's context window ({reason}). The raw source "
         "is still kept in the knowledge base for reference, but no summary "
         "or concept/entity extraction was generated for it."
     )
-    _write_summary(
-        wiki_dir,
-        doc_name,
-        body,
-        description="Not processed by the LLM \u2014 content too large for the context window.",
-    )
+    _write_summary(wiki_dir, doc_name, body, description=description)
+    _update_index(wiki_dir, doc_name, [], doc_brief=description)
 
 
 _SAFE_NAME_RE = re.compile(r"[^\w\-]")
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 4c7e274b8..a82cb7df9 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -1528,6 +1528,11 @@ async def test_skips_llm_call_when_preflight_detects_oversized_prompt(self, tmp_
         assert "too large" in text.lower()
         assert not list((wiki / "concepts").glob("*.md"))
 
+        # The stub must still get the usual index.md "## Documents" entry —
+        # otherwise it's an undiscoverable orphan (openkb lint's index-sync check).
+        index_text = (wiki / "index.md").read_text()
+        assert "[[summaries/doc]]" in index_text
+
     @pytest.mark.asyncio
     async def test_writes_stub_when_summary_call_raises_context_window_exceeded(self, tmp_path):
         wiki, source_path = TestCompileShortDocFallbacks._setup_kb(tmp_path)
@@ -1554,6 +1559,9 @@ async def test_writes_stub_when_summary_call_raises_context_window_exceeded(self
         assert "too large" in text.lower()
         assert not list((wiki / "concepts").glob("*.md"))
 
+        index_text = (wiki / "index.md").read_text()
+        assert "[[summaries/doc]]" in index_text
+
 
 class TestMaxInputTokens:
     """``_max_input_tokens`` must only ever be a best-effort hint: an

From 338d1bbf14c388ef54ef7cb654fc0769a24d3236 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Tue, 8 Sep 2026 16:05:09 +0200
Subject: [PATCH 7/9] fix(agent): skip context-window-exceeded concepts-plan
 without discarding the summary

The existing concept/entity index (_read_concept_briefs/_read_entity_briefs) grows unbounded with the KB and is appended on top of the already-cached document for the concepts-plan call -- a document whose summary call fit comfortably can still blow the context window once that index gets large enough. Previously this exception propagated uncaught out of _compile_concepts(), through compile_short_doc(), retried the whole doc (including a wasted repeat summary call) via cli.py's outer retry, and ultimately failed the add.

_compile_concepts() now catches _NON_RETRYABLE_LLM_ERRORS around the concepts-plan call. Since the summary already succeeded by this point (unlike the doc-too-large case in compile_short_doc's own preflight/summary-call catch), it reuses the same fallback already used for an unparseable/empty plan: write the real v1 summary (ghost-wikilink-stripped) and the normal index.md entry, just skipping concept/entity generation for this doc -- instead of discarding a summary that was already successfully generated. Applies to compile_long_doc() too, since it shares _compile_concepts().
---
 openkb/agent/compiler.py | 52 +++++++++++++++++++++++++---------------
 tests/test_compiler.py   | 47 ++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 19 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index b41628be2..8488c9d81 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -1871,25 +1871,6 @@ async def _compile_concepts(
     # (system + doc + summary) for the plan call and every concept call.
     summary_msg = {"role": "assistant", "content": _cached_text(summary)}
 
-    plan_raw = _llm_call(
-        model,
-        [
-            system_msg,
-            doc_msg,
-            summary_msg,
-            {
-                "role": "user",
-                "content": _CONCEPTS_PLAN_USER.format(
-                    concept_briefs=concept_briefs,
-                    entity_briefs=entity_briefs,
-                ).replace("__ENTITY_TYPES__", types_str),
-            },
-        ],
-        "concepts-plan",
-        response_format=_JSON_RESPONSE_FORMAT,
-        bundle=bundle,
-    )
-
     def _write_v1_summary_stripped() -> None:
         """Fallback writer for the v1 summary on early-return paths.
 
@@ -1911,6 +1892,39 @@ def _write_v1_summary_stripped() -> None:
             )
         _write_summary(wiki_dir, doc_name, cleaned, description=doc_brief)
 
+    try:
+        plan_raw = _llm_call(
+            model,
+            [
+                system_msg,
+                doc_msg,
+                summary_msg,
+                {
+                    "role": "user",
+                    "content": _CONCEPTS_PLAN_USER.format(
+                        concept_briefs=concept_briefs,
+                        entity_briefs=entity_briefs,
+                    ).replace("__ENTITY_TYPES__", types_str),
+                },
+            ],
+            "concepts-plan",
+            response_format=_JSON_RESPONSE_FORMAT,
+            bundle=bundle,
+        )
+    except _NON_RETRYABLE_LLM_ERRORS as exc:
+        # The existing concept/entity index grows with the KB (see
+        # _read_concept_briefs/_read_entity_briefs) and is added on top of the
+        # already-cached document — a doc that was fine for the summary call
+        # can still blow the window here once the index gets large enough.
+        # The summary itself already exists (unlike the doc-too-large case
+        # above), so it's kept — same fallback as an unparseable/empty plan,
+        # just skipping concept/entity generation for this doc.
+        logger.warning("Skipping concept/entity extraction for %s: %s", doc_name, exc)
+        if rewrite_summary:
+            _write_v1_summary_stripped()
+        _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type)
+        return
+
     try:
         parsed = _parse_json(plan_raw)
     except (json.JSONDecodeError, ValueError) as exc:
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index a82cb7df9..288077ba8 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -1503,6 +1503,53 @@ async def test_scalar_plan_handled_gracefully(self, tmp_path):
         # No concept pages produced from the unusable plan.
         assert not list((wiki / "concepts").glob("*.md"))
 
+    @pytest.mark.asyncio
+    async def test_concepts_plan_context_window_exceeded_keeps_real_summary(self, tmp_path):
+        """The summary call already succeeded with real content by the time
+        concepts-plan runs (unlike TestOversizedDocumentSkip, where the doc
+        itself is too large) — this must NOT be replaced with a generic
+        "too large" stub. Same fallback as an unparseable/empty plan: keep
+        the real v1 summary (ghost-stripped), index it, skip concept/entity
+        generation for this doc."""
+        wiki, source_path = self._setup_kb(tmp_path)
+
+        v1_summary_content = "# Summary\n\nDiscusses [[concepts/nonexistent]] here."
+        summary_response = json.dumps(
+            {"description": "A real summary", "content": v1_summary_content}
+        )
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 220670 tokens > 200000 maximum",
+            model="claude-sonnet-4-5",
+            llm_provider="anthropic",
+        )
+        call_count = {"n": 0}
+
+        def side_effect(*args, **kwargs):
+            idx = call_count["n"]
+            call_count["n"] += 1
+            if idx == 0:
+                return [_mock_response(summary_response)]
+            raise error
+
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.completion = MagicMock(side_effect=side_effect)
+            # Must not raise out of compile_short_doc.
+            await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5")
+            # summary + concepts-plan, no wasted retries on the doomed request.
+            assert mock_litellm.completion.call_count == 2
+
+        summary_path = wiki / "summaries" / "doc.md"
+        assert summary_path.exists()
+        text = summary_path.read_text()
+        assert "Discusses" in text  # real summary content kept, not a stub
+        assert "too large" not in text.lower()
+        assert "[[concepts/nonexistent]]" not in text  # ghost link stripped
+        assert "nonexistent" in text  # display text preserved
+
+        index_text = (wiki / "index.md").read_text()
+        assert "[[summaries/doc]]" in index_text
+        assert not list((wiki / "concepts").glob("*.md"))
+
 
 class TestOversizedDocumentSkip:
     """A document whose content can't fit an LLM call is treated like an

From d5bc8637045d2a595d1d8e9bec27ccebc0c102dc Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Thu, 10 Sep 2026 08:10:55 +0200
Subject: [PATCH 8/9] fix(agent): retry concepts-plan with summary as source
 before giving up

---
 openkb/agent/compiler.py | 61 +++++++++++++++++++++-----------
 tests/test_compiler.py   | 76 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 112 insertions(+), 25 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 8488c9d81..58aacad15 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -1892,21 +1892,17 @@ def _write_v1_summary_stripped() -> None:
             )
         _write_summary(wiki_dir, doc_name, cleaned, description=doc_brief)
 
+    concepts_plan_user_msg = {
+        "role": "user",
+        "content": _CONCEPTS_PLAN_USER.format(
+            concept_briefs=concept_briefs,
+            entity_briefs=entity_briefs,
+        ).replace("__ENTITY_TYPES__", types_str),
+    }
     try:
         plan_raw = _llm_call(
             model,
-            [
-                system_msg,
-                doc_msg,
-                summary_msg,
-                {
-                    "role": "user",
-                    "content": _CONCEPTS_PLAN_USER.format(
-                        concept_briefs=concept_briefs,
-                        entity_briefs=entity_briefs,
-                    ).replace("__ENTITY_TYPES__", types_str),
-                },
-            ],
+            [system_msg, doc_msg, summary_msg, concepts_plan_user_msg],
             "concepts-plan",
             response_format=_JSON_RESPONSE_FORMAT,
             bundle=bundle,
@@ -1916,14 +1912,39 @@ def _write_v1_summary_stripped() -> None:
         # _read_concept_briefs/_read_entity_briefs) and is added on top of the
         # already-cached document — a doc that was fine for the summary call
         # can still blow the window here once the index gets large enough.
-        # The summary itself already exists (unlike the doc-too-large case
-        # above), so it's kept — same fallback as an unparseable/empty plan,
-        # just skipping concept/entity generation for this doc.
-        logger.warning("Skipping concept/entity extraction for %s: %s", doc_name, exc)
-        if rewrite_summary:
-            _write_v1_summary_stripped()
-        _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type)
-        return
+        # Retry once without the full document: the plan prompt already asks
+        # the model to work "based on the summary above" (see
+        # _CONCEPTS_PLAN_USER), so dropping doc_msg usually shrinks the
+        # prompt back under the window without losing the plan's intent.
+        logger.warning(
+            "concepts plan exceeded context window for %s: %s; retrying with summary as source",
+            doc_name,
+            exc,
+        )
+        sys.stdout.write(
+            f"    [WARN] concepts plan exceeded context window for {doc_name} — "
+            "retrying with the summary as source instead of the full document.\n"
+        )
+        sys.stdout.flush()
+        try:
+            plan_raw = _llm_call(
+                model,
+                [system_msg, summary_msg, concepts_plan_user_msg],
+                "concepts-plan",
+                response_format=_JSON_RESPONSE_FORMAT,
+                bundle=bundle,
+            )
+        except _NON_RETRYABLE_LLM_ERRORS as retry_exc:
+            # Still too large even with just the summary (e.g. the existing
+            # concept/entity index alone is huge) — the summary itself
+            # already exists (unlike the doc-too-large case above), so it's
+            # kept — same fallback as an unparseable/empty plan, just
+            # skipping concept/entity generation for this doc.
+            logger.warning("Skipping concept/entity extraction for %s: %s", doc_name, retry_exc)
+            if rewrite_summary:
+                _write_v1_summary_stripped()
+            _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type)
+            return
 
     try:
         parsed = _parse_json(plan_raw)
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 288077ba8..02851bf7c 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -1508,9 +1508,11 @@ async def test_concepts_plan_context_window_exceeded_keeps_real_summary(self, tm
         """The summary call already succeeded with real content by the time
         concepts-plan runs (unlike TestOversizedDocumentSkip, where the doc
         itself is too large) — this must NOT be replaced with a generic
-        "too large" stub. Same fallback as an unparseable/empty plan: keep
-        the real v1 summary (ghost-stripped), index it, skip concept/entity
-        generation for this doc."""
+        "too large" stub. The concepts-plan call is retried once with the
+        summary as source (doc_msg dropped) before giving up; if that retry
+        ALSO hits the context window, same fallback as an unparseable/empty
+        plan: keep the real v1 summary (ghost-stripped), index it, skip
+        concept/entity generation for this doc."""
         wiki, source_path = self._setup_kb(tmp_path)
 
         v1_summary_content = "# Summary\n\nDiscusses [[concepts/nonexistent]] here."
@@ -1535,8 +1537,15 @@ def side_effect(*args, **kwargs):
             mock_litellm.completion = MagicMock(side_effect=side_effect)
             # Must not raise out of compile_short_doc.
             await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5")
-            # summary + concepts-plan, no wasted retries on the doomed request.
-            assert mock_litellm.completion.call_count == 2
+            # summary + concepts-plan (full doc) + concepts-plan retry (summary
+            # only) — both plan attempts are doomed here, no further retries.
+            assert mock_litellm.completion.call_count == 3
+
+        # The retry attempt must not still include the full document message
+        # (_SUMMARY_USER's "Full text:" marker only ever appears in doc_msg).
+        retry_messages = mock_litellm.completion.call_args_list[2].kwargs["messages"]
+        retry_text = json.dumps(retry_messages)
+        assert "Full text:" not in retry_text
 
         summary_path = wiki / "summaries" / "doc.md"
         assert summary_path.exists()
@@ -1550,6 +1559,63 @@ def side_effect(*args, **kwargs):
         assert "[[summaries/doc]]" in index_text
         assert not list((wiki / "concepts").glob("*.md"))
 
+    @pytest.mark.asyncio
+    async def test_concepts_plan_context_window_retry_succeeds_with_summary_only(self, tmp_path):
+        """When the retry (summary-only) succeeds, the plan is processed
+        normally — concepts are generated from the summary-derived plan
+        instead of the doc being treated as incomplete."""
+        wiki, source_path = self._setup_kb(tmp_path)
+
+        v1_summary_content = "# Summary\n\nDiscusses transformers."
+        summary_response = json.dumps(
+            {"description": "A real summary", "content": v1_summary_content}
+        )
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 220670 tokens > 200000 maximum",
+            model="claude-sonnet-4-5",
+            llm_provider="anthropic",
+        )
+        plan_response = json.dumps(
+            {
+                "create": [{"name": "transformer", "title": "Transformer"}],
+                "update": [],
+                "related": [],
+            }
+        )
+        concept_response = json.dumps({"description": "C", "content": "# T\n\nBody."})
+        call_count = {"n": 0}
+
+        def side_effect(*args, **kwargs):
+            idx = call_count["n"]
+            call_count["n"] += 1
+            if idx == 0:
+                return [_mock_response(summary_response)]
+            if idx == 1:
+                raise error
+            return [_mock_response(plan_response)]
+
+        with (
+            patch("openkb.agent.compiler.litellm") as mock_litellm,
+        ):
+            mock_litellm.completion = MagicMock(side_effect=side_effect)
+            mock_litellm.acompletion = AsyncMock(side_effect=_mock_acompletion([concept_response]))
+            await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5")
+            # summary + concepts-plan (full doc, fails) + concepts-plan retry
+            # (summary only, succeeds) + summary-rewrite (rewrite_summary=True).
+            assert mock_litellm.completion.call_count == 4
+
+        # The successful retry's messages must not include the full document
+        # message (_SUMMARY_USER's "Full text:" marker only ever appears in doc_msg).
+        retry_messages = mock_litellm.completion.call_args_list[2].kwargs["messages"]
+        retry_text = json.dumps(retry_messages)
+        assert "Full text:" not in retry_text
+
+        concept_path = wiki / "concepts" / "transformer.md"
+        assert concept_path.exists()
+
+        index_text = (wiki / "index.md").read_text()
+        assert "[[concepts/transformer]]" in index_text
+
 
 class TestOversizedDocumentSkip:
     """A document whose content can't fit an LLM call is treated like an

From 41a68acb67c703580de7098aa32641d5795c8584 Mon Sep 17 00:00:00 2001
From: Sebastian Braun <sebastian.braun@haufegroup.com>
Date: Thu, 10 Sep 2026 11:14:19 +0200
Subject: [PATCH 9/9] fix(agent): retry summary-rewrite with summary as source
 before giving up

known_targets_msg (the whitelist of every existing concept/entity page) grows
with the KB, same as concepts-plan's brief lists (#226) - a document whose
concepts-plan call still fit can push summary-rewrite over the context
window once that whitelist gets large enough. Previously any exception from
this call, including litellm.ContextWindowExceededError, fell straight
through to the v1-summary fallback, silently forfeiting the improved
cross-linking summary-rewrite exists to produce.

Mirrors the existing concepts-plan fix: on _NON_RETRYABLE_LLM_ERRORS, retry
once with doc_msg dropped (the rewrite prompt only needs summary_msg and
known_targets_msg, not the original document) before falling through to the
unchanged v1 fallback.

Resolves #256

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 openkb/agent/compiler.py |  51 +++++++++++++----
 tests/test_compiler.py   | 121 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 160 insertions(+), 12 deletions(-)

diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py
index 58aacad15..9aa3bb43d 100644
--- a/openkb/agent/compiler.py
+++ b/openkb/agent/compiler.py
@@ -2365,21 +2365,48 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]:
     # the full whitelist, so the summary is always written and never wiped.
     if rewrite_summary:
         candidate: str | None = None
+        summary_rewrite_user_msg = {"role": "user", "content": _SUMMARY_REWRITE_USER}
         try:
             # No max_tokens cap — matches the v1 summary call. The rewrite
             # prompt asks the model to keep length within ±20% of the v1.
-            rewrite_raw = _llm_call(
-                model,
-                [
-                    system_msg,
-                    doc_msg,  # cached (BP1)
-                    summary_msg,  # cached (BP2) — contains the v1 summary text
-                    known_targets_msg,  # cached (BP3) — whitelist
-                    {"role": "user", "content": _SUMMARY_REWRITE_USER},
-                ],
-                "summary-rewrite",
-                bundle=bundle,
-            )
+            try:
+                rewrite_raw = _llm_call(
+                    model,
+                    [
+                        system_msg,
+                        doc_msg,  # cached (BP1)
+                        summary_msg,  # cached (BP2) — contains the v1 summary text
+                        known_targets_msg,  # cached (BP3) — whitelist
+                        summary_rewrite_user_msg,
+                    ],
+                    "summary-rewrite",
+                    bundle=bundle,
+                )
+            except _NON_RETRYABLE_LLM_ERRORS as exc:
+                # known_targets_msg grows with the KB just like the
+                # concepts-plan index (#226) — a doc whose earlier calls fit
+                # can still blow the window here once the whitelist gets
+                # large enough. Retry once without the full document: the
+                # rewrite prompt only asks the model to reconcile the
+                # already-generated summary (summary_msg) against the
+                # whitelist, not the original document.
+                logger.warning(
+                    "summary-rewrite exceeded context window for %s: %s; "
+                    "retrying with summary as source",
+                    doc_name,
+                    exc,
+                )
+                sys.stdout.write(
+                    f"    [WARN] summary-rewrite exceeded context window for {doc_name} — "
+                    "retrying with the summary as source instead of the full document.\n"
+                )
+                sys.stdout.flush()
+                rewrite_raw = _llm_call(
+                    model,
+                    [system_msg, summary_msg, known_targets_msg, summary_rewrite_user_msg],
+                    "summary-rewrite",
+                    bundle=bundle,
+                )
             candidate = rewrite_raw.strip()
             # Strip frontmatter if the model added one anyway.
             cand_parts = frontmatter.split(candidate)
diff --git a/tests/test_compiler.py b/tests/test_compiler.py
index 02851bf7c..c0ecebaeb 100644
--- a/tests/test_compiler.py
+++ b/tests/test_compiler.py
@@ -1616,6 +1616,127 @@ def side_effect(*args, **kwargs):
         index_text = (wiki / "index.md").read_text()
         assert "[[concepts/transformer]]" in index_text
 
+    @pytest.mark.asyncio
+    async def test_summary_rewrite_context_window_exceeded_retries_with_summary_only(
+        self, tmp_path
+    ):
+        """summary-rewrite's prompt carries doc_msg plus the whitelist of
+        every existing concept/entity page (known_targets_msg), which grows
+        with the KB (#226) just like the concepts-plan index — so it can
+        blow the context window on a document whose concepts-plan call (a
+        smaller prompt, no whitelist) still fit. Mirrors the concepts-plan
+        fix: retry once with doc_msg dropped before falling back to v1."""
+        wiki, source_path = self._setup_kb(tmp_path)
+        (wiki / "concepts" / "transformer.md").write_text(
+            "---\ndescription: Existing\n---\n\nExisting content.", encoding="utf-8"
+        )
+
+        v1_summary_content = "# Summary\n\nDiscusses transformers."
+        summary_response = json.dumps(
+            {"description": "A real summary", "content": v1_summary_content}
+        )
+        plan_response = json.dumps(
+            {
+                "create": [],
+                "update": [{"name": "transformer", "title": "Transformer"}],
+                "related": [],
+            }
+        )
+        concept_response = json.dumps({"description": "C", "content": "# T\n\nUpdated body."})
+        rewritten_summary = "# Summary\n\nRewritten: discusses [[concepts/transformer]]."
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 200400 tokens > 200000 maximum",
+            model="claude-sonnet-4-5",
+            llm_provider="anthropic",
+        )
+        call_count = {"n": 0}
+
+        def side_effect(*args, **kwargs):
+            idx = call_count["n"]
+            call_count["n"] += 1
+            if idx == 0:
+                return [_mock_response(summary_response)]
+            if idx == 1:
+                return [_mock_response(plan_response)]
+            if idx == 2:
+                raise error
+            return [_mock_response(rewritten_summary)]
+
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.completion = MagicMock(side_effect=side_effect)
+            mock_litellm.acompletion = AsyncMock(side_effect=_mock_acompletion([concept_response]))
+            await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5")
+            # summary + concepts-plan + summary-rewrite (full doc, fails) +
+            # summary-rewrite retry (summary only, succeeds).
+            assert mock_litellm.completion.call_count == 4
+
+        # The retry attempt must not still include the full document message
+        # (_SUMMARY_USER's "Full text:" marker only ever appears in doc_msg).
+        retry_messages = mock_litellm.completion.call_args_list[3].kwargs["messages"]
+        retry_text = json.dumps(retry_messages)
+        assert "Full text:" not in retry_text
+
+        summary_path = wiki / "summaries" / "doc.md"
+        assert summary_path.exists()
+        text = summary_path.read_text()
+        assert "Rewritten" in text  # retried rewrite content used, not v1
+        assert "[[concepts/transformer]]" in text
+
+    @pytest.mark.asyncio
+    async def test_summary_rewrite_context_window_retry_also_fails_falls_back_to_v1(self, tmp_path):
+        """If the summary-only retry ALSO hits the context window (or any
+        other error), the existing v1 fallback still applies unchanged —
+        the real v1 summary (ghost-stripped) is written, never lost."""
+        wiki, source_path = self._setup_kb(tmp_path)
+        (wiki / "concepts" / "transformer.md").write_text(
+            "---\ndescription: Existing\n---\n\nExisting content.", encoding="utf-8"
+        )
+
+        v1_summary_content = (
+            "# Summary\n\nDiscusses [[concepts/transformer]] and [[concepts/ghost]]."
+        )
+        summary_response = json.dumps(
+            {"description": "A real summary", "content": v1_summary_content}
+        )
+        plan_response = json.dumps(
+            {
+                "create": [],
+                "update": [{"name": "transformer", "title": "Transformer"}],
+                "related": [],
+            }
+        )
+        concept_response = json.dumps({"description": "C", "content": "# T\n\nUpdated body."})
+        error = litellm.ContextWindowExceededError(
+            message="prompt is too long: 200400 tokens > 200000 maximum",
+            model="claude-sonnet-4-5",
+            llm_provider="anthropic",
+        )
+        call_count = {"n": 0}
+
+        def side_effect(*args, **kwargs):
+            idx = call_count["n"]
+            call_count["n"] += 1
+            if idx == 0:
+                return [_mock_response(summary_response)]
+            if idx == 1:
+                return [_mock_response(plan_response)]
+            raise error  # both the first rewrite attempt and its retry fail
+
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.completion = MagicMock(side_effect=side_effect)
+            mock_litellm.acompletion = AsyncMock(side_effect=_mock_acompletion([concept_response]))
+            # Must not raise out of compile_short_doc.
+            await compile_short_doc("doc", source_path, tmp_path, "claude-sonnet-4-5")
+            assert mock_litellm.completion.call_count == 4
+
+        summary_path = wiki / "summaries" / "doc.md"
+        assert summary_path.exists()
+        text = summary_path.read_text()
+        assert "Discusses" in text  # real v1 content kept, not lost
+        assert "[[concepts/transformer]]" in text  # valid link kept
+        assert "[[concepts/ghost]]" not in text  # ghost link stripped
+        assert "ghost" in text  # display text preserved
+
 
 class TestOversizedDocumentSkip:
     """A document whose content can't fit an LLM call is treated like an