From 32262455b460f51db5373eae3294b90596d2c7fb Mon Sep 17 00:00:00 2001
From: Peter Pan <Peter.Pan@daocloud.io>
Date: Thu, 11 Jun 2026 18:57:34 +0800
Subject: [PATCH 1/3] feat(router): add openai-compatible llm judge router

---
 custom_routers/llmjudgerouter/__init__.py |   1 +
 custom_routers/llmjudgerouter/config.yaml |  12 ++
 custom_routers/llmjudgerouter/router.py   | 187 ++++++++++++++++++++++
 llmrouter/serve/config.py                 |   3 +
 llmrouter/serve/server.py                 |  40 ++++-
 tests/test_llmjudgerouter.py              |  72 +++++++++
 6 files changed, 310 insertions(+), 5 deletions(-)
 create mode 100644 custom_routers/llmjudgerouter/__init__.py
 create mode 100644 custom_routers/llmjudgerouter/config.yaml
 create mode 100644 custom_routers/llmjudgerouter/router.py
 create mode 100644 tests/test_llmjudgerouter.py

diff --git a/custom_routers/llmjudgerouter/__init__.py b/custom_routers/llmjudgerouter/__init__.py
new file mode 100644
index 00000000..d72a2e62
--- /dev/null
+++ b/custom_routers/llmjudgerouter/__init__.py
@@ -0,0 +1 @@
+from .router import LLMJudgeRouter
diff --git a/custom_routers/llmjudgerouter/config.yaml b/custom_routers/llmjudgerouter/config.yaml
new file mode 100644
index 00000000..8161a332
--- /dev/null
+++ b/custom_routers/llmjudgerouter/config.yaml
@@ -0,0 +1,12 @@
+hparam:
+  small_model: "deepseek-flash"
+  large_model: "deepseek-pro"
+judge:
+  api_base: "http://127.0.0.1:11434/v1"
+  api_key: null
+  model: "qwen3:0.6b"
+  timeout_s: 5
+  max_tokens: 64
+  temperature: 0
+  reason_max_chars: 80
+  max_signals: 3
diff --git a/custom_routers/llmjudgerouter/router.py b/custom_routers/llmjudgerouter/router.py
new file mode 100644
index 00000000..0e9979f9
--- /dev/null
+++ b/custom_routers/llmjudgerouter/router.py
@@ -0,0 +1,187 @@
+import json
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+import httpx
+import torch.nn as nn
+
+from llmrouter.models.meta_router import MetaRouter
+
+
+class LLMJudgeRouter(MetaRouter):
+    def __init__(self, yaml_path: Optional[str] = None):
+        super().__init__(model=nn.Identity(), yaml_path=yaml_path)
+
+        hparam = (self.cfg or {}).get("hparam", {}) or {}
+        judge = (self.cfg or {}).get("judge", {}) or {}
+
+        self.small_model = hparam.get("small_model", "small-model")
+        self.large_model = hparam.get("large_model", "large-model")
+        self.judge_api_base = str(judge.get("api_base", "http://127.0.0.1:11434/v1")).rstrip("/")
+        self.judge_api_key = self._resolve_api_key(judge.get("api_key"))
+        self.judge_model = judge.get("model", "qwen3:0.6b")
+        self.timeout_s = float(judge.get("timeout_s", 5))
+        self.max_tokens = int(judge.get("max_tokens", 64))
+        self.temperature = float(judge.get("temperature", 0))
+        self.reason_max_chars = int(judge.get("reason_max_chars", 80))
+        self.max_signals = int(judge.get("max_signals", 3))
+
+    def _resolve_api_key(self, value: Optional[str]) -> Optional[str]:
+        if not isinstance(value, str):
+            return None
+        if value.startswith("${") and value.endswith("}"):
+            return os.environ.get(value[2:-1])
+        return value
+
+    def _extract_first_json_object(self, text: str) -> str:
+        start = text.find("{")
+        if start < 0:
+            raise ValueError("no json object found")
+        depth = 0
+        for i, ch in enumerate(text[start:], start=start):
+            if ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    return text[start:i + 1]
+        raise ValueError("unclosed json object")
+
+    def _compact_reason(self, text: str) -> Optional[str]:
+        one_line = " ".join((text or "").strip().split())
+        if not one_line:
+            return None
+        if len(one_line) <= self.reason_max_chars:
+            return one_line
+        return one_line[:self.reason_max_chars]
+
+    def _normalize_signals(self, signals: Any) -> Optional[List[str]]:
+        if not isinstance(signals, list):
+            return None
+        normalized = []
+        seen = set()
+        for item in signals:
+            if not isinstance(item, str):
+                continue
+            signal = " ".join(item.strip().split())
+            if not signal or signal in seen:
+                continue
+            seen.add(signal)
+            normalized.append(signal[:32])
+            if len(normalized) >= self.max_signals:
+                break
+        return normalized or None
+
+    def _extract_assistant_text(self, data: Dict[str, Any]) -> str:
+        choices = data.get("choices")
+        if not isinstance(choices, list) or not choices:
+            return ""
+        message = choices[0].get("message") or {}
+        content = message.get("content")
+        if isinstance(content, str):
+            return content.strip()
+        if isinstance(content, list):
+            parts = []
+            for item in content:
+                if isinstance(item, dict) and item.get("type") == "text":
+                    parts.append(item.get("text", ""))
+            return "".join(parts).strip()
+        return ""
+
+    def _judge(self, query: str) -> Dict[str, Any]:
+        # Keep the prompt short and generic so downstream users can tune it easily.
+        prompt = (
+            "You are a routing judge for a MaaS gateway.\n"
+            "Choose exactly one backend model for the user request.\n"
+            f"Available models: {self.small_model}, {self.large_model}.\n"
+            f"Prefer {self.small_model} for simple or routine requests.\n"
+            f"Choose {self.large_model} only when the request clearly needs stronger reasoning, planning, or reliability.\n"
+            f"If uncertain, prefer {self.small_model}.\n"
+            "Do not answer the user request.\n"
+            "Return only one JSON object with this schema:\n"
+            "{"
+            f"\"model\":\"{self.small_model}|{self.large_model}\","
+            "\"confidence\":0.0,"
+            "\"reason\":\"short routing reason\","
+            "\"signals\":[\"abstract_tag\"]"
+            "}\n"
+            "The reason must be short, non-sensitive, and must not include chain-of-thought.\n"
+            "Signals must be short abstract tags, not placeholders and not copied entities from the query."
+        )
+
+        headers = {"Content-Type": "application/json"}
+        if self.judge_api_key:
+            headers["Authorization"] = f"Bearer {self.judge_api_key}"
+
+        body = {
+            "model": self.judge_model,
+            "messages": [
+                {"role": "system", "content": prompt},
+                {"role": "user", "content": query},
+            ],
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "stream": False,
+        }
+
+        judge_start = time.perf_counter()
+        with httpx.Client(timeout=self.timeout_s) as client:
+            resp = client.post(
+                f"{self.judge_api_base}/chat/completions",
+                headers=headers,
+                json=body,
+            )
+            resp.raise_for_status()
+            data = resp.json()
+        judge_latency_ms = int((time.perf_counter() - judge_start) * 1000)
+
+        text = self._extract_assistant_text(data)
+        if not text:
+            raise ValueError("empty judge output")
+
+        try:
+            obj = json.loads(text)
+        except Exception:
+            obj = json.loads(self._extract_first_json_object(text))
+
+        if not isinstance(obj, dict):
+            raise ValueError("judge output is not a json object")
+
+        model = str(obj.get("model", "")).strip()
+        if model not in (self.small_model, self.large_model):
+            raise ValueError(f"invalid model from judge: {model}")
+
+        confidence = obj.get("confidence")
+        if isinstance(confidence, (int, float)) and 0.0 <= float(confidence) <= 1.0:
+            confidence = float(confidence)
+        else:
+            confidence = None
+
+        return {
+            "model": model,
+            "reason": self._compact_reason(obj.get("reason") or ""),
+            "signals": self._normalize_signals(obj.get("signals")),
+            "confidence": confidence,
+            "raw": text,
+            "judge_latency_ms": judge_latency_ms,
+        }
+
+    def route_single(self, query_input: Dict[str, Any]) -> Dict[str, Any]:
+        query = query_input.get("query", "") if isinstance(query_input, dict) else str(query_input)
+        judged = self._judge(query)
+        selected = judged["model"]
+        return {
+            "query": query,
+            "model_name": selected,
+            "predicted_llm": selected,
+            "predicted_llm_name": selected,
+            "method": "llm_judge",
+            "routing_reason": judged.get("reason"),
+            "routing_confidence": judged.get("confidence"),
+            "routing_signals": judged.get("signals"),
+            "routing_judge_latency_ms": judged.get("judge_latency_ms"),
+        }
+
+    def route_batch(self, batch):
+        return [self.route_single(item) for item in (batch or [])]
diff --git a/llmrouter/serve/config.py b/llmrouter/serve/config.py
index 0d316af4..20bcdb7c 100644
--- a/llmrouter/serve/config.py
+++ b/llmrouter/serve/config.py
@@ -43,6 +43,8 @@ class ServeConfig:
     # Show model name prefix
     show_model_prefix: bool = True
 
+    fail_on_routing_error: bool = False
+
     @classmethod
     def from_yaml(cls, yaml_path: str) -> "ServeConfig":
         """Load configuration from YAML file"""
@@ -59,6 +61,7 @@ def from_yaml(cls, yaml_path: str) -> "ServeConfig":
         config.host = serve_config.get("host", config.host)
         config.port = serve_config.get("port", config.port)
         config.show_model_prefix = serve_config.get("show_model_prefix", config.show_model_prefix)
+        config.fail_on_routing_error = serve_config.get("fail_on_routing_error", config.fail_on_routing_error)
 
         # Router settings
         router_config = data.get("router", {})
diff --git a/llmrouter/serve/server.py b/llmrouter/serve/server.py
index 6e96b03d..18435a35 100644
--- a/llmrouter/serve/server.py
+++ b/llmrouter/serve/server.py
@@ -58,10 +58,12 @@ class ChatRequest(BaseModel):
 class RouterAdapter:
     """LLMRouter adapter"""
 
-    def __init__(self, router_name: str, config_path: Optional[str] = None):
+    def __init__(self, router_name: str, config_path: Optional[str] = None, fail_on_error: bool = False):
         self.router_name = router_name
         self.config_path = config_path
+        self.fail_on_error = fail_on_error
         self.router = None
+        self.last_router_info = None
         self._load_router()
 
     def _load_router(self):
@@ -107,6 +109,13 @@ def route(self, query: str, available_models: List[str]) -> str:
         try:
             result = self.router.route_single({"query": query})
             model_name = result.get("model_name") or result.get("predicted_llm")
+            self.last_router_info = {
+                "method": result.get("method"),
+                "routing_reason": result.get("routing_reason"),
+                "routing_signals": result.get("routing_signals"),
+                "routing_confidence": result.get("routing_confidence"),
+                "routing_judge_latency_ms": result.get("routing_judge_latency_ms"),
+            }
 
             # Check if model is available
             if model_name in available_models:
@@ -121,7 +130,10 @@ def route(self, query: str, available_models: List[str]) -> str:
             return available_models[0]
 
         except Exception as e:
-            print(f"[Router] Error: {e}")
+            self.last_router_info = {"error": f"{type(e).__name__}: {e}"}
+            print(f"[Router] Error: {type(e).__name__}: {e}")
+            if self.fail_on_error:
+                raise
             return available_models[0]
 
 
@@ -232,7 +244,8 @@ def create_app(config: ServeConfig = None, config_path: str = None) -> FastAPI:
     # Initialize components
     router_adapter = RouterAdapter(
         router_name=config.router_name,
-        config_path=config.router_config_path
+        config_path=config.router_config_path,
+        fail_on_error=config.fail_on_routing_error
     )
     llm_backend = LLMBackend(config)
 
@@ -267,8 +280,25 @@ async def chat_completions(request: ChatRequest):
         # Select model
         available_models = list(config.llms.keys())
         if request.model == "auto" or request.model not in available_models:
-            selected_model = router_adapter.route(user_query, available_models)
-            print(f"[Router] Query: '{user_query[:50]}...' -> {selected_model}")
+            try:
+                selected_model = router_adapter.route(user_query, available_models)
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"routing failed: {type(e).__name__}: {e}")
+            info = router_adapter.last_router_info or {}
+            reason = info.get("routing_reason")
+            signals = info.get("routing_signals")
+            confidence = info.get("routing_confidence")
+            judge_ms = info.get("routing_judge_latency_ms")
+            extra = ""
+            if reason:
+                extra += f" reason={reason}"
+            if signals:
+                extra += f" signals={signals}"
+            if confidence is not None:
+                extra += f" confidence={confidence}"
+            if judge_ms is not None:
+                extra += f" judge_ms={judge_ms}"
+            print(f"[Router] Query: '{user_query[:50]}...' -> {selected_model}{extra}")
         else:
             selected_model = request.model
 
diff --git a/tests/test_llmjudgerouter.py b/tests/test_llmjudgerouter.py
new file mode 100644
index 00000000..1206527b
--- /dev/null
+++ b/tests/test_llmjudgerouter.py
@@ -0,0 +1,72 @@
+from unittest.mock import patch
+
+import torch.nn as nn
+
+from custom_routers.llmjudgerouter.router import LLMJudgeRouter
+
+
+def _build_router():
+    router = LLMJudgeRouter.__new__(LLMJudgeRouter)
+    nn.Module.__init__(router)
+    router.small_model = "deepseek-flash"
+    router.large_model = "deepseek-pro"
+    router.judge_api_base = "http://127.0.0.1:11434/v1"
+    router.judge_api_key = None
+    router.judge_model = "qwen3:0.6b"
+    router.timeout_s = 5
+    router.max_tokens = 64
+    router.temperature = 0
+    router.reason_max_chars = 80
+    router.max_signals = 3
+    return router
+
+
+def test_llm_judge_router_uses_openai_compatible_chat_completions():
+    router = _build_router()
+    captured = {}
+
+    class FakeResponse:
+        def raise_for_status(self):
+            return None
+
+        def json(self):
+            return {
+                "choices": [
+                    {
+                        "message": {
+                            "content": (
+                                '{"model":"deepseek-flash","confidence":0.82,'
+                                '"reason":"simple factual request","signals":["fact","brief"]}'
+                            )
+                        }
+                    }
+                ]
+            }
+
+    class FakeClient:
+        def __init__(self, timeout):
+            captured["timeout"] = timeout
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+        def post(self, url, headers=None, json=None):
+            captured["url"] = url
+            captured["headers"] = headers
+            captured["json"] = json
+            return FakeResponse()
+
+    with patch("custom_routers.llmjudgerouter.router.httpx.Client", FakeClient):
+        result = router.route_single({"query": "What is Goldbach's conjecture?"})
+
+    assert captured["url"] == "http://127.0.0.1:11434/v1/chat/completions"
+    assert captured["json"]["max_tokens"] == 64
+    assert captured["json"]["model"] == "qwen3:0.6b"
+    assert result["model_name"] == "deepseek-flash"
+    assert result["routing_confidence"] == 0.82
+    assert result["routing_reason"] == "simple factual request"
+    assert result["routing_signals"] == ["fact", "brief"]
+    assert isinstance(result["routing_judge_latency_ms"], int)

From 1d968cd79b8890932e6ed61fa7e0aa29a4b95c5a Mon Sep 17 00:00:00 2001
From: Peter Pan <Peter.Pan@daocloud.io>
Date: Thu, 11 Jun 2026 19:33:57 +0800
Subject: [PATCH 2/3] fix(router): add large-model fallback for judge budget
 risks

---
 custom_routers/llmjudgerouter/router.py | 53 ++++++++++++++++++-------
 tests/test_llmjudgerouter.py            | 28 +++++++++++++
 2 files changed, 67 insertions(+), 14 deletions(-)

diff --git a/custom_routers/llmjudgerouter/router.py b/custom_routers/llmjudgerouter/router.py
index 0e9979f9..ca3a8ac8 100644
--- a/custom_routers/llmjudgerouter/router.py
+++ b/custom_routers/llmjudgerouter/router.py
@@ -26,6 +26,9 @@ def __init__(self, yaml_path: Optional[str] = None):
         self.temperature = float(judge.get("temperature", 0))
         self.reason_max_chars = int(judge.get("reason_max_chars", 80))
         self.max_signals = int(judge.get("max_signals", 3))
+        self.fallback_to_large_on_judge_error = bool(judge.get("fallback_to_large_on_judge_error", True))
+        self.prompt_budget_chars_per_token = int(judge.get("prompt_budget_chars_per_token", 4))
+        self.prompt_budget_output_buffer = int(judge.get("prompt_budget_output_buffer", 128))
 
     def _resolve_api_key(self, value: Optional[str]) -> Optional[str]:
         if not isinstance(value, str):
@@ -92,23 +95,34 @@ def _extract_assistant_text(self, data: Dict[str, Any]) -> str:
     def _judge(self, query: str) -> Dict[str, Any]:
         # Keep the prompt short and generic so downstream users can tune it easily.
         prompt = (
-            "You are a routing judge for a MaaS gateway.\n"
-            "Choose exactly one backend model for the user request.\n"
-            f"Available models: {self.small_model}, {self.large_model}.\n"
-            f"Prefer {self.small_model} for simple or routine requests.\n"
-            f"Choose {self.large_model} only when the request clearly needs stronger reasoning, planning, or reliability.\n"
-            f"If uncertain, prefer {self.small_model}.\n"
-            "Do not answer the user request.\n"
-            "Return only one JSON object with this schema:\n"
+            "You are a MaaS routing judge.\n"
+            f"Pick exactly one model: {self.small_model} or {self.large_model}.\n"
+            f"Use {self.small_model} for simple requests.\n"
+            f"Use {self.large_model} only for clearly harder reasoning or reliability needs.\n"
+            f"If unsure, pick {self.small_model}.\n"
+            "Do not answer the user.\n"
+            "Return JSON only:\n"
             "{"
-            f"\"model\":\"{self.small_model}|{self.large_model}\","
+            f"\"model\":\"{self.small_model}\","
             "\"confidence\":0.0,"
-            "\"reason\":\"short routing reason\","
-            "\"signals\":[\"abstract_tag\"]"
+            "\"reason\":\"short\","
+            "\"signals\":[\"tag\"]"
             "}\n"
-            "The reason must be short, non-sensitive, and must not include chain-of-thought.\n"
-            "Signals must be short abstract tags, not placeholders and not copied entities from the query."
+            f"The model field must be exactly {self.small_model} or {self.large_model}."
         )
+        estimated_prompt_tokens = max(
+            1,
+            (len(prompt) + len(query)) // max(1, self.prompt_budget_chars_per_token),
+        )
+        if estimated_prompt_tokens > self.prompt_budget_output_buffer:
+            return {
+                "model": self.large_model,
+                "reason": "judge_budget_risk",
+                "signals": ["judge_budget"],
+                "confidence": None,
+                "raw": None,
+                "judge_latency_ms": 0,
+            }
 
         headers = {"Content-Type": "application/json"}
         if self.judge_api_key:
@@ -169,7 +183,18 @@ def _judge(self, query: str) -> Dict[str, Any]:
 
     def route_single(self, query_input: Dict[str, Any]) -> Dict[str, Any]:
         query = query_input.get("query", "") if isinstance(query_input, dict) else str(query_input)
-        judged = self._judge(query)
+        try:
+            judged = self._judge(query)
+        except Exception:
+            if not self.fallback_to_large_on_judge_error:
+                raise
+            judged = {
+                "model": self.large_model,
+                "reason": "judge_error_fallback",
+                "signals": ["judge_error"],
+                "confidence": None,
+                "judge_latency_ms": None,
+            }
         selected = judged["model"]
         return {
             "query": query,
diff --git a/tests/test_llmjudgerouter.py b/tests/test_llmjudgerouter.py
index 1206527b..1d53e544 100644
--- a/tests/test_llmjudgerouter.py
+++ b/tests/test_llmjudgerouter.py
@@ -18,6 +18,9 @@ def _build_router():
     router.temperature = 0
     router.reason_max_chars = 80
     router.max_signals = 3
+    router.fallback_to_large_on_judge_error = True
+    router.prompt_budget_chars_per_token = 4
+    router.prompt_budget_output_buffer = 128
     return router
 
 
@@ -70,3 +73,28 @@ def post(self, url, headers=None, json=None):
     assert result["routing_reason"] == "simple factual request"
     assert result["routing_signals"] == ["fact", "brief"]
     assert isinstance(result["routing_judge_latency_ms"], int)
+
+
+def test_llm_judge_router_uses_large_model_when_prompt_budget_is_risky():
+    router = _build_router()
+    router.max_tokens = 64
+    router.prompt_budget_output_buffer = 4
+
+    result = router.route_single({"query": "x" * 100})
+
+    assert result["model_name"] == "deepseek-pro"
+    assert result["routing_reason"] == "judge_budget_risk"
+    assert result["routing_signals"] == ["judge_budget"]
+    assert result["routing_judge_latency_ms"] == 0
+
+
+def test_llm_judge_router_falls_back_to_large_model_on_judge_error():
+    router = _build_router()
+
+    with patch.object(LLMJudgeRouter, "_judge", side_effect=ValueError("empty judge output")):
+        result = router.route_single({"query": "证明哥德巴赫猜想"})
+
+    assert result["model_name"] == "deepseek-pro"
+    assert result["routing_reason"] == "judge_error_fallback"
+    assert result["routing_signals"] == ["judge_error"]
+    assert result["routing_judge_latency_ms"] is None

From ae9b9fe3f655e48ea840298206300eecd126e5b4 Mon Sep 17 00:00:00 2001
From: Peter Pan <Peter.Pan@daocloud.io>
Date: Thu, 11 Jun 2026 19:39:04 +0800
Subject: [PATCH 3/3] fix(router): restore dynamic judge schema hints

---
 custom_routers/llmjudgerouter/router.py | 16 +++++++---------
 tests/test_llmjudgerouter.py            |  7 +++----
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/custom_routers/llmjudgerouter/router.py b/custom_routers/llmjudgerouter/router.py
index ca3a8ac8..3adae535 100644
--- a/custom_routers/llmjudgerouter/router.py
+++ b/custom_routers/llmjudgerouter/router.py
@@ -101,20 +101,18 @@ def _judge(self, query: str) -> Dict[str, Any]:
             f"Use {self.large_model} only for clearly harder reasoning or reliability needs.\n"
             f"If unsure, pick {self.small_model}.\n"
             "Do not answer the user.\n"
-            "Return JSON only:\n"
-            "{"
-            f"\"model\":\"{self.small_model}\","
-            "\"confidence\":0.0,"
-            "\"reason\":\"short\","
-            "\"signals\":[\"tag\"]"
-            "}\n"
-            f"The model field must be exactly {self.small_model} or {self.large_model}."
+            "Return one JSON object only.\n"
+            f"The `model` field must be exactly `{self.small_model}` or `{self.large_model}`.\n"
+            "The `confidence` field must be a number between 0.0 and 1.0.\n"
+            "The `reason` field must be a short routing reason.\n"
+            "The `signals` field must be a list of 1-3 short abstract tags.\n"
+            "Do not include markdown fences or any extra text."
         )
         estimated_prompt_tokens = max(
             1,
             (len(prompt) + len(query)) // max(1, self.prompt_budget_chars_per_token),
         )
-        if estimated_prompt_tokens > self.prompt_budget_output_buffer:
+        if estimated_prompt_tokens > self.max_tokens:
             return {
                 "model": self.large_model,
                 "reason": "judge_budget_risk",
diff --git a/tests/test_llmjudgerouter.py b/tests/test_llmjudgerouter.py
index 1d53e544..5177070a 100644
--- a/tests/test_llmjudgerouter.py
+++ b/tests/test_llmjudgerouter.py
@@ -14,7 +14,7 @@ def _build_router():
     router.judge_api_key = None
     router.judge_model = "qwen3:0.6b"
     router.timeout_s = 5
-    router.max_tokens = 64
+    router.max_tokens = 256
     router.temperature = 0
     router.reason_max_chars = 80
     router.max_signals = 3
@@ -66,7 +66,7 @@ def post(self, url, headers=None, json=None):
         result = router.route_single({"query": "What is Goldbach's conjecture?"})
 
     assert captured["url"] == "http://127.0.0.1:11434/v1/chat/completions"
-    assert captured["json"]["max_tokens"] == 64
+    assert captured["json"]["max_tokens"] == 256
     assert captured["json"]["model"] == "qwen3:0.6b"
     assert result["model_name"] == "deepseek-flash"
     assert result["routing_confidence"] == 0.82
@@ -77,8 +77,7 @@ def post(self, url, headers=None, json=None):
 
 def test_llm_judge_router_uses_large_model_when_prompt_budget_is_risky():
     router = _build_router()
-    router.max_tokens = 64
-    router.prompt_budget_output_buffer = 4
+    router.max_tokens = 16
 
     result = router.route_single({"query": "x" * 100})