"""The route engine: which entry answers, what a failure means, when we go back.
The behaviour under test is the one production lacked: a chain of
``model@target`` entries where a cloud failure moves the very next call to the
local model, quota exhaustion is a failover (not a dead turn), and the next
user turn returns to the primary once the background probe says it is healthy.
"""
import asyncio
import pytest
from navi.config import Settings
from navi.exceptions import (
LLMAuthError,
LLMBackendError,
LLMConnectionError,
LLMModelNotFoundError,
LLMQuotaExceededError,
LLMServerUnavailableError,
)
from navi.llm import fallback as fb
from navi.llm.base import LLMChunk, LLMResponse, Message
from navi.llm.config import LLMConfig, Route, RouteEntry, Target
from navi.llm.fallback import FallbackOllamaBackend
class FakeOllama:
"""Stands in for the per-target OllamaBackend."""
def __init__(self, error: Exception | None = None, fail_times: int = 0):
self.calls = 0
self.fail_times = fail_times
self.error = error or LLMConnectionError("connection refused")
self.seen: list[dict] = []
async def complete(self, messages, **kwargs):
self.calls += 1
self.seen.append(kwargs)
if self.calls <= self.fail_times:
raise self.error
return LLMResponse(
content="ok", tool_calls=None, finish_reason="stop", model=kwargs.get("model")
)
async def stream_complete(self, messages, **kwargs):
self.calls += 1
self.seen.append(kwargs)
if self.calls <= self.fail_times:
raise self.error
yield LLMChunk(delta="hello", model=kwargs.get("model"))
async def embed(self, texts, model=None, **kwargs):
self.calls += 1
self.seen.append({"model": model})
if self.calls <= self.fail_times:
raise self.error
return [[0.1] * 4 for _ in texts]
def build(clients: dict[str, FakeOllama], chain, *, route_id: str = "default"):
"""A backend over named targets, with one route through them.
``chain`` is a list of ``(model, target_name)`` or
``(model, target_name, context_window, options)``.
"""
targets = {name: Target(name=name, host=f"http://{name}.test") for name in clients}
entries = []
for item in chain:
model, target_name = item[0], item[1]
ctx = item[2] if len(item) > 2 else None
opts = item[3] if len(item) > 3 else {}
entries.append(
RouteEntry(model=model, target=targets[target_name],
context_window=ctx, options=opts)
)
config = LLMConfig(
targets=targets,
routes={route_id: Route(route_id, tuple(entries))},
probe_interval_sec=60,
probe_timeout_sec=5,
)
backend = FallbackOllamaBackend(config)
backend._client = lambda entry: clients[entry.target.name] # type: ignore[assignment]
return backend
@pytest.fixture(autouse=True)
def _fast_retries(monkeypatch):
monkeypatch.setattr(
fb,
"settings",
Settings(_env_file=None, llm_complete_retries=2, llm_retry_backoff_sec=0.0),
)
def _msg():
return [Message(role="user", content="hi")]
class TestFailover:
async def test_a_dead_primary_moves_to_the_backup_without_retrying(self):
"""The old engine retried per server *and* per model: 2 × 90 s.
With a healthy alternative one step down the chain, a connection error
must cost exactly one failed attempt.
"""
cloud = FakeOllama(fail_times=99)
home = FakeOllama()
backend = build(
{"cloud": cloud, "home": home},
[("glm-5.3-flash:cloud", "cloud"), ("qwen3.6:35b", "home")],
)
resp = await backend.complete(_msg(), route="default")
assert resp.content == "ok"
assert cloud.calls == 1
assert home.calls == 1
assert resp.target == "home"
assert resp.model == "qwen3.6:35b"
async def test_quota_exhaustion_fails_over(self):
"""429 used to surface as a generic backend error and kill the turn."""
cloud = FakeOllama(error=LLMQuotaExceededError("429 rate limit"), fail_times=99)
home = FakeOllama()
backend = build(
{"cloud": cloud, "home": home},
[("glm-5.3-flash:cloud", "cloud"), ("qwen3.6:35b", "home")],
)
resp = await backend.complete(_msg(), route="default")
assert resp.target == "home"
assert backend.health.snapshot()["targets"]["cloud"]["state"] == "cooling"
assert backend.health.snapshot()["targets"]["cloud"]["reason"] == "quota"
async def test_a_bad_key_is_recorded_as_auth(self):
cloud = FakeOllama(error=LLMAuthError("401"), fail_times=99)
home = FakeOllama()
backend = build(
{"cloud": cloud, "home": home},
[("m1", "cloud"), ("m2", "home")],
)
await backend.complete(_msg(), route="default")
assert backend.health.snapshot()["targets"]["cloud"]["state"] == "auth"
async def test_a_5xx_cools_the_target_down(self):
cloud = FakeOllama(error=LLMServerUnavailableError("503"), fail_times=99)
home = FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
await backend.complete(_msg(), route="default")
snapshot = backend.health.snapshot()["targets"]["cloud"]
assert snapshot["state"] == "cooling"
assert snapshot["reason"] == "server"
async def test_a_missing_model_is_recorded_per_model(self):
cloud = FakeOllama(error=LLMModelNotFoundError("404"), fail_times=99)
home = FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
await backend.complete(_msg(), route="default")
snapshot = backend.health.snapshot()
assert snapshot["missing_models"][0]["model"] == "m1"
async def test_the_last_entry_still_retries(self):
"""With nothing left to fall back to, a transient error is worth retrying."""
home = FakeOllama(fail_times=1)
backend = build({"home": home}, [("m", "home")])
resp = await backend.complete(_msg(), route="default")
assert resp.content == "ok"
assert home.calls == 2
async def test_a_missing_model_is_not_retried(self):
home = FakeOllama(error=LLMModelNotFoundError("404"), fail_times=99)
backend = build({"home": home}, [("m", "home")])
with pytest.raises(LLMBackendError):
await backend.complete(_msg(), route="default")
assert home.calls == 1
async def test_everything_down_raises_with_the_last_error(self):
home = FakeOllama(fail_times=99)
cloud = FakeOllama(fail_times=99)
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
with pytest.raises(LLMBackendError) as exc:
await backend.complete(_msg(), route="default")
assert "exhausted" in str(exc.value)
assert home.calls == 2 # the last entry retried once
class TestSwitchBack:
async def test_the_next_turn_returns_to_a_recovered_primary(self):
cloud = FakeOllama(fail_times=1)
home = FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
first = await backend.complete(_msg(), route="default")
assert first.target == "home"
# The next turn starts on the backup: the cloud is still marked down.
backend.begin_turn("default")
again = await backend.complete(_msg(), route="default")
assert again.target == "home"
# The background probe would call report_success on the primary.
backend.health.reset()
backend.begin_turn("default")
recovered = await backend.complete(_msg(), route="default")
assert recovered.target == "cloud"
async def test_a_running_turn_is_not_moved_back_mid_way(self):
cloud = FakeOllama(fail_times=1)
home = FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
await backend.complete(_msg(), route="default")
backend.health.reset() # probe says the cloud is back
# No begin_turn: the in-flight turn keeps using the entry it started on.
second = await backend.complete(_msg(), route="default")
assert second.target == "home"
backend.begin_turn("default")
third = await backend.complete(_msg(), route="default")
assert third.target == "cloud"
class TestRouteParameters:
async def test_the_entry_context_window_is_sent_as_num_ctx(self):
home = FakeOllama()
backend = build({"home": home}, [("m", "home", 49152, {})])
await backend.complete(_msg(), route="default", num_ctx=200_000)
assert home.seen[0]["num_ctx"] == 49152
async def test_a_smaller_caller_window_is_kept(self):
home = FakeOllama()
backend = build({"home": home}, [("m", "home", 49152, {})])
await backend.complete(_msg(), route="default", num_ctx=8000)
assert home.seen[0]["num_ctx"] == 8000
async def test_per_entry_options_reach_the_model(self):
home = FakeOllama()
backend = build({"home": home}, [("m", "home", None, {"presence_penalty": 0.0})])
await backend.complete(_msg(), route="default")
assert home.seen[0]["extra_options"] == {"presence_penalty": 0.0}
async def test_the_active_entry_window_is_exposed(self):
cloud, home = FakeOllama(fail_times=99), FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home", 49152)])
assert backend.active_window("default") is None # before the first call: entry 0
await backend.complete(_msg(), route="default")
assert backend.active_window("default") == 49152
assert backend.serving_primary("default", "m2") is False
class TestStreaming:
async def test_the_first_chunk_is_stamped_with_model_and_target(self):
cloud = FakeOllama(fail_times=99)
home = FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
chunks = [c async for c in backend.stream_complete(_msg(), route="default")]
assert chunks[0].model == "m2"
assert chunks[0].target == "home"
async def test_the_same_model_on_a_second_host_still_counts_as_a_fallback(self):
"""A route may run one model on two machines. Comparing model names
alone called the second host's answer primary, so the UI never showed
the turn had failed over."""
cloud, home = FakeOllama(fail_times=99), FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m", "cloud"), ("m", "home")])
chunks = [c async for c in backend.stream_complete(_msg(), route="default")]
assert chunks[0].target == "home"
assert chunks[0].is_fallback is True
async def test_a_failure_before_the_first_chunk_fails_over(self):
cloud = FakeOllama(error=LLMQuotaExceededError("429"), fail_times=99)
home = FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
chunks = [c async for c in backend.stream_complete(_msg(), route="default")]
assert chunks[0].delta == "hello"
class TestEmbed:
async def test_embeddings_follow_their_route(self):
home = FakeOllama()
cloud = FakeOllama(fail_times=99)
backend = build(
{"cloud": cloud, "home": home},
[("nomic-embed-text", "home"), ("nomic-embed-text", "cloud")],
route_id="embedding",
)
vectors = await backend.embed(["hello"], route="embedding")
assert len(vectors) == 1
assert home.calls == 1
async def test_an_embedding_backup_is_used_when_the_first_is_down(self):
home = FakeOllama(fail_times=99)
cloud = FakeOllama()
backend = build(
{"cloud": cloud, "home": home},
[("nomic-embed-text", "home"), ("nomic-embed-text", "cloud")],
route_id="embedding",
)
vectors = await backend.embed(["hello"], route="embedding")
assert len(vectors) == 1
assert cloud.calls == 1
class TestLegacyModelList:
"""Without a route the old "list of model names" path must still work."""
async def test_a_model_list_resolves_each_name_to_its_configured_target(self):
cloud = FakeOllama(fail_times=99)
home = FakeOllama()
backend = build(
{"cloud": cloud, "home": home},
[("cloud-only", "cloud"), ("local-only", "home")],
)
resp = await backend.complete(_msg(), model=["cloud-only", "local-only"])
assert resp.model == "local-only"
assert resp.target == "home"
assert cloud.calls == 1
async def test_an_unknown_model_goes_to_the_first_target(self):
only = FakeOllama()
backend = build({"only": only}, [("known", "only")])
resp = await backend.complete(_msg(), model=["never-heard-of-it"])
assert resp.model == "never-heard-of-it"
assert resp.target == "only"
class TestHealthSnapshot:
async def test_a_clean_backend_reports_ok(self):
home = FakeOllama()
backend = build({"home": home}, [("m", "home")])
await backend.complete(_msg(), route="default")
assert backend.health.snapshot()["targets"]["home"]["state"] == "ok"
assert backend.active_report()["default"]["primary"] is True
async def test_reset_forgets_every_failure(self):
cloud, home = FakeOllama(fail_times=99), FakeOllama()
backend = build({"cloud": cloud, "home": home}, [("m1", "cloud"), ("m2", "home")])
await backend.complete(_msg(), route="default")
assert backend.health.snapshot()["targets"]["cloud"]["state"] == "down"
backend.reset_health()
assert backend.health.snapshot()["targets"]["cloud"]["state"] == "ok"
assert backend.active_entry("default").model == "m1"
class TestHealthStates:
def test_an_entry_is_skipped_until_its_retry_time(self):
home = FakeOllama()
backend = build({"home": home}, [("m", "home")])
entry = backend.config.routes["default"].entries[0]
backend.health.report_failure(entry, LLMConnectionError("refused"))
assert backend.health.entry_usable(entry, now=0.0) is False
def test_a_healthy_entry_is_usable(self):
home = FakeOllama()
backend = build({"home": home}, [("m", "home")])
entry = backend.config.routes["default"].entries[0]
assert backend.health.entry_usable(entry) is True
