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>

Author: Sebastian Braun

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)

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