fix(agentic): tolerate LLM call failures in attempt_answer and ensure trace completion

apps/api/tests/contract/test_agentic_answer_policy_contract.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from shared.services.retrieval.agentic.budget import BudgetExceeded
 from shared.services.retrieval.agentic.policy import attempt_answer
 from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState
 
@@ -10,6 +11,14 @@ async def _malformed_json_wrapper(_prompt: str) -> str:
     return '{"status": "DONE", "answer": "truncated"'
 
 
+async def _raise_runtime_error(_prompt: str) -> str:
+    raise RuntimeError("simulated LLM failure")
+
+
+async def _raise_budget_exceeded(_prompt: str) -> str:
+    raise BudgetExceeded("budget exhausted")
+
+
 @pytest.mark.asyncio
 async def test_attempt_answer_should_not_expose_malformed_json_wrapper() -> None:
     status, answer, reason = await attempt_answer(
@@ -23,3 +32,53 @@ async def test_attempt_answer_should_not_expose_malformed_json_wrapper() -> None
     assert status == "NOT_FOUND"
     assert answer == ""
     assert reason == "attempt_answer returned malformed JSON"
+
+
+@pytest.mark.asyncio
+async def test_attempt_answer_text_llm_error_returns_not_found() -> None:
+    """Text LLM failures should return NOT_FOUND instead of crashing."""
+    status, answer, reason = await attempt_answer(
+        _raise_runtime_error,
+        query="What changed?",
+        evidence_text="┈ evidence",
+        state=AgentState(),
+        config=AgentRunConfig(),
+    )
+
+    assert status == "NOT_FOUND"
+    assert answer == ""
+    assert "LLM error" in reason
+    assert "simulated LLM failure" in reason
+
+
+@pytest.mark.asyncio
+async def test_attempt_answer_vlm_and_fallback_both_fail_returns_not_found() -> None:
+    """When VLM fails and the text LLM fallback also fails, return NOT_FOUND."""
+    status, answer, reason = await attempt_answer(
+        _raise_runtime_error,  # llm_fn — used as the text fallback
+        query="What changed?",
+        evidence_text="┈ evidence",
+        state=AgentState(),
+        config=AgentRunConfig(),
+        vlm_fn=_raise_runtime_error,
+        image_urls=["https://example.com/test.png"],
+    )
+
+    assert status == "NOT_FOUND"
+    assert answer == ""
+    assert "both failed" in reason
+
+
+@pytest.mark.asyncio
+async def test_attempt_answer_budget_exceeded_in_fallback_propagates() -> None:
+    """BudgetExceeded during the text LLM fallback must propagate, not be swallowed."""
+    with pytest.raises(BudgetExceeded):
+        await attempt_answer(
+            _raise_budget_exceeded,  # llm_fn — used as the text fallback
+            query="What changed?",
+            evidence_text="┈ evidence",
+            state=AgentState(),
+            config=AgentRunConfig(),
+            vlm_fn=_raise_runtime_error,
+            image_urls=["https://example.com/test.png"],
+        )

packages/shared-python/shared/services/retrieval/agentic/orchestrator.py

@@ -843,6 +843,8 @@ async def _context_llm_call(prompt):
         revision_hint: str | None = None
         stop_reason = 'max_revisions'
         failure_reason = ''
+        all_refs: list[dict[str, str]] = []
+        router_used = 'agentic_discovery_only'
 
         for round_idx in range(config.max_revisions + 1):
             if state.elapsed_ms >= config.latency_budget_ms:
@@ -1224,6 +1226,13 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
                 status, answer_text, reason = 'NOT_FOUND', '', 'context budget exhausted'
                 stop_reason = 'context_budget'
                 break
+            except Exception as exc:
+                logger.warning(f'  agentic: attempt_answer failed: {exc}')
+                if trace_enabled:
+                    trace.record_budget_stop('llm_error')
+                status, answer_text, reason = 'NOT_FOUND', '', f'LLM error: {exc}'
+                stop_reason = 'llm_error'
+                break
             state.step_count += 1
 
             if trace_enabled:
@@ -1315,29 +1324,63 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
         # ══════════════════════════════════════════════════════════════════
         # Final Assembly
         # ══════════════════════════════════════════════════════════════════
-        router_used = (
-            'agentic_llm' if any(t.has_content() for t in state.doc_trees.values())
-            else 'agentic_discovery_only'
-        )
+        try:
+            router_used = (
+                'agentic_llm' if any(t.has_content() for t in state.doc_trees.values())
+                else 'agentic_discovery_only'
+            )
 
-        # Collect referenced chunk IDs from all doc trees
-        all_refs: list[dict[str, str]] = []
-        seen_ref_ids: set[str] = set()
-        for doc_id, doc_tree in state.doc_trees.items():
-            doc_name = state.doc_id_to_name.get(doc_id, doc_id)
-            for ref in doc_tree.collect_referenced_ids(document_name=doc_name):
-                cid = ref.get('chunk_id', '')
-                if cid and cid not in seen_ref_ids:
-                    seen_ref_ids.add(cid)
-                    all_refs.append(ref)
-
-        # Re-render final evidence (may have been updated in last revision)
-        if not evidence_text or evidence_text == '(no evidence collected)':
-            evidence_text = await _render_evidence(
-                db,
-                state.doc_trees, state.doc_id_to_name,
+            # Collect referenced chunk IDs from all doc trees
+            all_refs.clear()
+            seen_ref_ids: set[str] = set()
+            for doc_id, doc_tree in state.doc_trees.items():
+                doc_name = state.doc_id_to_name.get(doc_id, doc_id)
+                for ref in doc_tree.collect_referenced_ids(document_name=doc_name):
+                    cid = ref.get('chunk_id', '')
+                    if cid and cid not in seen_ref_ids:
+                        seen_ref_ids.add(cid)
+                        all_refs.append(ref)
+
+            # Re-render final evidence (may have been updated in last revision)
+            if not evidence_text or evidence_text == '(no evidence collected)':
+                evidence_text = await _render_evidence(
+                    db,
+                    state.doc_trees, state.doc_id_to_name,
+                )
+
+            result = AgenticResult(
+                evidence_text=evidence_text,
+                answer_text=answer_text,
+                referenced_chunks=all_refs,
+                router_used=router_used,
+                budget_snapshot=state.ledger.snapshot() if state.ledger else None,
+                stop_reason=stop_reason,
+                failure_reason=failure_reason,
+            )
+
+            logger.info(
+                f'agentic retrieval DONE: {len(all_refs)} referenced chunks, '
+                f'evidence_text={len(evidence_text)} chars, '
+                f'answer_text={len(answer_text)} chars, '
+                f'router={router_used}, steps={state.step_count}, '
+                f'stop_reason={stop_reason}, revisions={state.revision_count}, '
+                f'{state.elapsed_ms}ms'
             )
 
+            return result
+        except Exception as exc:
+            logger.error(f'agentic retrieval failed: {exc}')
+            stop_reason = 'llm_error'
+            failure_reason = str(exc)
+        finally:
+            if trace_enabled:
+                await trace.complete(
+                    all_refs,
+                    router_used,
+                    budget_snapshot=state.ledger.snapshot() if state.ledger else None,
+                )
+
+        # Fallback result when an unexpected exception escapes the pipeline
         result = AgenticResult(
             evidence_text=evidence_text,
             answer_text=answer_text,
@@ -1349,19 +1392,12 @@ async def vlm_context_call(prompt, _vlm_fn=vlm_fn):
         )
 
         logger.info(
-            f'agentic retrieval DONE: {len(all_refs)} referenced chunks, '
+            f'agentic retrieval DONE (degraded): {len(all_refs)} referenced chunks, '
             f'evidence_text={len(evidence_text)} chars, '
             f'answer_text={len(answer_text)} chars, '
             f'router={router_used}, steps={state.step_count}, '
             f'stop_reason={stop_reason}, revisions={state.revision_count}, '
             f'{state.elapsed_ms}ms'
         )
 
-        if trace_enabled:
-            await trace.complete(
-                all_refs,
-                router_used,
-                budget_snapshot=state.ledger.snapshot() if state.ledger else None,
-            )
-
         return result

packages/shared-python/shared/services/retrieval/agentic/policy.py

@@ -165,9 +165,18 @@ async def attempt_answer(
         raise
     except Exception as exc:
         if effective_fn is llm_fn:
-            raise
+            logger.warning(f'  [attempt_answer] text LLM failed: {exc}')
+            return 'NOT_FOUND', '', f'LLM error: {exc}'
         logger.warning(f'  [attempt_answer] VLM failed, falling back to text LLM: {exc}')
-        raw_response = await llm_fn(prompt_text)
+        try:
+            raw_response = await llm_fn(prompt_text)
+        except BudgetExceeded:
+            raise
+        except Exception as fallback_exc:
+            logger.warning(
+                f'  [attempt_answer] text LLM fallback also failed: {fallback_exc}'
+            )
+            return 'NOT_FOUND', '', f'VLM and text LLM both failed: {fallback_exc}'
     logger.info(f'  [attempt_answer] raw={repr(raw_response[:300])}')
 
     if verbose: