fix(llm): enhance fallback logic to handle empty completions and impr… (#5491)

* fix(llm): enhance fallback logic to handle empty completions and improve metadata handling

* fix(llm): stream tool call deltas immediately

Author: Abhiiishek44

src/llm_core.py

@@ -2769,8 +2769,9 @@ async def stream_llm_with_fallback(candidates, messages, **kwargs):
     """Wrap stream_llm with an ordered fallback chain.
 
     `candidates` is a list of (url, model, headers). Each is tried in order,
-    but only retried on a *pre-content* failure — i.e. an ``event: error``
-    that arrives before any assistant text / tool-call data has been yielded.
+    but only retried on a *pre-content* failure — an ``event: error`` or an
+    empty completion before any assistant text / completed tool call is yielded.
+    Metadata is held until substantive output commits the candidate.
     Once a candidate has emitted real output we never switch (that would
     duplicate streamed tokens); a later error from that candidate passes
     through unchanged. The dead-host cooldown in stream_llm makes repeat
@@ -2789,6 +2790,7 @@ async def stream_llm_with_fallback(candidates, messages, **kwargs):
         is_last = (i == len(cands) - 1)
         emitted = False
         retried = False
+        pending_metadata = []
         async for chunk in stream_llm(url, model, messages, headers=headers, **kwargs):
             if chunk.startswith("event: error"):
                 if not emitted and not is_last:
@@ -2801,34 +2803,67 @@ async def stream_llm_with_fallback(candidates, messages, **kwargs):
                     else:
                         logger.warning(f"[fallback] candidate {model} failed; trying next")
                     break
+                if not emitted:
+                    # A last-candidate error is already the clearest terminal
+                    # result; do not append an empty-completion error as well.
+                    yield chunk
+                    return
                 yield chunk
                 continue
-            # Any data chunk other than the terminal [DONE] means real output.
-            if chunk.startswith("data: ") and not chunk.startswith("data: [DONE]"):
+
+            event_data = {}
+            is_done = chunk.startswith("data: [DONE]")
+            if chunk.startswith("data: ") and not is_done:
                 try:
                     event_data = json.loads(chunk[6:])
                 except Exception:
-                    event_data = {}
-                if event_data.get("type") == "model_actual":
-                    yield chunk
-                    continue
+                    pass
+
+            delta = event_data.get("delta")
+            event_type = event_data.get("type")
+            substantive = (
+                isinstance(delta, str) and bool(delta)
+            ) or (
+                event_type == "tool_call_delta"
+            ) or (
+                event_type == "tool_calls"
+                and bool(event_data.get("calls"))
+            )
+
+            if substantive and not emitted:
                 # First real output from a NON-primary candidate: tell the client
                 # the selected model failed and another answered. Without this the
                 # fallback is invisible — a misconfigured provider looks like it
                 # works because the reply is shown under the originally selected
                 # model's name (e.g. a Bedrock/Claude endpoint that 400s every
                 # request but appears fine because another model silently answered).
-                if not emitted and i > 0:
+                if i > 0:
                     yield ('data: ' + json.dumps({
                         "type": "fallback",
                         "selected_model": primary_model,
                         "answered_by": model,
                         "reason": _summarize_stream_error(last_error),
                     }) + '\n\n')
+                # Metadata must not commit a candidate. Once real output arrives,
+                # flush it after any fallback notice and before the output itself.
+                for metadata_chunk in pending_metadata:
+                    yield metadata_chunk
+                pending_metadata.clear()
                 emitted = True
-            yield chunk
-        if not retried:
-            return  # candidate finished (success, or terminal error already sent)
-    # Every candidate failed pre-content — surface the last error.
-    if last_error:
-        yield last_error
+
+            if substantive or emitted:
+                yield chunk
+            elif not is_done:
+                pending_metadata.append(chunk)
+
+        if emitted:
+            return
+        if retried:
+            continue
+        if not is_last:
+            last_error = f'event: error\ndata: {json.dumps({"error": f"Model {model} returned no substantive output", "status": 502})}\n\n'
+            tag = "primary" if i == 0 else "candidate"
+            logger.warning(f"[fallback] {tag} {model} returned no substantive output; trying next")
+            continue
+        yield f'event: error\ndata: {json.dumps({"error": "All model candidates returned no substantive output", "status": 502})}\n\n'
+        return

tests/test_llm_core_fallback.py

@@ -8,6 +8,8 @@
 import json
 import asyncio
 
+import pytest
+
 from src import llm_core
 
 
@@ -55,6 +57,145 @@ def per_model(model):
     assert not any('"fallback"' in c for c in chunks)
 
 
+def test_done_only_primary_invokes_fallback(monkeypatch):
+    calls = []
+
+    def per_model(model):
+        calls.append(model)
+        if model == "primary":
+            return ["data: [DONE]\n\n"]
+        return [
+            'data: {"type": "model_actual", "requested_model": "backup", "model": "backup-v2"}\n\n',
+            'data: {"delta": "backup answer"}\n\n',
+            "data: [DONE]\n\n",
+        ]
+
+    chunks = _run_fallback(monkeypatch, per_model)
+    assert calls == ["primary", "backup"]
+    assert any('"delta": "backup answer"' in c for c in chunks)
+    model_idx = next(i for i, c in enumerate(chunks) if '"model_actual"' in c)
+    fallback_idx = next(i for i, c in enumerate(chunks) if '"fallback"' in c)
+    answer_idx = next(i for i, c in enumerate(chunks) if '"delta": "backup answer"' in c)
+    assert fallback_idx < model_idx < answer_idx
+
+
+def test_usage_then_done_primary_invokes_fallback_and_discards_usage(monkeypatch):
+    calls = []
+
+    def per_model(model):
+        calls.append(model)
+        if model == "primary":
+            return [
+                'data: {"type": "usage", "data": {"input_tokens": 4, "output_tokens": 0}}\n\n',
+                "data: [DONE]\n\n",
+            ]
+        return ['data: {"delta": "backup answer"}\n\n', "data: [DONE]\n\n"]
+
+    chunks = _run_fallback(monkeypatch, per_model)
+    assert calls == ["primary", "backup"]
+    assert not any('"type": "usage"' in c for c in chunks)
+
+
+@pytest.mark.parametrize(
+    "output_chunk",
+    [
+        'data: {"delta": "visible text"}\n\n',
+        'data: {"delta": "reasoning", "thinking": true}\n\n',
+    ],
+)
+def test_text_or_reasoning_output_prevents_fallback(monkeypatch, output_chunk):
+    calls = []
+
+    def per_model(model):
+        calls.append(model)
+        return [output_chunk, "data: [DONE]\n\n"]
+
+    chunks = _run_fallback(monkeypatch, per_model)
+    assert calls == ["primary"]
+    assert output_chunk in chunks
+    assert not any('"fallback"' in c for c in chunks)
+
+
+def test_whitespace_only_delta_prevents_fallback(monkeypatch):
+    calls = []
+    whitespace = 'data: {"delta": "   "}\n\n'
+
+    def per_model(model):
+        calls.append(model)
+        return [whitespace, "data: [DONE]\n\n"]
+
+    chunks = _run_fallback(monkeypatch, per_model)
+    assert calls == ["primary"]
+    assert whitespace in chunks
+    assert not any('"fallback"' in c for c in chunks)
+
+
+def test_completed_tool_call_output_prevents_fallback(monkeypatch):
+    calls = []
+    tool_calls = 'data: {"type": "tool_calls", "calls": [{"id": "c1", "name": "bash", "arguments": "{}"}]}\n\n'
+
+    def per_model(model):
+        calls.append(model)
+        return [tool_calls, "data: [DONE]\n\n"]
+
+    chunks = _run_fallback(monkeypatch, per_model)
+    assert calls == ["primary"]
+    assert tool_calls in chunks
+    assert not any('"fallback"' in c for c in chunks)
+
+
+def test_tool_call_delta_is_forwarded_immediately_and_prevents_fallback(monkeypatch):
+    calls = []
+    advanced_past_delta = False
+    tool_delta = 'data: {"type": "tool_call_delta", "index": 0, "arg_delta": "{\\"path\\":"}\n\n'
+    tool_calls = 'data: {"type": "tool_calls", "calls": [{"id": "c1", "name": "write_file", "arguments": "{\\"path\\":\\"x\\"}"}]}\n\n'
+
+    async def fake_stream(url, model, messages, **kw):
+        nonlocal advanced_past_delta
+        calls.append(model)
+        yield tool_delta
+        advanced_past_delta = True
+        yield tool_calls
+        yield "data: [DONE]\n\n"
+
+    monkeypatch.setattr(llm_core, "stream_llm", fake_stream)
+
+    async def run():
+        stream = llm_core.stream_llm_with_fallback(
+            [("u1", "primary", {}), ("u2", "backup", {})],
+            [{"role": "user", "content": "hi"}],
+        )
+        first = await anext(stream)
+        assert first == tool_delta
+        assert not advanced_past_delta
+        chunks = [first]
+        async for chunk in stream:
+            chunks.append(chunk)
+        return chunks
+
+    chunks = asyncio.run(run())
+    assert calls == ["primary"]
+    assert tool_calls in chunks
+    assert not any('"type": "fallback"' in c for c in chunks)
+
+
+def test_empty_final_candidate_surfaces_terminal_error(monkeypatch):
+    calls = []
+
+    def per_model(model):
+        calls.append(model)
+        if model == "primary":
+            return []  # clean EOF without substantive output
+        return ["data: [DONE]\n\n"]
+
+    chunks = _run_fallback(monkeypatch, per_model)
+    assert calls == ["primary", "backup"]
+    errors = [c for c in chunks if c.startswith("event: error")]
+    assert len(errors) == 1
+    assert "All model candidates returned no substantive output" in errors[0]
+    assert '"status": 502' in errors[0]
+
+
 def test_dedupe_candidates_keeps_first_of_each_route():
     """(url, model) is the route key; later repeats are dropped, order preserved,
     the first tuple (with its headers) kept, malformed entries filtered."""