fix(cli): retry failed concept/entity generations before giving up, not the whole document

_compile_concepts now runs a deferred, end-of-first-pass sweep for concepts/
entities that failed their first attempt, under "normal" and "fail-at-end"
insert_mode (not "fail-fast", which already aborts before the rest of the
batch runs). This gives transient failures a real second chance, with the
prompt cache still warm, before insert_mode's strict modes decide the
document is incomplete.

_run_compile_with_retry (cli.py) no longer retries the whole document when
ConceptCompilationError is raised: insert_mode already exhausted the
per-call and per-item retries, so a full recompile would just repeat the
same failures. Other exceptions keep their existing 2-attempt retry.

Author: Sebastian Braun

openkb/agent/compiler.py

@@ -1603,6 +1603,27 @@ def _format_known_targets(targets: set[str]) -> str:
     return "\n".join(f"- {t}" for t in sorted(targets))
 
 
+async def _sweep_failed_generations(results: list, factories: list, kind: str) -> list:
+    """Retry, once, only the items that failed the first pass.
+
+    Not used under ``insert_mode="fail-fast"`` — that mode already aborts on
+    the first failure without waiting for the rest of the batch.
+    """
+    failed_idx = [i for i, r in enumerate(results) if isinstance(r, Exception)]
+    if not failed_idx:
+        return results
+    logger.warning(
+        "Retrying %d failed %s generation(s) after the first pass...", len(failed_idx), kind
+    )
+    retry_results = await asyncio.gather(
+        *(factories[i]() for i in failed_idx), return_exceptions=True
+    )
+    results = list(results)
+    for idx, r in zip(failed_idx, retry_results):
+        results[idx] = r
+    return results
+
+
 async def _compile_concepts(
     wiki_dir: Path,
     kb_dir: Path,
@@ -2025,6 +2046,12 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]:
     tasks.extend(asyncio.create_task(_gen_create(c)) for c in create_items)
     tasks.extend(asyncio.create_task(_gen_update(c)) for c in update_items)
 
+    # Zero-arg factories, same order as `tasks`, so a failed item can be
+    # re-run in isolation by the end-of-first-pass sweep below.
+    concept_factories = [(lambda c=c: _gen_create(c)) for c in create_items] + [
+        (lambda c=c: _gen_update(c)) for c in update_items
+    ]
+
     # --- Step 3 (entities): build the entity task list up front so it can be
     # gathered concurrently with the concept tasks below. Entity coroutines
     # return 4-arity tuples (name, content, brief, type), so their results are
@@ -2035,6 +2062,10 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]:
     entity_tasks.extend(asyncio.create_task(_gen_entity_create(e)) for e in entity_create)
     entity_tasks.extend(asyncio.create_task(_gen_entity_update(e)) for e in entity_update)
 
+    entity_factories = [(lambda e=e: _gen_entity_create(e)) for e in entity_create] + [
+        (lambda e=e: _gen_entity_update(e)) for e in entity_update
+    ]
+
     concept_names: list[str] = []
     concept_briefs_map: dict[str, str] = {}
     pending_writes: list[tuple[str, str, bool, str]] = []
@@ -2086,6 +2117,14 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]:
                 asyncio.gather(*tasks, return_exceptions=True),
                 asyncio.gather(*entity_tasks, return_exceptions=True),
             )
+            # One more chance for exactly the items that failed, now that the
+            # rest of the batch has run (transient conditions get real
+            # wall-clock time to clear, prompt cache is still warm). Not used
+            # under "fail-fast" — that mode already aborted above.
+            results = await _sweep_failed_generations(results, concept_factories, "concept")
+            entity_results = await _sweep_failed_generations(
+                entity_results, entity_factories, "entity"
+            )
 
     failure_types: list[str] = []
     if tasks:

openkb/cli.py

@@ -429,11 +429,17 @@ def _snapshot_add_paths(
 
 
 def _run_compile_with_retry(coro_factory, label: str) -> None:
+    from openkb.agent.compiler import ConceptCompilationError
+
     click.echo(f"  {label}...")
     for attempt in range(2):
         try:
             asyncio.run(coro_factory())
             return
+        except ConceptCompilationError:
+            # insert_mode already decided this document is incomplete —
+            # a full recompile would just repeat the same failures.
+            raise
         except Exception as exc:
             if attempt == 0:
                 click.echo("  Retrying compilation in 2s...")

tests/test_cli.py

@@ -303,6 +303,60 @@ def test_init_model_prompt_accepts_input(tmp_path):
         assert config["model"] == "anthropic/claude-opus-4-6"
 
 
+class TestRunCompileWithRetry:
+    """insert_mode already decided a document is incomplete when it raises
+    ConceptCompilationError, so _run_compile_with_retry must not spend a
+    second full-document recompile repeating the same failures — but every
+    other exception keeps its original 2-attempt retry."""
+
+    def test_concept_compilation_error_is_not_retried(self):
+        from openkb.agent.compiler import ConceptCompilationError
+        from openkb.cli import _run_compile_with_retry
+
+        calls = 0
+
+        async def _factory():
+            nonlocal calls
+            calls += 1
+            raise ConceptCompilationError("insert_mode said no")
+
+        with (
+            patch("openkb.cli.time.sleep"),
+            pytest.raises(ConceptCompilationError),
+        ):
+            _run_compile_with_retry(_factory, "compiling")
+        assert calls == 1
+
+    def test_other_exceptions_still_get_one_retry(self):
+        from openkb.cli import _run_compile_with_retry
+
+        calls = 0
+
+        async def _factory():
+            nonlocal calls
+            calls += 1
+            if calls == 1:
+                raise RuntimeError("transient")
+
+        with patch("openkb.cli.time.sleep"):
+            _run_compile_with_retry(_factory, "compiling")
+        assert calls == 2
+
+    def test_other_exceptions_raise_after_exhausting_retries(self):
+        from openkb.cli import _run_compile_with_retry
+
+        calls = 0
+
+        async def _factory():
+            nonlocal calls
+            calls += 1
+            raise RuntimeError("still broken")
+
+        with patch("openkb.cli.time.sleep"), pytest.raises(RuntimeError, match="still broken"):
+            _run_compile_with_retry(_factory, "compiling")
+        assert calls == 2
+
+
 class TestQueryStreamGate:
     """Regression tests for issue #34.
 

tests/test_compiler.py

@@ -2138,6 +2138,31 @@ def _plan_response(self):
             }
         )
 
+    @staticmethod
+    def _flaky_once_acompletion():
+        """Fails the first time "concept-b" is generated, succeeds on any
+        later call — simulates a transient error that clears up by the time
+        the end-of-first-pass sweep retries it."""
+        attempts: dict[str, int] = {}
+
+        def _call(*args, **kwargs):
+            messages = kwargs.get("messages") or (args[1] if len(args) > 1 else [])
+            last_content = messages[-1]["content"] if messages else ""
+            if "concept-b" in last_content:
+                attempts["concept-b"] = attempts.get("concept-b", 0) + 1
+                if attempts["concept-b"] == 1:
+                    raise RuntimeError("boom: transient concept-b failure")
+            mock_resp = MagicMock()
+            mock_resp.choices = [MagicMock()]
+            mock_resp.choices[0].message.content = json.dumps(
+                {"brief": "b", "content": "# Concept\n\nBody."}
+            )
+            mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50)
+            mock_resp.usage.prompt_tokens_details = None
+            return mock_resp
+
+        return _call
+
     @pytest.mark.asyncio
     async def test_normal_mode_keeps_partial_success_silent(self, tmp_path):
         """Default/omitted insert_mode: unchanged behavior — no exception,
@@ -2244,6 +2269,59 @@ async def test_fail_at_end_no_failures_does_not_raise(self, tmp_path):
             )
         assert (wiki / "concepts" / "concept-a.md").exists()
 
+    @pytest.mark.asyncio
+    async def test_normal_mode_sweep_recovers_transient_failure(self, tmp_path):
+        """A concept that fails its first attempt but succeeds on retry gets
+        picked up by the end-of-first-pass sweep (a tier above the low-level
+        _llm_call retry) and ends up written, with no exception raised."""
+        wiki = self._setup_wiki(tmp_path)
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.completion = MagicMock(
+                side_effect=_mock_completion([self._plan_response()])
+            )
+            mock_litellm.acompletion = AsyncMock(side_effect=self._flaky_once_acompletion())
+            await _compile_concepts(
+                wiki,
+                tmp_path,
+                "gpt-4o-mini",
+                {"role": "system", "content": "s"},
+                {"role": "user", "content": "d"},
+                "summary",
+                "test-doc",
+                5,
+                insert_mode="normal",
+            )
+        assert (wiki / "concepts" / "concept-a.md").exists()
+        assert (wiki / "concepts" / "concept-b.md").exists()
+
+    @pytest.mark.asyncio
+    async def test_fail_fast_does_not_get_a_sweep_retry(self, tmp_path):
+        """insert_mode="fail-fast" aborts on the first failure without
+        waiting for the rest of the batch, so it never reaches the sweep —
+        even a failure that would have cleared up on retry still aborts."""
+        from openkb.agent.compiler import ConceptCompilationError
+
+        wiki = self._setup_wiki(tmp_path)
+        with patch("openkb.agent.compiler.litellm") as mock_litellm:
+            mock_litellm.completion = MagicMock(
+                side_effect=_mock_completion([self._plan_response()])
+            )
+            mock_litellm.acompletion = AsyncMock(side_effect=self._flaky_once_acompletion())
+            with pytest.raises(ConceptCompilationError):
+                await _compile_concepts(
+                    wiki,
+                    tmp_path,
+                    "gpt-4o-mini",
+                    {"role": "system", "content": "s"},
+                    {"role": "user", "content": "d"},
+                    "summary",
+                    "test-doc",
+                    5,
+                    insert_mode="fail-fast",
+                )
+        assert not (wiki / "concepts" / "concept-a.md").exists()
+        assert not (wiki / "concepts" / "concept-b.md").exists()
+
     @pytest.mark.asyncio
     async def test_compile_short_doc_reads_insert_mode_from_kb_config(self, tmp_path):
         """End-to-end wiring check: compile_short_doc reads insert_mode from