From 45b4e4a08dfd77f4cebebf89b9ac0f9b97457ec8 Mon Sep 17 00:00:00 2001 From: Foster Lipkey Date: Wed, 26 Aug 2026 10:03:17 -0400 Subject: [PATCH 1/6] feat: add LLM model router as a first-class runtime (issue #506) Adds a "router" runtime that picks a target runtime+model per message using a configurable "brain" LLM call (itself dispatched through the existing runtime infrastructure), with a configurable routing prompt, an explicit allowlist, per-runtime cooldown on infra failures, safe fallback to the agent's primary pair, and stickiness that resumes the previously-routed runtime's own sub-session so prompt caching still applies. Explicit /runtime or /model selection simply switches the session away from "router" (and back), so nothing is silently overridden. - llm_router.py: config load/save/validate, decision parsing, cooldown tracking, and LLMRouter.route() (pure, agent_manager-free, fully unit-testable). - config_schemas.py: RouterConfigSchema for the PUT /api/v1/router-config API validator. - agent_manager.py: "router" runtime registration, run_router() + supporting helpers, dispatch branch, and GET/PUT /api/v1/router-config, POST /api/v1/router/test, GET /api/v1/router/status. - tests/test_issue506_llm_router.py: 50 tests covering valid/invalid decisions, timeouts, allowlist enforcement, cooldown/fallback retry, stickiness, config validation, and the #84 "auto"-runtime guard. - webui/: reference RouterSettingsPanel.tsx + API client, following the existing AgentSettingsPanel.tsx convention (live WebUI wiring into webui/dist/app.js is a follow-up). Co-Authored-By: Claude Sonnet 5 --- agent_manager.py | 403 ++++++++++++++ config_schemas.py | 164 ++++++ llm_router.py | 377 +++++++++++++ tests/test_issue506_llm_router.py | 524 +++++++++++++++++++ webui/src/api/routerConfig.ts | 128 +++++ webui/src/components/RouterSettingsPanel.tsx | 349 ++++++++++++ webui/src/styles/RouterSettingsPanel.css | 232 ++++++++ 7 files changed, 2177 insertions(+) create mode 100644 llm_router.py create mode 100644 tests/test_issue506_llm_router.py create mode 100644 webui/src/api/routerConfig.ts create mode 100644 webui/src/components/RouterSettingsPanel.tsx create mode 100644 webui/src/styles/RouterSettingsPanel.css diff --git a/agent_manager.py b/agent_manager.py index a9773da7..df6c95da 100644 --- a/agent_manager.py +++ b/agent_manager.py @@ -25,6 +25,8 @@ from typing import Dict, List, Optional, Tuple from uuid import uuid4 +import llm_router + # Native Wee tool calls run in-process, so they cannot rely on subprocess # environment variables to retain the originating authenticated session. @@ -1304,6 +1306,15 @@ def check_runtime_available(runtime: str) -> bool: "wee": "copilot", # Copilot SDK BYOK (OpenAI fallback during migration), } + # Issue #506: the "router" runtime is available only when routing is + # enabled (config or WEE_ROUTER_ENABLED) AND its configured brain runtime + # is itself available β€” a router with no working brain is not available. + if runtime == "router": + if not is_router_enabled(): + return False + brain_runtime = get_router_config().load().get("brain", {}).get("runtime") + return bool(brain_runtime) and brain_runtime != "router" and check_runtime_available(brain_runtime) + executable_name = runtime_map.get(runtime) if not executable_name: return False @@ -1415,6 +1426,44 @@ def get_disabled_runtimes_manager() -> DisabledRuntimesManager: return _disabled_runtimes_manager +# --- Issue #506: LLM model router ------------------------------------------ +# Singletons mirror the DisabledRuntimesManager pattern above: module-level, +# lazily constructed, shared across SessionManager instances/threads. + +_router_config: "llm_router.RouterConfig | None" = None +_runtime_cooldown_tracker: "llm_router.RuntimeCooldownTracker | None" = None +_llm_router: "llm_router.LLMRouter | None" = None + + +def get_router_config() -> "llm_router.RouterConfig": + global _router_config + if _router_config is None: + _router_config = llm_router.RouterConfig() + return _router_config + + +def get_runtime_cooldown_tracker() -> "llm_router.RuntimeCooldownTracker": + global _runtime_cooldown_tracker + if _runtime_cooldown_tracker is None: + _runtime_cooldown_tracker = llm_router.RuntimeCooldownTracker() + return _runtime_cooldown_tracker + + +def get_llm_router() -> "llm_router.LLMRouter": + global _llm_router + if _llm_router is None: + _llm_router = llm_router.LLMRouter(get_router_config(), get_runtime_cooldown_tracker()) + return _llm_router + + +def is_router_enabled() -> bool: + """WEE_ROUTER_ENABLED env var overrides the config file's 'enabled' flag.""" + env_override = os.environ.get("WEE_ROUTER_ENABLED") + if env_override is not None: + return env_override.strip().lower() in ("1", "true", "yes", "on") + return bool(get_router_config().load().get("enabled", False)) + + def get_all_runtimes() -> List[Dict[str, str]]: """Return all known runtimes regardless of availability or disabled state.""" return [ @@ -1428,6 +1477,7 @@ def get_all_runtimes() -> List[Dict[str, str]]: {"id": "devin", "label": "devin"}, {"id": "cursor", "label": "cursor", "icon": "πŸ–±οΈ"}, {"id": "wee", "label": "wee", "icon": "πŸ€"}, + {"id": "router", "label": "router", "icon": "🧭"}, ] def get_available_runtimes() -> List[Dict[str, str]]: @@ -1447,6 +1497,7 @@ def get_available_runtimes() -> List[Dict[str, str]]: {"id": "devin", "label": "devin"}, {"id": "cursor", "label": "cursor", "icon": "πŸ–±οΈ"}, {"id": "wee", "label": "wee", "icon": "πŸ€"}, + {"id": "router", "label": "router", "icon": "🧭"}, ] available = [rt for rt in all_runtimes if check_runtime_available(rt["id"])] @@ -5050,6 +5101,7 @@ def get_models_for_runtime(self, runtime: str) -> Dict: "cursor": self.fetch_cursor_models, "wee": self.fetch_wee_models, "ollama": self.fetch_ollama_models, + "router": self._get_router_pseudo_models, } fetcher = dispatch.get(runtime) if fetcher is None: @@ -11440,6 +11492,229 @@ def agent_workspace_error(self, agent: str) -> Optional[str]: return None return self.describe_missing_agent_workspace(agent, path) + def _get_router_pseudo_models(self) -> Dict: + """Synthetic model list for the 'router' runtime (issue #506): one + entry per configured allowlist pair, so UIs that expect a 'model' + picker for every runtime have something sensible to show. Whatever + is picked here is cosmetic only β€” run_router() ignores it and + chooses the real target purely from routing policy.""" + cfg = get_router_config().load() + entries = [] + for item in cfg.get("allowlist", []): + label = f"{item.get('runtime')}/{item.get('model')}" + entries.append((label, item.get("hint") or label, [label])) + return {"routed": entries} if entries else {} + + def _agent_primary_pair(self, agent: str) -> Tuple[str, str]: + """Last-resort runtime+model pair when routing can't proceed at all + (router disabled, or even the configured fallback is unavailable).""" + agent_cfg = self.AGENTS.get(agent, self.AGENTS.get("orchestrator", {})) + rt = agent_cfg.get("primary_runtime") or get_default_runtime() + model = agent_cfg.get("primary_model") or get_default_model(rt) + return rt, model + + def _resume_state_for_runtime( + self, runtime: str, session_id: Optional[str], n8n_session_id: str + ) -> bool: + """Compute can_resume for an arbitrary target runtime. + + Mirrors the per-runtime branching in execute() (agent_manager.py, + the block starting 'Check if we can resume'), which normally only + needs to handle the session's single current_runtime. run_router() + needs the same logic for whichever runtime routing just picked, so + it's factored out here rather than duplicated inline. + """ + if runtime == "gemini": + return True + if runtime in ("devin", "cursor"): + return ( + self.session_exists(session_id, runtime, n8n_session_id=n8n_session_id) + if session_id + else self.session_exists("", runtime, n8n_session_id=n8n_session_id) + ) + if runtime == "wee": + return self.session_exists(session_id, runtime, n8n_session_id=n8n_session_id) + return self.session_exists(session_id, runtime) if session_id else False + + def _dispatch_router_target( + self, + target_runtime: str, + target_model: str, + prompt: str, + agent: str, + n8n_session_id: str, + effective_timeout: int, + render_type: str, + mode: str, + session_data: dict, + ) -> str: + """Dispatch a router-selected request to its target runtime, resuming + that runtime's own sub-session when one exists so repeated routing to + the same target reuses its prompt cache (issue #506 stickiness). + + Per-target sub-session ids are tracked in session_data["router_sessions"] + (keyed by runtime id) rather than the session's top-level "session_id" + field, because several run_* implementations (copilot, claude, + claude-sdk, codex) write a freshly-established session id directly into + session_data["session_id"] as a side effect. That field is snapshotted + before dispatch and restored after, so a router-selected sub-target's + session id never leaks into the top-level field an explicit + `/runtime ` switch would otherwise inherit. + """ + router_sessions = dict(session_data.get("router_sessions") or {}) + sub_session_id = router_sessions.get(target_runtime) + original_session_id = session_data.get("session_id") + + can_resume = self._resume_state_for_runtime( + target_runtime, sub_session_id, n8n_session_id + ) + + output = self._dispatch_single_runtime( + target_runtime, + prompt, + target_model, + agent, + sub_session_id, + can_resume, + n8n_session_id, + effective_timeout, + render_type, + mode, + ) + + fresh_session_data = self.get_or_create_session_data(n8n_session_id) + new_sub_id = fresh_session_data.get("session_id") + if not can_resume and target_runtime in ("copilot", "opencode", "gemini"): + discovered = self.get_most_recent_session_id(target_runtime, agent) + if discovered: + new_sub_id = discovered + if new_sub_id and new_sub_id != original_session_id: + router_sessions[target_runtime] = new_sub_id + self.update_session_field(n8n_session_id, "router_sessions", router_sessions) + + # Restore the shared field regardless of the branch above β€” some + # run_* implementations write to it unconditionally. + self.update_session_field(n8n_session_id, "session_id", original_session_id) + + return output + + def run_router( + self, + prompt: str, + model: str, + agent: str, + session_id: Optional[str], + can_resume: bool, + n8n_session_id: str, + effective_timeout: int, + render_type: str, + mode: str = "restricted", + ) -> str: + """Dispatch via the LLM router (issue #506). + + Picks a runtime+model pair for this request using a small 'brain' LLM + call (itself invoked via a full runtime dispatch, per session config), + then dispatches the real prompt to that pair. Never raises: every + failure path (router disabled, brain timeout/error, invalid/disallowed + decision, even a failed fallback) degrades to the agent's configured + primary_runtime/primary_model rather than breaking the request. + """ + session_data = self.get_or_create_session_data(n8n_session_id) + cfg = get_router_config().load() + + if not is_router_enabled(): + agent_rt, agent_model = self._agent_primary_pair(agent) + print( + f"[Router] session={n8n_session_id} disabled -> agent primary " + f"{agent_rt}/{agent_model}", + file=sys.stderr, + ) + return self._dispatch_router_target( + agent_rt, agent_model, prompt, agent, n8n_session_id, + effective_timeout, render_type, mode, session_data, + ) + + router = get_llm_router() + disabled_mgr = get_disabled_runtimes_manager() + + def runtime_available(rt: str) -> bool: + return check_runtime_available(rt) and not disabled_mgr.is_disabled(rt) + + def invoke_brain(brain_runtime: str, brain_model: str, brain_prompt: str, timeout: float) -> Optional[str]: + # One-shot, no history: a fixed (not random) sub-session id per + # origin session keeps the brain's own runtime storage tidy + # without accumulating routing chatter into the user's actual + # conversation history. + brain_session_id = f"{n8n_session_id}__router_brain" + return self._dispatch_single_runtime( + brain_runtime, brain_prompt, brain_model, agent, + None, False, brain_session_id, int(timeout), "plain", mode, + ) + + def resolve_model(name: str, rt: str) -> Optional[str]: + return self.get_model_from_name(name, rt) + + last_routed = session_data.get("router_last") + decision = router.route( + prompt=prompt, + last_routed=last_routed, + runtime_available=runtime_available, + invoke_brain=invoke_brain, + resolve_model=resolve_model, + ) + + if not decision.runtime: + agent_rt, agent_model = self._agent_primary_pair(agent) + print( + f"[Router] session={n8n_session_id} decision=none reason={decision.reason!r} " + f"-> agent primary {agent_rt}/{agent_model}", + file=sys.stderr, + ) + return self._dispatch_router_target( + agent_rt, agent_model, prompt, agent, n8n_session_id, + effective_timeout, render_type, mode, session_data, + ) + + debug = os.environ.get("WEE_ROUTER_DEBUG") == "1" + print( + f"[Router] session={n8n_session_id} decision={decision.runtime}/{decision.model} " + f"source={decision.source} latency_ms={decision.latency_ms}" + + (f" reason={decision.reason!r}" if debug else f" reason_len={len(decision.reason)}"), + file=sys.stderr, + ) + + output = self._dispatch_router_target( + decision.runtime, decision.model, prompt, agent, n8n_session_id, + effective_timeout, render_type, mode, session_data, + ) + + if llm_router.is_infra_failure_text(output): + get_runtime_cooldown_tracker().mark_failure( + decision.runtime, "dispatch failure", cfg.get("cooldown_seconds", 300) + ) + fallback = cfg.get("fallback", {}) + if fallback.get("runtime") and fallback.get("runtime") != decision.runtime: + print( + f"[Router] session={n8n_session_id} target {decision.runtime} failed " + f"(infra) -> retrying once with fallback " + f"{fallback['runtime']}/{fallback.get('model')}", + file=sys.stderr, + ) + output = self._dispatch_router_target( + fallback["runtime"], fallback.get("model", ""), prompt, agent, + n8n_session_id, effective_timeout, render_type, mode, session_data, + ) + decision = llm_router.RouteDecision( + runtime=fallback["runtime"], model=fallback.get("model", ""), + reason="post-dispatch infra fallback", source="fallback", + ) + + self.update_session_field( + n8n_session_id, "router_last", + {"runtime": decision.runtime, "model": decision.model, "ts": time.time()}, + ) + return output + def _dispatch_single_runtime( self, runtime: str, @@ -11589,6 +11864,18 @@ def _dispatch_single_runtime( ) finally: _wee_dispatch_context.reset(context_token) + elif runtime == "router": + result = self.run_router( + prompt, + model, + agent, + session_id, + can_resume, + n8n_session_id, + effective_timeout, + render_type, + mode, + ) else: return f"Error: Unknown runtime '{runtime}'" @@ -18575,6 +18862,122 @@ async def put_env_file(request: Request): except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to save .env: {e}") + @app.get("/api/v1/router-config") + async def get_router_config_endpoint(request: Request): + """Return the LLM router config (issue #506), plus validation warnings + and whether routing is effectively enabled once the WEE_ROUTER_ENABLED + env override is applied.""" + auth = await authenticate( + request, + authorization=request.headers.get("authorization"), + x_user_identity=request.headers.get("x-user-identity"), + x_auth_channel=request.headers.get("x-auth-channel"), + ) + cfg = get_router_config().load() + return { + "config": cfg, + "enabled_effective": is_router_enabled(), + "validation_errors": llm_router.RouterConfig.validate(cfg), + } + + @app.put("/api/v1/router-config") + async def put_router_config_endpoint(request: Request): + """Validate and persist the LLM router config.""" + auth = await authenticate( + request, + authorization=request.headers.get("authorization"), + x_user_identity=request.headers.get("x-user-identity"), + x_auth_channel=request.headers.get("x-auth-channel"), + ) + body = await request.json() + config = body.get("config", body) + try: + from config_schemas import validate_router_config + + validate_router_config(config) + except Exception as e: + raise HTTPException(status_code=422, detail=f"Invalid router config: {e}") + + try: + get_router_config().save(config) + return {"saved": True, "config": config} + except llm_router.RouterConfigError as e: + raise HTTPException(status_code=422, detail=str(e)) + except Exception as e: + raise HTTPException(status_code=500, detail=f"Failed to save router config: {e}") + + @app.post("/api/v1/router/test") + async def test_router_endpoint(request: Request): + """Dry-run a routing decision for a prompt, without mutating any + session. Lets the router prompt/allowlist be iterated on from the + WebUI before enabling routing for real traffic.""" + auth = await authenticate( + request, + authorization=request.headers.get("authorization"), + x_user_identity=request.headers.get("x-user-identity"), + x_auth_channel=request.headers.get("x-auth-channel"), + ) + body = await request.json() + prompt = body.get("prompt", "") + if not prompt or not isinstance(prompt, str): + raise HTTPException(status_code=400, detail="prompt is required") + + disabled_mgr = get_disabled_runtimes_manager() + + def runtime_available(rt: str) -> bool: + return check_runtime_available(rt) and not disabled_mgr.is_disabled(rt) + + def invoke_brain(brain_runtime, brain_model, brain_prompt, timeout): + return session_mgr._dispatch_single_runtime( + brain_runtime, brain_prompt, brain_model, "orchestrator", + None, False, f"router_test_{uuid4()}", int(timeout), "plain", + ) + + def resolve_model(name, rt): + return session_mgr.get_model_from_name(name, rt) + + router = get_llm_router() + cfg = router.config.load() + eligible = router.eligible_pairs(cfg.get("allowlist", []), runtime_available) + start = time.time() + decision = router.route( + prompt=prompt, + last_routed=None, + runtime_available=runtime_available, + invoke_brain=invoke_brain, + resolve_model=resolve_model, + ) + return { + "decision": { + "runtime": decision.runtime, + "model": decision.model, + "reason": decision.reason, + "source": decision.source, + "latency_ms": decision.latency_ms, + }, + "eligible_pairs": eligible, + "total_ms": int((time.time() - start) * 1000), + } + + @app.get("/api/v1/router/status") + async def get_router_status_endpoint(request: Request): + """Router health snapshot: enabled state, brain runtime reachability, + and current per-runtime cooldowns (issue #506).""" + auth = await authenticate( + request, + authorization=request.headers.get("authorization"), + x_user_identity=request.headers.get("x-user-identity"), + x_auth_channel=request.headers.get("x-auth-channel"), + ) + cfg = get_router_config().load() + brain_runtime = cfg.get("brain", {}).get("runtime") + return { + "enabled": is_router_enabled(), + "brain": cfg.get("brain"), + "brain_available": bool(brain_runtime) and check_runtime_available(brain_runtime), + "cooldowns": get_runtime_cooldown_tracker().status(), + } + @app.get("/api/v1/settings/model-manifest") async def get_model_manifest_settings(request: Request, runtime: str = "claude"): """Return the editable model list for a runtime. diff --git a/config_schemas.py b/config_schemas.py index 334e7897..d04c63ee 100644 --- a/config_schemas.py +++ b/config_schemas.py @@ -302,6 +302,170 @@ def validate_webex_config(data: dict) -> WebEXConfigSchema: return WebEXConfigSchema.model_validate(data) +# --------------------------------------------------------------------------- +# router_config.json schema (issue #506 β€” LLM model router) +# --------------------------------------------------------------------------- + +_ROUTER_BRAIN_KNOWN = {"runtime", "model"} + + +class RouterBrainConfig(BaseModel): + """Runtime+model pair used to invoke the router's decision-making LLM.""" + + model_config = ConfigDict(extra="allow") + + runtime: str + model: str + + @model_validator(mode="before") + @classmethod + def warn_unknown(cls, values: Any) -> Any: + _warn_unknown_keys(values, _ROUTER_BRAIN_KNOWN, "router_config.brain") + return values + + @model_validator(mode="after") + def no_recursion(self) -> "RouterBrainConfig": + if self.runtime == "router": + raise ValueError("brain.runtime cannot be 'router' (no recursive routing)") + return self + + +_ROUTER_ALLOWLIST_ENTRY_KNOWN = {"runtime", "model", "hint"} + + +class RouterAllowlistEntry(BaseModel): + """A single routable runtime+model target.""" + + model_config = ConfigDict(extra="allow") + + runtime: str + model: str + hint: Optional[str] = "" + + @model_validator(mode="before") + @classmethod + def warn_unknown(cls, values: Any) -> Any: + _warn_unknown_keys(values, _ROUTER_ALLOWLIST_ENTRY_KNOWN, "router_config.allowlist[]") + return values + + @model_validator(mode="after") + def no_recursion(self) -> "RouterAllowlistEntry": + if self.runtime == "router": + raise ValueError("allowlist entries cannot target 'router' itself") + return self + + +_ROUTER_FALLBACK_KNOWN = {"runtime", "model"} + + +class RouterFallbackConfig(BaseModel): + """Safe fallback pair used when the router brain fails, times out, or + returns an invalid/disallowed decision.""" + + model_config = ConfigDict(extra="allow") + + runtime: str + model: str + + @model_validator(mode="before") + @classmethod + def warn_unknown(cls, values: Any) -> Any: + _warn_unknown_keys(values, _ROUTER_FALLBACK_KNOWN, "router_config.fallback") + return values + + @model_validator(mode="after") + def no_recursion(self) -> "RouterFallbackConfig": + if self.runtime == "router": + raise ValueError("fallback.runtime cannot be 'router'") + return self + + +_ROUTER_STICKINESS_KNOWN = {"enabled", "prefer_same_runtime", "window_seconds"} + + +class RouterStickinessConfig(BaseModel): + """Preference for reusing the previously-routed runtime/model so the + underlying runtime's own session (and its prompt cache) is reused.""" + + model_config = ConfigDict(extra="allow") + + enabled: bool = True + prefer_same_runtime: bool = True + window_seconds: int = 900 + + @model_validator(mode="before") + @classmethod + def warn_unknown(cls, values: Any) -> Any: + _warn_unknown_keys(values, _ROUTER_STICKINESS_KNOWN, "router_config.stickiness") + return values + + +_ROUTER_CONFIG_KNOWN = { + "enabled", + "brain", + "timeout_seconds", + "prompt_template", + "allowlist", + "fallback", + "stickiness", + "cooldown_seconds", +} + +_REQUIRED_ROUTER_TEMPLATE_PLACEHOLDERS = ("{allowlist_table}", "{user_message}") + + +class RouterConfigSchema(BaseModel): + """Schema for config/router_config.json (issue #506). + + This is the API-facing validator used by PUT /api/v1/router-config to + return structured 422 errors. llm_router.RouterConfig.validate() runs the + same checks independently as a defense-in-depth guard for callers that + write the file directly (e.g. tests, manual edits) without going through + the API. + """ + + model_config = ConfigDict(extra="allow") + + enabled: bool = False + brain: RouterBrainConfig + timeout_seconds: float = 30 + prompt_template: str + allowlist: List[RouterAllowlistEntry] + fallback: RouterFallbackConfig + stickiness: RouterStickinessConfig = Field(default_factory=RouterStickinessConfig) + cooldown_seconds: float = 300 + + @model_validator(mode="before") + @classmethod + def warn_unknown(cls, values: Any) -> Any: + _warn_unknown_keys(values, _ROUTER_CONFIG_KNOWN, "router_config.json") + return values + + @model_validator(mode="after") + def check_semantics(self) -> "RouterConfigSchema": + if not self.allowlist: + raise ValueError("allowlist must contain at least one runtime/model pair") + for placeholder in _REQUIRED_ROUTER_TEMPLATE_PLACEHOLDERS: + if placeholder not in self.prompt_template: + raise ValueError(f"prompt_template missing required placeholder {placeholder}") + if self.timeout_seconds <= 0: + raise ValueError("timeout_seconds must be a positive number") + if self.cooldown_seconds < 0: + raise ValueError("cooldown_seconds must be a non-negative number") + return self + + +def validate_router_config(data: dict) -> RouterConfigSchema: + """Validate router_config.json data. + + Raises pydantic.ValidationError on structural or semantic failures + (missing brain/allowlist/fallback, recursive 'router' targets, a + prompt_template missing required placeholders, non-positive timeouts). + Unknown keys produce UserWarnings. + """ + return RouterConfigSchema.model_validate(data) + + # --------------------------------------------------------------------------- # Registry: maps config file basename -> validator function # Used by base_connector.BaseConfig._load_config to auto-validate on load. diff --git a/llm_router.py b/llm_router.py new file mode 100644 index 00000000..f1789f89 --- /dev/null +++ b/llm_router.py @@ -0,0 +1,377 @@ +"""llm_router.py β€” Configuration, prompt construction, and decision validation +for the Wee "router" runtime (issue #506). + +Design note: this module is intentionally free of any agent_manager / +dispatch imports. The actual call to the router "brain" (an LLM invoked via +a full runtime dispatch β€” e.g. `wee` + a local Ollama model, or `copilot` + +haiku) is made by `SessionManager.run_router()` in agent_manager.py, which +injects an `invoke_brain` callable into `LLMRouter.route()`. Keeping the +invocation out of this module makes prompt-building and decision validation +trivially unit-testable without spinning up any runtime. +""" + +from __future__ import annotations + +import json +import os +import re +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Optional + +CONFIG_DIR = Path(os.path.dirname(os.path.abspath(__file__))) / "config" +ROUTER_CONFIG_PATH = CONFIG_DIR / "router_config.json" + +REQUIRED_TEMPLATE_PLACEHOLDERS = ("{allowlist_table}", "{user_message}") + +DEFAULT_ROUTER_CONFIG = { + "enabled": False, + "brain": {"runtime": "wee", "model": "ollama/qwen3:8b"}, + "timeout_seconds": 30, + "prompt_template": ( + "You are a model router for an AI assistant platform. Given the " + "candidate runtime/model pairs below and the user's message, pick " + "the single best pair for this request.\n\n" + "Candidates:\n{allowlist_table}\n\n" + "{stickiness_hint}\n" + "User message:\n{user_message}\n\n" + "Respond with ONLY a JSON object, no other text:\n" + '{{"runtime": "...", "model": "...", "reason": "..."}}' + ), + "allowlist": [], + "fallback": {"runtime": "copilot", "model": "auto"}, + "stickiness": {"enabled": True, "prefer_same_runtime": True, "window_seconds": 900}, + "cooldown_seconds": 300, +} + +# Infra-failure detection, mirroring agent_manager.py's background-task +# fallback patterns (issue #219/#243) so a runtime that is rate-limited, +# unauthorized, unreachable, or timing out gets a router-side cooldown using +# the same signal already trusted for background-task fallback. Kept as its +# own copy here β€” not imported β€” because the original lives as a local +# closure inside agent_manager's background-task endpoint, not a module-level +# helper. +_INFRA_FAILURE_PATTERNS = [ + re.compile(p, re.IGNORECASE) + for p in [ + r"\b429\b", + r"\brate[\s\-]?limit(?:ed|ing)?\b", + r"\bquota[\s\-]exceeded\b", + r"\b401\b", + r"\bunauthorized\b", + r"\bmissing[\s\-]authentication\b", + r"\bapi[\s_\-]?key[\s_\-]?(?:invalid|expired|missing)\b", + r"\b503\b", + r"\bservice[\s\-]unavailable\b", + r"\b502\b", + r"\bbad[\s\-]gateway\b", + r"\bconnection[\s\-]refused\b", + r"\btimed?\s*out\b", + r"\betimedout\b", + r"\boverloaded\b", + ] +] +_INFRA_EXCLUSION_RE = re.compile( + r"^(?:assert(?:ion)?error|typeerror|valueerror|keyerror" + r"|attributeerror|nameerror|runtimeerror)\s*:", + re.IGNORECASE, +) + + +def is_infra_failure_text(error_text: Optional[str]) -> bool: + """True if error_text looks like an infra failure (rate limit, auth, 5xx, timeout).""" + if not error_text: + return False + if _INFRA_EXCLUSION_RE.match(error_text.strip()): + return False + return any(pat.search(error_text) for pat in _INFRA_FAILURE_PATTERNS) + + +@dataclass +class RouteDecision: + runtime: str + model: str + reason: str = "" + latency_ms: int = 0 + source: str = "router" # "router" | "single" | "fallback" + + +class RouterConfigError(ValueError): + """Raised by RouterConfig.save() when the config fails validation.""" + + +class RouterConfig: + """Loads/saves/validates config/router_config.json, mtime-cached like + the model-manifest loader so edits apply live without a restart.""" + + def __init__(self, path: Path = ROUTER_CONFIG_PATH): + self._path = Path(path) + self._cache: Optional[dict] = None + self._cache_mtime: Optional[float] = None + + def load(self) -> dict: + try: + mtime = self._path.stat().st_mtime + except FileNotFoundError: + return json.loads(json.dumps(DEFAULT_ROUTER_CONFIG)) + if self._cache is not None and self._cache_mtime == mtime: + return self._cache + try: + with open(self._path) as f: + data = json.load(f) + except (json.JSONDecodeError, OSError): + return json.loads(json.dumps(DEFAULT_ROUTER_CONFIG)) + merged = {**DEFAULT_ROUTER_CONFIG, **data} + self._cache = merged + self._cache_mtime = mtime + return merged + + def save(self, config: dict) -> None: + errors = self.validate(config) + if errors: + raise RouterConfigError("; ".join(errors)) + self._path.parent.mkdir(parents=True, exist_ok=True) + if self._path.exists(): + try: + self._path.replace(self._path.with_suffix(".json.bak")) + except OSError: + pass + with open(self._path, "w") as f: + json.dump(config, f, indent=2) + self._cache = None + self._cache_mtime = None + + @staticmethod + def validate(config: dict) -> list: + errors = [] + brain = config.get("brain") or {} + if not brain.get("runtime"): + errors.append("brain.runtime is required") + elif brain.get("runtime") == "router": + errors.append("brain.runtime cannot be 'router' (no recursive routing)") + if not brain.get("model"): + errors.append("brain.model is required") + + allowlist = config.get("allowlist") or [] + if not allowlist: + errors.append("allowlist must contain at least one runtime/model pair") + for i, entry in enumerate(allowlist): + if not entry.get("runtime") or not entry.get("model"): + errors.append(f"allowlist[{i}] must have both runtime and model") + elif entry.get("runtime") == "router": + errors.append(f"allowlist[{i}] cannot target 'router' itself") + + fallback = config.get("fallback") or {} + if not fallback.get("runtime") or not fallback.get("model"): + errors.append("fallback.runtime and fallback.model are required") + elif fallback.get("runtime") == "router": + errors.append("fallback.runtime cannot be 'router'") + + template = config.get("prompt_template") or "" + for placeholder in REQUIRED_TEMPLATE_PLACEHOLDERS: + if placeholder not in template: + errors.append(f"prompt_template missing required placeholder {placeholder}") + + timeout = config.get("timeout_seconds") + if not isinstance(timeout, (int, float)) or timeout <= 0: + errors.append("timeout_seconds must be a positive number") + + cooldown = config.get("cooldown_seconds") + if not isinstance(cooldown, (int, float)) or cooldown < 0: + errors.append("cooldown_seconds must be a non-negative number") + + return errors + + +class RuntimeCooldownTracker: + """Tracks runtimes temporarily excluded from routing after an infra failure.""" + + def __init__(self): + self._cooldowns = {} # runtime -> {"until": ts, "reason": str} + + def mark_failure(self, runtime: str, reason: str, cooldown_seconds: float) -> None: + self._cooldowns[runtime] = {"until": time.time() + cooldown_seconds, "reason": reason} + + def is_cooling(self, runtime: str) -> bool: + entry = self._cooldowns.get(runtime) + if not entry: + return False + if time.time() >= entry["until"]: + del self._cooldowns[runtime] + return False + return True + + def status(self) -> dict: + now = time.time() + return { + rt: {"reason": e["reason"], "seconds_remaining": max(0, round(e["until"] - now))} + for rt, e in self._cooldowns.items() + if e["until"] > now + } + + +def build_allowlist_table(allowlist: list) -> str: + lines = [] + for entry in allowlist: + hint = entry.get("hint", "") + line = f"- runtime={entry['runtime']} model={entry['model']}" + if hint: + line += f" β€” {hint}" + lines.append(line) + return "\n".join(lines) if lines else "(none available)" + + +def build_stickiness_hint(last_routed: Optional[dict], stickiness_cfg: dict) -> str: + if not stickiness_cfg.get("enabled") or not last_routed: + return "" + age = time.time() - last_routed.get("ts", 0) + window = stickiness_cfg.get("window_seconds", 900) + if age > window: + return "" + pair = f"runtime={last_routed.get('runtime')} model={last_routed.get('model')}" + if stickiness_cfg.get("prefer_same_runtime", True): + return ( + f"The previous message in this conversation was routed to {pair}. " + "Prefer staying on the same runtime/model unless the new request " + "clearly needs a different one β€” this reuses cached context and " + "is cheaper and faster." + ) + return f"The previous message in this conversation was routed to {pair}." + + +_JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) + + +def parse_decision_json(raw: Optional[str]) -> Optional[dict]: + """Tolerantly extract a JSON object from an LLM's free-form reply.""" + if not raw: + return None + text = raw.strip() + if text.startswith("```"): + text = re.sub(r"^```(?:json)?\s*", "", text) + text = re.sub(r"\s*```$", "", text) + match = _JSON_OBJECT_RE.search(text) + if not match: + return None + try: + data = json.loads(match.group(0)) + except json.JSONDecodeError: + return None + return data if isinstance(data, dict) else None + + +class LLMRouter: + """Stateless decision engine: builds the router prompt, validates the + brain's reply against the eligible allowlist, and falls back safely on + any failure. `route()` never raises β€” every failure path returns a + RouteDecision with source="fallback" (empty runtime/model only if even + the fallback pair is unavailable, in which case the caller must supply + an agent-primary last resort).""" + + def __init__(self, config: RouterConfig, cooldowns: RuntimeCooldownTracker): + self.config = config + self.cooldowns = cooldowns + + def eligible_pairs(self, allowlist: list, runtime_available: Callable[[str], bool]) -> list: + return [ + entry + for entry in allowlist + if entry.get("runtime") != "router" + and not self.cooldowns.is_cooling(entry["runtime"]) + and runtime_available(entry["runtime"]) + ] + + def route( + self, + *, + prompt: str, + last_routed: Optional[dict], + runtime_available: Callable[[str], bool], + invoke_brain: Callable[[str, str, str, float], Optional[str]], + resolve_model: Callable[[str, str], Optional[str]], + ) -> RouteDecision: + cfg = self.config.load() + allowlist = cfg.get("allowlist") or [] + fallback = cfg.get("fallback") or {} + eligible = self.eligible_pairs(allowlist, runtime_available) + + if not eligible: + return self._fallback_decision(fallback, runtime_available, "no eligible runtimes") + + if len(eligible) == 1: + only = eligible[0] + return RouteDecision( + runtime=only["runtime"], model=only["model"], + reason="only eligible pair", source="single", + ) + + stickiness_hint = build_stickiness_hint(last_routed, cfg.get("stickiness", {})) + allowlist_table = build_allowlist_table(eligible) + truncated_prompt = prompt[:1500] + try: + full_prompt = cfg["prompt_template"].format( + allowlist_table=allowlist_table, + stickiness_hint=stickiness_hint, + user_message=truncated_prompt, + ) + except (KeyError, IndexError): + return self._fallback_decision(fallback, runtime_available, "invalid prompt_template") + + start = time.time() + try: + raw = invoke_brain( + cfg["brain"]["runtime"], cfg["brain"]["model"], full_prompt, + cfg.get("timeout_seconds", 30), + ) + except Exception as exc: # the brain call must never break the request + return self._fallback_decision( + fallback, runtime_available, f"brain error: {type(exc).__name__}", + ) + latency_ms = int((time.time() - start) * 1000) + + data = parse_decision_json(raw) + if not data: + return self._fallback_decision( + fallback, runtime_available, "unparseable brain reply", latency_ms, + ) + + chosen_runtime = str(data.get("runtime", "")).strip() + chosen_model_raw = str(data.get("model", "")).strip() + reason = str(data.get("reason", ""))[:300] + + eligible_by_runtime = {e["runtime"]: e for e in eligible} + if chosen_runtime not in eligible_by_runtime: + return self._fallback_decision( + fallback, runtime_available, + f"runtime '{chosen_runtime}' not eligible", latency_ms, + ) + + allowed_entry = eligible_by_runtime[chosen_runtime] + resolved_model = resolve_model(chosen_model_raw, chosen_runtime) or chosen_model_raw + if resolved_model != allowed_entry["model"] and chosen_model_raw != allowed_entry["model"]: + # Runtime is eligible but the brain named a model outside that + # pair's allowlisted model β€” trust the allowlist, not free text. + return RouteDecision( + runtime=chosen_runtime, model=allowed_entry["model"], + reason=reason or "model normalized to allowlisted pair", + latency_ms=latency_ms, source="router", + ) + + return RouteDecision( + runtime=chosen_runtime, model=allowed_entry["model"], + reason=reason, latency_ms=latency_ms, source="router", + ) + + def _fallback_decision( + self, fallback: dict, runtime_available: Callable[[str], bool], + reason: str, latency_ms: int = 0, + ) -> RouteDecision: + rt = fallback.get("runtime") + model = fallback.get("model") + if rt and model and not self.cooldowns.is_cooling(rt) and runtime_available(rt): + return RouteDecision(runtime=rt, model=model, reason=reason, latency_ms=latency_ms, source="fallback") + return RouteDecision( + runtime="", model="", reason=f"{reason}; fallback also unavailable", + latency_ms=latency_ms, source="fallback", + ) diff --git a/tests/test_issue506_llm_router.py b/tests/test_issue506_llm_router.py new file mode 100644 index 00000000..a2ed08ab --- /dev/null +++ b/tests/test_issue506_llm_router.py @@ -0,0 +1,524 @@ +""" +Regression / feature tests for Issue #506 β€” LLM model router. + +Part A tests llm_router.py in isolation (no agent_manager import at all): +config load/save/validate, cooldown tracking, prompt building, tolerant JSON +parsing, and LLMRouter.route()'s full decision/validation/fallback logic. + +Part B tests the agent_manager.py integration points: runtime registration +("router" appears in / is excluded from available runtimes based on config), +config_schemas validation, and SessionManager.run_router() built as a minimal +double (per the tests/test_issue125_429_retry.py `_make_mgr` pattern) with +_dispatch_single_runtime mocked so no real runtime is ever invoked. +""" + +import json +import os +import sys +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +os.environ.setdefault("API_SHARED_KEY", "test_key_123") + +REPO = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO)) + +import llm_router # noqa: E402 + + +# --------------------------------------------------------------------------- +# Part A β€” llm_router.py in isolation +# --------------------------------------------------------------------------- + + +def _cfg(**overrides): + cfg = json.loads(json.dumps(llm_router.DEFAULT_ROUTER_CONFIG)) + cfg["allowlist"] = [ + {"runtime": "claude-sdk", "model": "opus", "hint": "complex coding"}, + {"runtime": "copilot", "model": "haiku", "hint": "quick questions"}, + ] + cfg["fallback"] = {"runtime": "copilot", "model": "auto"} + cfg.update(overrides) + return cfg + + +class _FakeConfig: + """Stand-in for RouterConfig that returns a fixed dict, no filesystem.""" + + def __init__(self, cfg): + self._cfg = cfg + + def load(self): + return self._cfg + + +def _identity_resolver(name, rt): + return name + + +class TestRouterConfig: + def test_load_missing_file_returns_disabled_defaults(self, tmp_path): + cfg = llm_router.RouterConfig(tmp_path / "nope.json") + data = cfg.load() + assert data["enabled"] is False + assert data["allowlist"] == [] + + def test_save_then_load_round_trips(self, tmp_path): + path = tmp_path / "router_config.json" + cfg = llm_router.RouterConfig(path) + valid = _cfg(enabled=True) + cfg.save(valid) + reloaded = llm_router.RouterConfig(path).load() + assert reloaded["enabled"] is True + assert reloaded["allowlist"][0]["runtime"] == "claude-sdk" + + def test_save_rejects_invalid_config(self, tmp_path): + cfg = llm_router.RouterConfig(tmp_path / "router_config.json") + with pytest.raises(llm_router.RouterConfigError): + cfg.save({"enabled": True, "allowlist": []}) + + def test_validate_rejects_recursive_brain(self): + bad = _cfg(brain={"runtime": "router", "model": "x"}) + errors = llm_router.RouterConfig.validate(bad) + assert any("brain.runtime" in e for e in errors) + + def test_validate_rejects_recursive_allowlist_entry(self): + bad = _cfg() + bad["allowlist"].append({"runtime": "router", "model": "x"}) + errors = llm_router.RouterConfig.validate(bad) + assert any("cannot target 'router'" in e for e in errors) + + def test_validate_rejects_missing_template_placeholders(self): + bad = _cfg(prompt_template="no placeholders here") + errors = llm_router.RouterConfig.validate(bad) + assert any("prompt_template" in e for e in errors) + + def test_validate_rejects_bad_timeout(self): + bad = _cfg(timeout_seconds=0) + errors = llm_router.RouterConfig.validate(bad) + assert any("timeout_seconds" in e for e in errors) + + def test_validate_rejects_empty_allowlist(self): + bad = _cfg(allowlist=[]) + errors = llm_router.RouterConfig.validate(bad) + assert any("allowlist" in e for e in errors) + + +class TestRuntimeCooldownTracker: + def test_mark_and_check(self): + tracker = llm_router.RuntimeCooldownTracker() + assert not tracker.is_cooling("copilot") + tracker.mark_failure("copilot", "429", cooldown_seconds=60) + assert tracker.is_cooling("copilot") + assert "copilot" in tracker.status() + + def test_cooldown_expires(self, monkeypatch): + tracker = llm_router.RuntimeCooldownTracker() + base = 1000.0 + monkeypatch.setattr(llm_router.time, "time", lambda: base) + tracker.mark_failure("copilot", "429", cooldown_seconds=10) + assert tracker.is_cooling("copilot") + monkeypatch.setattr(llm_router.time, "time", lambda: base + 11) + assert not tracker.is_cooling("copilot") + assert tracker.status() == {} + + +class TestParsing: + def test_parse_plain_json(self): + assert llm_router.parse_decision_json('{"runtime":"copilot","model":"auto","reason":"x"}') == { + "runtime": "copilot", "model": "auto", "reason": "x", + } + + def test_parse_fenced_json(self): + raw = '```json\n{"runtime": "copilot", "model": "auto", "reason": "y"}\n```' + assert llm_router.parse_decision_json(raw)["runtime"] == "copilot" + + def test_parse_json_with_surrounding_prose(self): + raw = 'Sure, here is my pick:\n{"runtime": "claude-sdk", "model": "opus", "reason": "z"}\nHope that helps!' + assert llm_router.parse_decision_json(raw)["model"] == "opus" + + def test_parse_garbage_returns_none(self): + assert llm_router.parse_decision_json("not json at all") is None + + def test_parse_empty_returns_none(self): + assert llm_router.parse_decision_json("") is None + assert llm_router.parse_decision_json(None) is None + + +class TestIsInfraFailureText: + @pytest.mark.parametrize("text", [ + "Error 429: rate limited", "Error: quota exceeded", "401 unauthorized", + "503 Service Unavailable", "connection refused", "request timed out", + "model overloaded, please retry", + ]) + def test_detects_infra_failures(self, text): + assert llm_router.is_infra_failure_text(text) + + def test_ignores_application_exceptions(self): + assert not llm_router.is_infra_failure_text("AssertionError: expected 429 in output") + + def test_ignores_unrelated_text(self): + assert not llm_router.is_infra_failure_text("Here are the search results you asked for.") + + def test_handles_empty(self): + assert not llm_router.is_infra_failure_text("") + assert not llm_router.is_infra_failure_text(None) + + +class TestBuildHelpers: + def test_allowlist_table_includes_hints(self): + table = llm_router.build_allowlist_table(_cfg()["allowlist"]) + assert "claude-sdk" in table and "complex coding" in table + + def test_stickiness_hint_within_window(self): + last = {"runtime": "copilot", "model": "auto", "ts": time.time()} + hint = llm_router.build_stickiness_hint(last, {"enabled": True, "prefer_same_runtime": True, "window_seconds": 900}) + assert "copilot" in hint + + def test_stickiness_hint_expired(self): + last = {"runtime": "copilot", "model": "auto", "ts": time.time() - 10000} + hint = llm_router.build_stickiness_hint(last, {"enabled": True, "window_seconds": 900}) + assert hint == "" + + def test_stickiness_hint_disabled(self): + last = {"runtime": "copilot", "model": "auto", "ts": time.time()} + hint = llm_router.build_stickiness_hint(last, {"enabled": False}) + assert hint == "" + + +class TestLLMRouterRoute: + def _router(self, cfg): + return llm_router.LLMRouter(_FakeConfig(cfg), llm_router.RuntimeCooldownTracker()) + + def test_valid_decision_in_allowlist(self): + router = self._router(_cfg()) + decision = router.route( + prompt="write me a recursive quicksort", + last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=lambda rt, m, p, t: '{"runtime":"claude-sdk","model":"opus","reason":"coding task"}', + resolve_model=_identity_resolver, + ) + assert decision.runtime == "claude-sdk" + assert decision.model == "opus" + assert decision.source == "router" + + def test_disallowed_runtime_falls_back(self): + router = self._router(_cfg()) + decision = router.route( + prompt="hi", + last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=lambda rt, m, p, t: '{"runtime":"devin","model":"x","reason":"nope"}', + resolve_model=_identity_resolver, + ) + assert decision.source == "fallback" + assert decision.runtime == "copilot" + assert decision.model == "auto" + + def test_invalid_json_falls_back(self): + router = self._router(_cfg()) + decision = router.route( + prompt="hi", + last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=lambda rt, m, p, t: "I refuse to answer in JSON", + resolve_model=_identity_resolver, + ) + assert decision.source == "fallback" + assert decision.reason == "unparseable brain reply" + + def test_brain_exception_falls_back_never_raises(self): + router = self._router(_cfg()) + + def boom(rt, m, p, t): + raise TimeoutError("brain took too long") + + decision = router.route( + prompt="hi", last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=boom, + resolve_model=_identity_resolver, + ) + assert decision.source == "fallback" + assert "TimeoutError" in decision.reason + + def test_model_outside_pair_is_normalized_not_trusted(self): + router = self._router(_cfg()) + # Brain picks an eligible runtime but a model that isn't the + # allowlisted model for that runtime β€” must not be trusted verbatim. + decision = router.route( + prompt="hi", last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=lambda rt, m, p, t: '{"runtime":"copilot","model":"some-random-model","reason":"x"}', + resolve_model=lambda name, rt: name, + ) + assert decision.runtime == "copilot" + assert decision.model == "haiku" # allowlisted model for copilot, not the brain's free text + + def test_no_eligible_runtimes_uses_fallback(self): + # Allowlist candidates are unavailable; the fallback pair (copilot, + # not itself in the allowlist) is what should be used. + cfg = _cfg(allowlist=[{"runtime": "claude-sdk", "model": "opus", "hint": "x"}]) + router = self._router(cfg) + decision = router.route( + prompt="hi", last_routed=None, + runtime_available=lambda rt: rt == "copilot", # only fallback's runtime is up + invoke_brain=lambda *a: (_ for _ in ()).throw(AssertionError("brain should not be called")), + resolve_model=_identity_resolver, + ) + assert decision.source == "fallback" + assert decision.runtime == "copilot" + + def test_single_eligible_pair_skips_brain_call(self): + cfg = _cfg(allowlist=[{"runtime": "copilot", "model": "auto", "hint": "only one"}]) + router = self._router(cfg) + decision = router.route( + prompt="hi", last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=lambda *a: (_ for _ in ()).throw(AssertionError("brain should not be called")), + resolve_model=_identity_resolver, + ) + assert decision.source == "single" + assert decision.runtime == "copilot" + + def test_zero_eligible_and_no_fallback_returns_empty_decision(self): + router = self._router(_cfg()) + decision = router.route( + prompt="hi", last_routed=None, + runtime_available=lambda rt: False, # nothing is up, not even fallback + invoke_brain=lambda *a: (_ for _ in ()).throw(AssertionError("brain should not be called")), + resolve_model=_identity_resolver, + ) + assert decision.runtime == "" + assert decision.model == "" + assert "fallback also unavailable" in decision.reason + + def test_cooldown_excludes_runtime_from_eligibility(self): + cooldowns = llm_router.RuntimeCooldownTracker() + cooldowns.mark_failure("claude-sdk", "429", cooldown_seconds=60) + router = llm_router.LLMRouter(_FakeConfig(_cfg()), cooldowns) + # Only copilot remains eligible -> single-pair shortcut, no brain call. + decision = router.route( + prompt="hi", last_routed=None, + runtime_available=lambda rt: True, + invoke_brain=lambda *a: (_ for _ in ()).throw(AssertionError("brain should not be called")), + resolve_model=_identity_resolver, + ) + assert decision.runtime == "copilot" + assert decision.source == "single" + + def test_router_itself_never_eligible(self): + cfg = _cfg() + cfg["allowlist"].append({"runtime": "router", "model": "x"}) + router = self._router(cfg) + eligible = router.eligible_pairs(cfg["allowlist"], lambda rt: True) + assert all(e["runtime"] != "router" for e in eligible) + + +# --------------------------------------------------------------------------- +# Part B β€” agent_manager.py integration points +# --------------------------------------------------------------------------- + +os.environ.setdefault("WEE_AGENT_DIR", "/opt/wee-dev") + +import agent_manager # noqa: E402 + + +def _make_mgr(): + """Minimal SessionManager double, mirroring tests/test_issue125_429_retry.py.""" + mgr = agent_manager.SessionManager.__new__(agent_manager.SessionManager) + mgr.session_map = {} + mgr._session_map_lock = threading.Lock() + mgr.command_timeout = 60 + mgr._stream_buffers = {} + mgr.AGENTS = { + "orchestrator": {"path": "/opt", "primary_runtime": "claude", "primary_model": "haiku"}, + } + return mgr + + +@pytest.fixture(autouse=True) +def _reset_router_singletons(): + """Router config/cooldown singletons are module-level; reset between tests.""" + agent_manager._router_config = None + agent_manager._runtime_cooldown_tracker = None + agent_manager._llm_router = None + yield + agent_manager._router_config = None + agent_manager._runtime_cooldown_tracker = None + agent_manager._llm_router = None + + +class TestRuntimeRegistration: + def test_router_absent_when_disabled(self, monkeypatch): + monkeypatch.setenv("WEE_ROUTER_ENABLED", "0") + assert agent_manager.check_runtime_available("router") is False + ids = [rt["id"] for rt in agent_manager.get_available_runtimes()] + assert "router" not in ids + + def test_router_present_when_enabled_and_brain_available(self, monkeypatch, tmp_path): + cfg_path = tmp_path / "router_config.json" + cfg = _cfg(enabled=True, brain={"runtime": "copilot", "model": "auto"}) + agent_manager._router_config = llm_router.RouterConfig(cfg_path) + agent_manager._router_config.save(cfg) + monkeypatch.setenv("WEE_ROUTER_ENABLED", "1") + # "copilot" (the configured brain) availability shouldn't depend on + # whether the real CLI is installed on the test machine. + monkeypatch.setattr(agent_manager.shutil, "which", lambda name: "/usr/bin/copilot" if name == "copilot" else None) + + assert agent_manager.check_runtime_available("router") is True + ids = [rt["id"] for rt in agent_manager.get_available_runtimes()] + assert "router" in ids + + def test_router_absent_when_brain_unavailable(self, monkeypatch, tmp_path): + cfg_path = tmp_path / "router_config.json" + cfg = _cfg(enabled=True, brain={"runtime": "devin", "model": "x"}) + agent_manager._router_config = llm_router.RouterConfig(cfg_path) + agent_manager._router_config.save(cfg) + monkeypatch.setenv("WEE_ROUTER_ENABLED", "1") + monkeypatch.setattr(agent_manager.shutil, "which", lambda name: None) + + assert agent_manager.check_runtime_available("router") is False + + def test_env_override_wins_over_config_file(self, monkeypatch, tmp_path): + cfg_path = tmp_path / "router_config.json" + agent_manager._router_config = llm_router.RouterConfig(cfg_path) + agent_manager._router_config.save(_cfg(enabled=False)) + monkeypatch.setenv("WEE_ROUTER_ENABLED", "true") + assert agent_manager.is_router_enabled() is True + + monkeypatch.setenv("WEE_ROUTER_ENABLED", "false") + assert agent_manager.is_router_enabled() is False + + +class TestRunRouter: + def test_disabled_falls_back_to_agent_primary(self, monkeypatch): + mgr = _make_mgr() + monkeypatch.setattr(agent_manager, "is_router_enabled", lambda: False) + monkeypatch.setattr(mgr, "get_or_create_session_data", lambda sid: {}) + monkeypatch.setattr(mgr, "_resume_state_for_runtime", lambda *a, **k: False) + monkeypatch.setattr(mgr, "get_most_recent_session_id", lambda *a, **k: None) + monkeypatch.setattr(mgr, "update_session_field", lambda *a, **k: None) + + captured = {} + + def fake_dispatch(runtime, prompt, model, agent, session_id, can_resume, n8n_sid, timeout, render, mode): + captured["runtime"] = runtime + captured["model"] = model + return "ok" + + monkeypatch.setattr(mgr, "_dispatch_single_runtime", fake_dispatch) + + out = mgr.run_router("hello", "auto", "orchestrator", None, False, "sess1", 60, "plain") + assert out == "ok" + assert captured["runtime"] == "claude" + assert captured["model"] == "haiku" + + def test_valid_route_dispatches_to_decision_and_persists_router_last(self, monkeypatch, tmp_path): + mgr = _make_mgr() + cfg_path = tmp_path / "router_config.json" + agent_manager._router_config = llm_router.RouterConfig(cfg_path) + agent_manager._router_config.save(_cfg(enabled=True)) + monkeypatch.setenv("WEE_ROUTER_ENABLED", "1") + + monkeypatch.setattr(mgr, "get_or_create_session_data", lambda sid: {"router_last": None}) + monkeypatch.setattr(mgr, "_resume_state_for_runtime", lambda *a, **k: False) + monkeypatch.setattr(mgr, "get_most_recent_session_id", lambda *a, **k: None) + monkeypatch.setattr(mgr, "get_model_from_name", lambda name, rt: name) + monkeypatch.setattr(agent_manager, "check_runtime_available", lambda rt: True) + monkeypatch.setattr(agent_manager, "get_disabled_runtimes_manager", lambda: MagicMock(is_disabled=lambda rt: False)) + + persisted = {} + monkeypatch.setattr(mgr, "update_session_field", lambda sid, field, value: persisted.__setitem__(field, value)) + + dispatch_calls = [] + + def fake_dispatch(runtime, prompt, model, agent, session_id, can_resume, n8n_sid, timeout, render, mode=None): + dispatch_calls.append(runtime) + if runtime == "claude-sdk": # the real target + return "solved it" + return '{"runtime":"claude-sdk","model":"opus","reason":"coding task"}' # the brain + + monkeypatch.setattr(mgr, "_dispatch_single_runtime", fake_dispatch) + + out = mgr.run_router("write me quicksort", "auto", "orchestrator", None, False, "sess2", 60, "plain") + assert out == "solved it" + assert "claude-sdk" in dispatch_calls + assert persisted["router_last"]["runtime"] == "claude-sdk" + + def test_infra_failure_triggers_cooldown_and_single_fallback_retry(self, monkeypatch, tmp_path): + mgr = _make_mgr() + cfg_path = tmp_path / "router_config.json" + cfg = _cfg(enabled=True, allowlist=[{"runtime": "claude-sdk", "model": "opus", "hint": "x"}]) + agent_manager._router_config = llm_router.RouterConfig(cfg_path) + agent_manager._router_config.save(cfg) + monkeypatch.setenv("WEE_ROUTER_ENABLED", "1") + + monkeypatch.setattr(mgr, "get_or_create_session_data", lambda sid: {"router_last": None}) + monkeypatch.setattr(mgr, "_resume_state_for_runtime", lambda *a, **k: False) + monkeypatch.setattr(mgr, "get_most_recent_session_id", lambda *a, **k: None) + monkeypatch.setattr(mgr, "get_model_from_name", lambda name, rt: name) + monkeypatch.setattr(agent_manager, "check_runtime_available", lambda rt: True) + monkeypatch.setattr(agent_manager, "get_disabled_runtimes_manager", lambda: MagicMock(is_disabled=lambda rt: False)) + monkeypatch.setattr(mgr, "update_session_field", lambda *a, **k: None) + + dispatch_calls = [] + + def fake_dispatch(runtime, prompt, model, agent, session_id, can_resume, n8n_sid, timeout, render, mode=None): + dispatch_calls.append(runtime) + if runtime == "claude-sdk": + return "Error: 429 rate limited" + if runtime == "copilot": # fallback + return "recovered via fallback" + return "unused" + + monkeypatch.setattr(mgr, "_dispatch_single_runtime", fake_dispatch) + + out = mgr.run_router("hi", "auto", "orchestrator", None, False, "sess3", 60, "plain") + assert out == "recovered via fallback" + assert dispatch_calls.count("claude-sdk") == 1 # single eligible pair -> no brain call, one attempt + assert dispatch_calls.count("copilot") == 1 # exactly one fallback retry, not a loop + + tracker = agent_manager.get_runtime_cooldown_tracker() + assert tracker.is_cooling("claude-sdk") + + +class TestRouterConfigSchemaIntegration: + def test_config_schemas_validate_router_config_accepts_valid(self): + from config_schemas import validate_router_config + + validate_router_config(_cfg()) # should not raise + + def test_config_schemas_validate_router_config_rejects_recursive_brain(self): + from config_schemas import validate_router_config + import pydantic + + with pytest.raises(pydantic.ValidationError): + validate_router_config(_cfg(brain={"runtime": "router", "model": "x"})) + + def test_config_schemas_validate_router_config_rejects_missing_placeholders(self): + from config_schemas import validate_router_config + import pydantic + + with pytest.raises(pydantic.ValidationError): + validate_router_config(_cfg(prompt_template="no placeholders")) + + +class TestNoAutoRuntimeReintroduced: + """Guards against reintroducing the 'auto' runtime name removed in #84 + (see tests/test_auto_runtime_removed.py) while adding 'router'.""" + + def test_auto_still_absent_router_present_when_enabled(self, monkeypatch, tmp_path): + cfg_path = tmp_path / "router_config.json" + agent_manager._router_config = llm_router.RouterConfig(cfg_path) + agent_manager._router_config.save(_cfg(enabled=True, brain={"runtime": "copilot", "model": "auto"})) + monkeypatch.setenv("WEE_ROUTER_ENABLED", "1") + + ids = [rt["id"] for rt in agent_manager.get_all_runtimes()] + assert "auto" not in ids + assert "router" in ids diff --git a/webui/src/api/routerConfig.ts b/webui/src/api/routerConfig.ts new file mode 100644 index 00000000..827e037b --- /dev/null +++ b/webui/src/api/routerConfig.ts @@ -0,0 +1,128 @@ +/** + * API utilities for reading, writing, and testing the LLM router + * configuration (issue #506). Backed by /api/v1/router-config, + * /api/v1/router/test, and /api/v1/router/status. + */ + +const API_BASE = '/api/v1'; + +export interface RouterAllowlistEntry { + runtime: string; + model: string; + hint?: string; +} + +export interface RouterBrain { + runtime: string; + model: string; +} + +export interface RouterFallback { + runtime: string; + model: string; +} + +export interface RouterStickiness { + enabled: boolean; + prefer_same_runtime: boolean; + window_seconds: number; +} + +export interface RouterConfig { + enabled: boolean; + brain: RouterBrain; + timeout_seconds: number; + prompt_template: string; + allowlist: RouterAllowlistEntry[]; + fallback: RouterFallback; + stickiness: RouterStickiness; + cooldown_seconds: number; +} + +export interface RouterConfigResponse { + config: RouterConfig; + enabled_effective: boolean; + validation_errors: string[]; +} + +export interface RouterTestDecision { + runtime: string; + model: string; + reason: string; + source: 'router' | 'single' | 'fallback' | ''; + latency_ms: number; +} + +export interface RouterTestResponse { + decision: RouterTestDecision; + eligible_pairs: RouterAllowlistEntry[]; + total_ms: number; +} + +export interface RouterStatusResponse { + enabled: boolean; + brain: RouterBrain; + brain_available: boolean; + cooldowns: Record; +} + +/** Load the current router config, plus validation warnings and whether + * routing is effectively enabled after the WEE_ROUTER_ENABLED env override. */ +export async function loadRouterConfig(): Promise { + const resp = await fetch(`${API_BASE}/router-config`, { + headers: { 'Accept': 'application/json' }, + credentials: 'include', + }); + if (!resp.ok) { + throw new Error(`Failed to load router config: ${resp.status} ${resp.statusText}`); + } + return resp.json() as Promise; +} + +/** Save the router config. Throws with the server's validation message on 422. */ +export async function saveRouterConfig(config: RouterConfig): Promise<{ saved: boolean; config: RouterConfig }> { + const resp = await fetch(`${API_BASE}/router-config`, { + method: 'PUT', + headers: { 'Content-Type': 'application/json', 'Accept': 'application/json' }, + credentials: 'include', + body: JSON.stringify({ config }), + }); + if (!resp.ok) { + const detail = await resp.json().catch(() => null); + throw new Error(detail?.detail ?? `Failed to save router config: ${resp.status} ${resp.statusText}`); + } + return resp.json(); +} + +/** Dry-run a routing decision for a prompt without mutating any session β€” + * used by the "Test route" box to iterate on the prompt/allowlist. */ +export async function testRouterPrompt(prompt: string): Promise { + const resp = await fetch(`${API_BASE}/router/test`, { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'Accept': 'application/json' }, + credentials: 'include', + body: JSON.stringify({ prompt }), + }); + if (!resp.ok) { + const detail = await resp.json().catch(() => null); + throw new Error(detail?.detail ?? `Router test failed: ${resp.status} ${resp.statusText}`); + } + return resp.json() as Promise; +} + +/** Router health snapshot: enabled state, brain reachability, cooldowns. */ +export async function loadRouterStatus(): Promise { + const resp = await fetch(`${API_BASE}/router/status`, { + headers: { 'Accept': 'application/json' }, + credentials: 'include', + }); + if (!resp.ok) { + throw new Error(`Failed to load router status: ${resp.status} ${resp.statusText}`); + } + return resp.json() as Promise; +} + +/** Empty allowlist entry for the "+ Add pair" control. */ +export function emptyAllowlistEntry(): RouterAllowlistEntry { + return { runtime: '', model: '', hint: '' }; +} diff --git a/webui/src/components/RouterSettingsPanel.tsx b/webui/src/components/RouterSettingsPanel.tsx new file mode 100644 index 00000000..2862c576 --- /dev/null +++ b/webui/src/components/RouterSettingsPanel.tsx @@ -0,0 +1,349 @@ +/** + * RouterSettingsPanel β€” React/TypeScript component for configuring the LLM + * model router (issue #506). + * + * This is the reference/future-build implementation, following the same + * convention as AgentSettingsPanel.tsx in this directory: the live WebUI is + * currently vanilla JS (webui/dist/app.js). Wiring an equivalent panel into + * that bundle is tracked as a follow-up. + * + * Usage: + * import { RouterSettingsPanel } from './components/RouterSettingsPanel'; + * setOpen(false)} /> + */ + +import React, { useState, useEffect, useCallback } from 'react'; +import type { RouterConfig, RouterAllowlistEntry, RouterTestResponse } from '../api/routerConfig'; +import { loadRouterConfig, saveRouterConfig, testRouterPrompt, emptyAllowlistEntry } from '../api/routerConfig'; +import '../styles/RouterSettingsPanel.css'; + +// ─── Sub-components ────────────────────────────────────────────────────────── + +interface AllowlistRowProps { + entry: RouterAllowlistEntry; + onChange: (entry: RouterAllowlistEntry) => void; + onRemove: () => void; +} + +function AllowlistRow({ entry, onChange, onRemove }: AllowlistRowProps) { + return ( +
+ onChange({ ...entry, runtime: e.target.value })} + /> + onChange({ ...entry, model: e.target.value })} + /> + onChange({ ...entry, hint: e.target.value })} + /> + +
+ ); +} + +// ─── Toast ──────────────────────────────────────────────────────────────────── + +interface ToastProps { + message: string; + type: 'success' | 'error' | 'warning'; + onDismiss: () => void; +} + +function Toast({ message, type, onDismiss }: ToastProps) { + useEffect(() => { + const t = setTimeout(onDismiss, 4000); + return () => clearTimeout(t); + }, [onDismiss]); + + return ( +
+ {message} + +
+ ); +} + +// ─── Main component ─────────────────────────────────────────────────────────── + +interface RouterSettingsPanelProps { + onClose: () => void; +} + +function deepClone(obj: T): T { + return JSON.parse(JSON.stringify(obj)); +} + +export function RouterSettingsPanel({ onClose }: RouterSettingsPanelProps) { + const [draft, setDraft] = useState(null); + const [enabledEffective, setEnabledEffective] = useState(false); + const [serverValidationErrors, setServerValidationErrors] = useState([]); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [toast, setToast] = useState<{ message: string; type: 'success' | 'error' | 'warning' } | null>(null); + + const [testPrompt, setTestPrompt] = useState(''); + const [testing, setTesting] = useState(false); + const [testResult, setTestResult] = useState(null); + + // ── Load on mount ─────────────────────────────────────────────────────────── + useEffect(() => { + loadRouterConfig() + .then(resp => { + setDraft(deepClone(resp.config)); + setEnabledEffective(resp.enabled_effective); + setServerValidationErrors(resp.validation_errors); + }) + .catch(e => setToast({ message: (e as Error).message, type: 'error' })) + .finally(() => setLoading(false)); + }, []); + + const setField = (key: K, value: RouterConfig[K]) => { + setDraft(prev => prev ? { ...prev, [key]: value } : null); + }; + + const setAllowlistEntry = (idx: number, entry: RouterAllowlistEntry) => { + setDraft(prev => { + if (!prev) return null; + const allowlist = prev.allowlist.map((e, i) => (i === idx ? entry : e)); + return { ...prev, allowlist }; + }); + }; + + const addAllowlistEntry = () => { + setDraft(prev => prev ? { ...prev, allowlist: [...prev.allowlist, emptyAllowlistEntry()] } : null); + }; + + const removeAllowlistEntry = (idx: number) => { + setDraft(prev => prev ? { ...prev, allowlist: prev.allowlist.filter((_, i) => i !== idx) } : null); + }; + + // ── Save ──────────────────────────────────────────────────────────────────── + const handleSave = async () => { + if (!draft) return; + setSaving(true); + try { + const resp = await saveRouterConfig(draft); + setDraft(deepClone(resp.config)); + setServerValidationErrors([]); + setToast({ message: 'Router config saved.', type: 'success' }); + } catch (e: unknown) { + setToast({ message: (e as Error).message, type: 'error' }); + } finally { + setSaving(false); + } + }; + + // ── Test route ────────────────────────────────────────────────────────────── + const handleTest = useCallback(async () => { + if (!testPrompt.trim()) return; + setTesting(true); + setTestResult(null); + try { + const resp = await testRouterPrompt(testPrompt); + setTestResult(resp); + } catch (e: unknown) { + setToast({ message: (e as Error).message, type: 'error' }); + } finally { + setTesting(false); + } + }, [testPrompt]); + + // ── Render ────────────────────────────────────────────────────────────────── + return ( +
+
+ +
+

🧭 LLM Router

+ +
+ +
+ {loading &&
Loading router config…
} + + {serverValidationErrors.length > 0 && ( +
+ Current config has issues: + {serverValidationErrors.map((e, i) =>
{e}
)} +
+ )} + + {draft && ( + <> +
+
+ + + {enabledEffective ? 'Active' : 'Inactive'} + {!draft.enabled && enabledEffective ? ' (env override)' : ''} + +
+

+ When enabled, sessions with runtime router pick their + target runtime/model per message. Switch a session away with + /runtime <name>, and back with /runtime router. +

+
+ +
+

Brain (decision-making LLM)

+
+
+ + setField('brain', { ...draft.brain, runtime: e.target.value })} + /> +
+
+ + setField('brain', { ...draft.brain, model: e.target.value })} + /> +
+
+ + setField('timeout_seconds', Number(e.target.value))} + /> +
+
+
+ +
+

Allowlist β€” candidate runtime/model pairs

+
+ {draft.allowlist.map((entry, i) => ( + setAllowlistEntry(i, e)} + onRemove={() => removeAllowlistEntry(i)} + /> + ))} +
+ +
+ +
+

Fallback pair (used on brain failure / invalid decision)

+
+
+ + setField('fallback', { ...draft.fallback, runtime: e.target.value })} + /> +
+
+ + setField('fallback', { ...draft.fallback, model: e.target.value })} + /> +
+
+
+ +
+

Stickiness & cooldown

+
+ +
+ + setField('stickiness', { ...draft.stickiness, window_seconds: Number(e.target.value) })} + /> +
+
+ + setField('cooldown_seconds', Number(e.target.value))} + /> +
+
+
+ +
+

Routing prompt

+

+ Must include {'{allowlist_table}'} and {'{user_message}'}; + {'{stickiness_hint}'} is optional. +

+