diff --git a/console/web/src/hooks/use-llm-router-status.test.ts b/console/web/src/hooks/use-llm-router-status.test.ts new file mode 100644 index 000000000..f004222cc --- /dev/null +++ b/console/web/src/hooks/use-llm-router-status.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, it } from 'vitest' +import { + isLlmRouterAvailable, + LLM_ROUTER_WATCH_FN, + LLM_ROUTER_WORKER_NAME, +} from './use-llm-router-status' + +describe('llm-router presence probe wiring', () => { + it('probes the llm-router worker with its own watch handler id', () => { + // The model picker's router::* reads gate on THIS worker — gating on the + // harness blanked the picker whenever the harness was slow or absent. + expect(LLM_ROUTER_WORKER_NAME).toBe('llm-router') + // Must be unique per presence probe so the browser-local `worker` trigger + // handlers never collide (shell uses console::shell-watch). + expect(LLM_ROUTER_WATCH_FN).toBe('console::llm-router-watch') + }) + + it('gates on both presence and the initial probe settling', () => { + expect(isLlmRouterAvailable({ present: true, loading: false })).toBe(true) + expect(isLlmRouterAvailable({ present: true, loading: true })).toBe(false) + expect(isLlmRouterAvailable({ present: false, loading: false })).toBe(false) + }) +}) diff --git a/console/web/src/hooks/use-llm-router-status.ts b/console/web/src/hooks/use-llm-router-status.ts new file mode 100644 index 000000000..63ce69bf8 --- /dev/null +++ b/console/web/src/hooks/use-llm-router-status.ts @@ -0,0 +1,43 @@ +import { + isWorkerPresent, + useWorkerPresence, + type WorkerPresence, +} from './use-worker-presence' + +/** + * Presence probe for the `llm-router` worker. The router owns every + * `router::*` RPC the model picker reads (`provider::list`, `models::list`) + * and the change triggers it subscribes to, so provider/model UI gates on + * THIS worker's presence — gating on the harness starved the picker whenever + * the harness was slow or absent while the router was healthy. Thin wrapper + * over the generic worker-presence probe. + */ + +/** Engine worker name for the llm-router worker. */ +export const LLM_ROUTER_WORKER_NAME = 'llm-router' +/** Base id for the browser-local handler bound to the `worker` trigger. */ +export const LLM_ROUTER_WATCH_FN = 'console::llm-router-watch' + +export type LlmRouterStatus = WorkerPresence + +/** + * @param enabled - only run against the real backend; pass `false` for the + * mock/Storybook backend (treats the router as present so the picker shows + * in isolation). + */ +export function useLlmRouterStatus(enabled: boolean): LlmRouterStatus { + return useWorkerPresence({ + workerName: LLM_ROUTER_WORKER_NAME, + watchFnId: LLM_ROUTER_WATCH_FN, + enabled, + }) +} + +/** + * Whether the router's `router::*` functions are registered and safe to + * trigger. False during the initial presence probe and while the router is + * absent. + */ +export function isLlmRouterAvailable(status: LlmRouterStatus): boolean { + return isWorkerPresent(status) +} diff --git a/console/web/src/hooks/use-model-picker-source.ts b/console/web/src/hooks/use-model-picker-source.ts index 6f688e6e1..acb83fb9c 100644 --- a/console/web/src/hooks/use-model-picker-source.ts +++ b/console/web/src/hooks/use-model-picker-source.ts @@ -1,5 +1,6 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { onHarnessConfigSaved } from '@/lib/harness-config-events' +import { getIiiClient } from '@/lib/iii-client' import { catalogRowsToModelOptions, fetchModelsCatalog, @@ -23,11 +24,15 @@ import type { ModelOption } from '@/types/chat' * current by provider lifecycle events. Catalog refreshes never poll the * provider list. * - * `harnessAvailable` gates harness-owned RPCs until the worker is connected. + * `routerAvailable` gates the router-owned RPCs on the llm-router worker + * being connected — every read here is `router::*`. When it flips false→true + * (router installed, restarted, or just slow to boot) every effect below + * re-runs, so the picker recovers without a page reload. A WebSocket + * reconnect re-pulls both reads for the same reason. */ export function useModelPickerSource( backendId: string, - harnessAvailable = true, + routerAvailable = true, ): { modelOptions: ModelOption[] catalogKeys: string[] @@ -40,8 +45,14 @@ export function useModelPickerSource( [], ) const providerEventVersion = useRef(0) + // Mirror of `presentProviders` for event handlers: React state updaters may + // run deferred, so membership checks must not live inside them. + const providersRef = useRef([]) + useEffect(() => { + providersRef.current = presentProviders + }, [presentProviders]) const [catalogLoading, setCatalogLoading] = useState( - backendId === 'real' && harnessAvailable, + backendId === 'real' && routerAvailable, ) const refresh = useCallback(async () => { @@ -50,7 +61,7 @@ export function useModelPickerSource( setCatalogLoading(false) return } - if (!harnessAvailable) { + if (!routerAvailable) { setModelOptions([]) setCatalogLoading(false) return @@ -64,43 +75,46 @@ export function useModelPickerSource( } finally { setCatalogLoading(false) } - }, [backendId, harnessAvailable]) + }, [backendId, routerAvailable]) useEffect(() => { void refresh() }, [refresh]) - // One initial snapshot. Subsequent availability changes are applied from - // `router::provider::changed`, so model refreshes never re-read this list. - useEffect(() => { - if (backendId !== 'real' || !harnessAvailable) { + // Re-read `router::provider::list`, dropping the result if a newer provider + // event (or a newer snapshot) has advanced the version since we started. + const refreshProviders = useCallback(async () => { + if (backendId !== 'real' || !routerAvailable) { setPresentProviders([]) return } - let cancelled = false const snapshotVersion = providerEventVersion.current - void fetchProviderList() - .then((providers) => { - if (!cancelled && providerEventVersion.current === snapshotVersion) { - setPresentProviders(providers) - } - }) - .catch(() => { - if (!cancelled && providerEventVersion.current === snapshotVersion) { - setPresentProviders([]) - } - }) - return () => { - cancelled = true + try { + const providers = await fetchProviderList() + if (providerEventVersion.current === snapshotVersion) { + setPresentProviders(providers) + } + } catch { + if (providerEventVersion.current === snapshotVersion) { + setPresentProviders([]) + } } - }, [backendId, harnessAvailable]) + }, [backendId, routerAvailable]) - // Live updates: re-pull the catalog when the harness signals a model change - // (provider configured/cleared, refresh_models, CLI edits). The harness + // Initial snapshot (re-run when the router (re)appears). Availability flips + // are applied from `router::provider::changed`; an event for a provider the + // snapshot has never seen triggers a full re-read instead, so late-arriving + // providers render with their declared display name and capabilities. + useEffect(() => { + void refreshProviders() + }, [refreshProviders]) + + // Live updates: re-pull the catalog when the router signals a model change + // (provider configured/cleared, refresh_models, CLI edits). The router // coalesces bursts; the short trailing debounce here collapses any remaining // back-to-back pushes into a single re-read. useEffect(() => { - if (backendId !== 'real' || !harnessAvailable) return + if (backendId !== 'real' || !routerAvailable) return let disposed = false const disposers: (() => void)[] = [] let timer: ReturnType | null = null @@ -113,50 +127,81 @@ export function useModelPickerSource( }, 150) } - void subscribeModelChanges(onModelsChanged).then((dispose) => { - if (disposed) dispose() - else disposers.push(dispose) - }) + void subscribeModelChanges(onModelsChanged) + .then((dispose) => { + if (disposed) dispose() + else disposers.push(dispose) + }) + // Setup failure degrades to manual refresh; never an unhandled rejection. + .catch(() => {}) void subscribeProviderChanges(({ provider, op }) => { providerEventVersion.current += 1 - setPresentProviders((current) => { - const available = op !== 'unavailable' - const existing = current.find((entry) => entry.id === provider) - if (existing) { - if (existing.available === available) return current - return current.map((entry) => - entry.id === provider ? { ...entry, available } : entry, - ) - } - return [ - ...current, - { - id: provider, - display_name: provider, - supports_model_listing: true, - available, - }, - ] - }) - }).then((dispose) => { - if (disposed) dispose() - else disposers.push(dispose) + if (op === 'unregister') { + setPresentProviders((current) => + current.filter((entry) => entry.id !== provider), + ) + return + } + // A provider the snapshot never saw: re-read the list rather than + // inventing a degraded entry (raw id as display name, guessed + // capabilities). + if (!providersRef.current.some((entry) => entry.id === provider)) { + void refreshProviders() + return + } + const available = op !== 'unavailable' + setPresentProviders((current) => + current.map((entry) => + entry.id === provider && entry.available !== available + ? { ...entry, available } + : entry, + ), + ) }) + .then((dispose) => { + if (disposed) dispose() + else disposers.push(dispose) + }) + // Setup failure degrades to snapshot re-reads; never an unhandled rejection. + .catch(() => {}) return () => { disposed = true if (timer !== null) clearTimeout(timer) for (const d of disposers) d() } - }, [backendId, harnessAvailable, refresh]) + }, [backendId, routerAvailable, refresh, refreshProviders]) + + // A WebSocket drop loses any change events fired while disconnected; + // re-pull both reads when the connection comes back. + useEffect(() => { + if (backendId !== 'real' || !routerAvailable) return + let disposed = false + let offConn: (() => void) | null = null + getIiiClient() + .then((client) => { + if (disposed) return + offConn = client.addConnectionStateListener((state) => { + if (state === 'connected') { + void refresh() + void refreshProviders() + } + }) + }) + .catch(() => {}) + return () => { + disposed = true + offConn?.() + } + }, [backendId, routerAvailable, refresh, refreshProviders]) useEffect(() => { - if (backendId !== 'real' || !harnessAvailable) return + if (backendId !== 'real' || !routerAvailable) return return onHarnessConfigSaved(() => { void refresh() }) - }, [backendId, harnessAvailable, refresh]) + }, [backendId, routerAvailable, refresh]) const catalogKeys = useMemo( () => modelOptions.map((o) => o.id), diff --git a/console/web/src/lib/conversations-context.tsx b/console/web/src/lib/conversations-context.tsx index da47ec59b..e9d814ddc 100644 --- a/console/web/src/lib/conversations-context.tsx +++ b/console/web/src/lib/conversations-context.tsx @@ -15,9 +15,12 @@ import { } from '@/hooks/use-conversations' import { type HarnessStatus, - isHarnessAvailable, useHarnessStatus, } from '@/hooks/use-harness-status' +import { + isLlmRouterAvailable, + useLlmRouterStatus, +} from '@/hooks/use-llm-router-status' import { isMemoryAvailable, useMemoryStatus } from '@/hooks/use-memory-status' import { useModelPickerSource } from '@/hooks/use-model-picker-source' import { isShellAvailable, useShellStatus } from '@/hooks/use-shell-status' @@ -102,7 +105,11 @@ export function ConversationsProvider({ children, }: ConversationsProviderProps) { const harnessStatus = useHarnessStatus(backend.id === 'real') - const harnessAvailable = isHarnessAvailable(harnessStatus) + // The model picker reads router-owned RPCs; gate them on llm-router, not + // the harness — the harness being slow or absent must not blank the picker. + const routerAvailable = isLlmRouterAvailable( + useLlmRouterStatus(backend.id === 'real'), + ) const approvalGateAvailable = isApprovalGateAvailable( useApprovalGateStatus(backend.id === 'real'), ) @@ -119,7 +126,7 @@ export function ConversationsProvider({ catalogLoading, presentProviders, refresh, - } = useModelPickerSource(backend.id, harnessAvailable) + } = useModelPickerSource(backend.id, routerAvailable) // Conversations are backed by the session-manager worker on the real // backend; mocks stay in-memory. const api = useConversations( @@ -130,7 +137,7 @@ export function ConversationsProvider({ const [refreshingModels, setRefreshingModels] = useState(false) const refreshModels = useCallback(async () => { - if (!harnessAvailable) return + if (!routerAvailable) return setRefreshingModels(true) try { if (backend.id === 'real') { @@ -146,7 +153,7 @@ export function ConversationsProvider({ } finally { setRefreshingModels(false) } - }, [harnessAvailable, refresh, presentProviders]) + }, [routerAvailable, refresh, presentProviders]) const value: ConversationsContextValue = { ...api, diff --git a/console/web/src/pages/Configuration/tabs/WorkersTab/hooks.ts b/console/web/src/pages/Configuration/tabs/WorkersTab/hooks.ts index ea6ebb348..ebc47eb0d 100644 --- a/console/web/src/pages/Configuration/tabs/WorkersTab/hooks.ts +++ b/console/web/src/pages/Configuration/tabs/WorkersTab/hooks.ts @@ -5,9 +5,11 @@ */ import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query' +import { useEffect } from 'react' import { useWorkerLifecycle } from '@/hooks/use-worker-lifecycle' import { getDefaultBackend } from '@/lib/backend' import { notifyHarnessConfigSaved } from '@/lib/harness-config-events' +import { subscribeProviderChanges } from '@/lib/models-catalog' import { type ConfigurationSchemaView, getConfiguration, @@ -54,11 +56,20 @@ export function useConfigurationsList() { const WORKER_REGISTRY_WATCH_FN = 'console::workers-config-watch' /** - * Keep the worker configuration list fresh when workers are added or removed - * out of band — e.g. `iii worker add harness` run in a terminal — purely off - * the engine's `worker` lifecycle trigger (its push channel). No polling. - * Each event re-pulls `configuration::list` so a freshly-installed worker's - * entry appears once it registers. + * Keep the worker configuration registry fresh when workers change out of + * band — e.g. `iii worker add harness` run in a terminal — purely off the + * engine's `worker` lifecycle trigger (its push channel). No polling. + * + * Invalidation covers the whole `['configuration']` namespace, not just the + * list: a worker arriving can re-compose an entry SCHEMA an open editor is + * rendering (llm-router recomposes its entry from every provider + * declaration), and a schema query that never refetches keeps showing zero + * provider cards. + * + * llm-router's provider registrations don't always ride the `worker` + * lifecycle (a provider process restart replays no `add`), so the hook also + * listens to `router::provider::changed` and invalidates the llm-router + * entry's queries on every provider register/unregister. */ export function useWorkerRegistryReactivity(): void { const qc = useQueryClient() @@ -69,9 +80,35 @@ export function useWorkerRegistryReactivity(): void { fnId: WORKER_REGISTRY_WATCH_FN, operations: ['add', 'remove'], onEvent: () => { - qc.invalidateQueries({ queryKey: configurationKeys.list() }) + qc.invalidateQueries({ queryKey: configurationKeys.all }) }, }) + + useEffect(() => { + if (!enabled) return + let disposed = false + const disposers: (() => void)[] = [] + void subscribeProviderChanges(() => { + qc.invalidateQueries({ queryKey: configurationKeys.schema('llm-router') }) + qc.invalidateQueries({ + queryKey: configurationKeys.rawValue('llm-router'), + }) + qc.invalidateQueries({ + queryKey: configurationKeys.expandedValue('llm-router'), + }) + }) + .then((dispose) => { + if (disposed) dispose() + else disposers.push(dispose) + }) + // Setup failure (engine unreachable) degrades to the lifecycle-driven + // invalidation above; never an unhandled rejection. + .catch(() => {}) + return () => { + disposed = true + for (const d of disposers) d() + } + }, [enabled, qc]) } export function useConfigurationSchema(id: string | null | undefined) { @@ -108,6 +145,11 @@ export function useSetConfiguration(id: string | null | undefined) { qc.invalidateQueries({ queryKey: configurationKeys.expandedValue(targetId), }) + // A save can change what the entry's schema derives (llm-router + // recomposes provider cards from configured credentials). + qc.invalidateQueries({ + queryKey: configurationKeys.schema(targetId), + }) } qc.invalidateQueries({ queryKey: configurationKeys.list() }) }, diff --git a/crates/console-ui/src/lib.rs b/crates/console-ui/src/lib.rs index 27831c43a..691cb863d 100644 --- a/crates/console-ui/src/lib.rs +++ b/crates/console-ui/src/lib.rs @@ -201,8 +201,11 @@ impl ConsoleUi { } /// Register the content function and one trigger per asset; spawn the - /// dev watcher when the watch env var is set. Call once, after the - /// worker's regular functions are registered. + /// dev watcher when the watch env var is set. Call once at boot — order + /// relative to the worker's regular functions is free, and registering + /// the UI FIRST is preferred: every moment before the console:script + /// announcement is a window where an open console renders its generic + /// fallback instead of this worker's injected UI. /// /// Trigger registration failures are warn-logged, not fatal — injected /// UI is an accessory, and the SDK replays surviving registrations on @@ -261,11 +264,24 @@ impl ConsoleUi { prev: spec.content.clone(), }); } - Err(e) => tracing::warn!( - error = %e, - path = spec.path, - "failed to register console ui trigger" - ), + Err(e) => { + tracing::warn!( + error = %e, + path = spec.path, + "failed to register console ui trigger; retrying in background" + ); + // One failed attempt used to mean the generic fallback UI + // for the console's whole session. Retry until it lands — + // the asset is static, so late is strictly better than + // never. (Late assets skip the dev watcher; it is a + // dev-only convenience.) + spawn_registration_retry( + iii.clone(), + self.content_function_id.clone(), + spec.kind, + spec.path.clone(), + ); + } } } @@ -375,6 +391,42 @@ impl Served { } } +/// Keep retrying a failed asset-trigger registration with capped backoff. +/// The handle can be dropped: unregistration is explicit, never on drop. +/// No-ops (with a warning) outside a Tokio runtime — `register` is a sync +/// API and `tokio::spawn` would panic there; behavior then matches the +/// pre-retry contract of one attempt. +fn spawn_registration_retry( + iii: Arc, + function_id: String, + kind: AssetKind, + path: String, +) { + let Ok(handle) = tokio::runtime::Handle::try_current() else { + tracing::warn!( + path, + "console ui trigger registration failed outside a tokio runtime; not retrying" + ); + return; + }; + handle.spawn(async move { + let mut delay = std::time::Duration::from_secs(2); + loop { + tokio::time::sleep(delay).await; + match register_asset_trigger(&iii, &function_id, kind, &path) { + Ok(_handle) => { + tracing::info!(path, "registered console ui asset (after retry)"); + return; + } + Err(e) => { + tracing::debug!(error = %e, path, "console ui trigger retry failed"); + delay = (delay * 2).min(std::time::Duration::from_secs(30)); + } + } + } + }); +} + fn register_asset_trigger( iii: &Arc, function_id: &str, diff --git a/eval/Cargo.lock b/eval/Cargo.lock index ce0d706cc..0ebaf8a10 100644 --- a/eval/Cargo.lock +++ b/eval/Cargo.lock @@ -532,7 +532,7 @@ dependencies = [ [[package]] name = "harness" -version = "1.7.4" +version = "1.8.1" dependencies = [ "anyhow", "async-trait", diff --git a/eval/src/runtime.rs b/eval/src/runtime.rs index 7be7d4490..74e9904f9 100644 --- a/eval/src/runtime.rs +++ b/eval/src/runtime.rs @@ -678,11 +678,11 @@ async fn send_finalization( fn send_options(job: &EvalJobRecordV1, variant: &crate::contract::EvalVariantV1) -> SendOptions { SendOptions { system_prompt: variant.system_prompt.clone(), - system_prompt_strategy: if variant.system_prompt.is_none() { + system_prompt_strategy: Some(if variant.system_prompt.is_none() { harness::prompt::SystemPromptStrategy::Disabled } else { job.request.model.system_prompt_strategy - }, + }), mode: job.request.model.mode, max_turns: Some(job.request.limits.execution.max_turns), max_output_tokens: Some(job.request.limits.execution.max_output_tokens_per_call), diff --git a/eval/tests/golden/schemas/eval.assert.exact.json b/eval/tests/golden/schemas/eval.assert.exact.json index 42071cae9..472186ab1 100644 --- a/eval/tests/golden/schemas/eval.assert.exact.json +++ b/eval/tests/golden/schemas/eval.assert.exact.json @@ -39,6 +39,14 @@ "null" ] }, + "session_cost_usd": { + "description": "Running cost of the whole session in USD, accumulated across every generation step. `usage.cost_usd` is one step's bill — on providers with steep cache discounts the per-step number swings two orders of magnitude, so a chip showing it alone reads as a bouncing total.", + "format": "double", + "type": [ + "number", + "null" + ] + }, "session_id": { "type": "string" }, diff --git a/eval/tests/golden/schemas/eval.assert.normalized_text.json b/eval/tests/golden/schemas/eval.assert.normalized_text.json index d5c8744df..9d6842fdc 100644 --- a/eval/tests/golden/schemas/eval.assert.normalized_text.json +++ b/eval/tests/golden/schemas/eval.assert.normalized_text.json @@ -39,6 +39,14 @@ "null" ] }, + "session_cost_usd": { + "description": "Running cost of the whole session in USD, accumulated across every generation step. `usage.cost_usd` is one step's bill — on providers with steep cache discounts the per-step number swings two orders of magnitude, so a chip showing it alone reads as a bouncing total.", + "format": "double", + "type": [ + "number", + "null" + ] + }, "session_id": { "type": "string" }, diff --git a/eval/tests/golden/schemas/eval.compare-sessions.json b/eval/tests/golden/schemas/eval.compare-sessions.json index 16c9a4bd0..a8dbe79a8 100644 --- a/eval/tests/golden/schemas/eval.compare-sessions.json +++ b/eval/tests/golden/schemas/eval.compare-sessions.json @@ -63,6 +63,14 @@ "null" ] }, + "session_cost_usd": { + "description": "Running cost of the whole session in USD, accumulated across every generation step. `usage.cost_usd` is one step's bill — on providers with steep cache discounts the per-step number swings two orders of magnitude, so a chip showing it alone reads as a bouncing total.", + "format": "double", + "type": [ + "number", + "null" + ] + }, "session_id": { "type": "string" }, diff --git a/eval/tests/golden/schemas/eval.result.json b/eval/tests/golden/schemas/eval.result.json index 814c8abbf..2d98d6318 100644 --- a/eval/tests/golden/schemas/eval.result.json +++ b/eval/tests/golden/schemas/eval.result.json @@ -157,6 +157,14 @@ "null" ] }, + "session_cost_usd": { + "description": "Running cost of the whole session in USD, accumulated across every generation step. `usage.cost_usd` is one step's bill — on providers with steep cache discounts the per-step number swings two orders of magnitude, so a chip showing it alone reads as a bouncing total.", + "format": "double", + "type": [ + "number", + "null" + ] + }, "session_id": { "type": "string" }, diff --git a/harness/src/budget.rs b/harness/src/budget.rs index f69f9281b..c030526ce 100644 --- a/harness/src/budget.rs +++ b/harness/src/budget.rs @@ -232,8 +232,10 @@ pub async fn reserve( let Some(pricing) = model.and_then(|model| model.pricing) else { return Ok(ReserveOutcome::Rejected(BudgetRejection::Unavailable( format!( - "cannot enforce max_cost_usd for model {} because no pricing is configured in \ - the model catalog; configure input/output pricing or remove max_cost_usd", + "cannot enforce max_cost_usd for model {}: the model catalog returned no \ + pricing (llm-router absent/unreachable, provider not registered yet, or no \ + input/output pricing configured); retry once the router is up, configure \ + pricing, or remove max_cost_usd", record.options.model ), ))); diff --git a/harness/src/clients/router.rs b/harness/src/clients/router.rs index b770ba70c..3e727ce31 100644 --- a/harness/src/clients/router.rs +++ b/harness/src/clients/router.rs @@ -380,13 +380,19 @@ impl RouterClient { .clone() .or_else(|| terminal_error.clone()) .or_else(|| Some("router produced no terminal frame".to_string())); - // A dispatch/ack rejection (e.g. provider unavailable) is - // authoritative — retrying in-turn cannot help. Only a - // frame-less-but-acked stream stays classified transient. - m.error_kind = Some(if response_error.is_some() { - crate::types::event::ErrorKind::Permanent - } else { - crate::types::event::ErrorKind::Transient + // Topology rejections (router absent, provider not registered + // yet, provider marked down) are transient by construction: they + // are exactly what a router/provider startup race produces, and + // the router's re-discovery heals them within seconds. Killing + // the turn permanently for those raced every boot. Other + // dispatch/ack rejections stay authoritative — retrying in-turn + // cannot help. + m.error_kind = Some(match &response_error { + Some(msg) if dispatch_error_is_transient(msg) => { + crate::types::event::ErrorKind::Transient + } + Some(_) => crate::types::event::ErrorKind::Permanent, + None => crate::types::event::ErrorKind::Transient, }); m }); @@ -766,6 +772,24 @@ fn enrich_streaming_args( out } +/// Dispatch/ack rejections that describe TOPOLOGY, not the request: the +/// router function is absent (engine `function_not_found`), the provider is +/// not (yet) registered, or the provider is marked down. All of these are the +/// startup/restart race the router's re-discovery repairs; a bounded +/// transient resume rides it out. A genuinely wrong provider name also +/// matches "unknown provider" — it burns the bounded resumes and then fails, +/// which is the acceptable cost of not killing every boot-race turn. +fn dispatch_error_is_transient(msg: &str) -> bool { + let m = msg.to_ascii_lowercase(); + m.contains("function_not_found") + || m.contains("function not found") + || m.contains("provider_unavailable") + // "provider unavailable" — the router's dispatch-time wording. + || (m.contains("provider") && m.contains("unavailable")) + || m.contains("unknown provider") + || m.contains("no provider registered") +} + fn utf8_tail(s: &str, max: usize) -> &str { if s.len() <= max { return s; @@ -783,6 +807,25 @@ mod streaming_args_tests { use crate::types::content::ContentBlock; use std::collections::HashMap; + #[test] + fn topology_rejections_are_transient_other_rejections_permanent() { + assert!(dispatch_error_is_transient( + "Function not found: router::chat" + )); + assert!(dispatch_error_is_transient("router/function_not_found")); + assert!(dispatch_error_is_transient("unknown provider anthropic")); + assert!(dispatch_error_is_transient( + "provider anthropic unavailable (worker not connected)" + )); + assert!(dispatch_error_is_transient( + "no provider registered for model gpt-5" + )); + assert!(!dispatch_error_is_transient("invalid request: bad payload")); + assert!(!dispatch_error_is_transient( + "upstream 401: invalid api key" + )); + } + #[test] fn enrich_injects_tail_only_while_args_are_unparsed() { let mut m = empty_assistant("p", "m"); diff --git a/harness/src/functions/send.rs b/harness/src/functions/send.rs index f8c60c6ee..379bf64ba 100644 --- a/harness/src/functions/send.rs +++ b/harness/src/functions/send.rs @@ -224,6 +224,30 @@ pub async fn start(deps: &Deps, req: SendRequest) -> Result, mode: Option) -> bool { + prompt + == Some( + prompt::build_system_prompt(prompt::SystemPromptOpts { + mode, + identity: None, + }) + .as_str(), + ) +} + /// Default a BRAND-NEW session's working directory (MOT-3897): when the very /// first turn arrives without `metadata.fs_scope.root`, scope it to the /// configured default (the stack's launch folder). Existing sessions are never @@ -804,6 +845,41 @@ pub(crate) async fn seed_new( mod tests { use super::*; + /// The steer-time heal only fires on the exact embedded fallback: a + /// session born during a router outage re-resolves, everything a caller + /// or provider shaped stays frozen as designed. + #[test] + fn pure_fallback_prompt_is_detected_for_reresolve() { + let fallback = prompt::build_system_prompt(prompt::SystemPromptOpts { + mode: None, + identity: None, + }); + assert!(inherited_prompt_is_pure_fallback(Some(&fallback), None)); + + // A mode-prefixed fallback matches only under the same mode. + let ask = prompt::build_system_prompt(prompt::SystemPromptOpts { + mode: Some(Mode::Ask), + identity: None, + }); + assert!(inherited_prompt_is_pure_fallback( + Some(&ask), + Some(Mode::Ask) + )); + assert!(!inherited_prompt_is_pure_fallback(Some(&ask), None)); + + // Provider identity, caller override, and disabled prompts never match. + let identity = prompt::build_system_prompt(prompt::SystemPromptOpts { + mode: None, + identity: Some("You are a provider identity."), + }); + assert!(!inherited_prompt_is_pure_fallback(Some(&identity), None)); + assert!(!inherited_prompt_is_pure_fallback( + Some("custom prompt"), + None + )); + assert!(!inherited_prompt_is_pure_fallback(None, None)); + } + #[test] fn string_message_becomes_user_text() { let m = normalize_message(MessageInput::Text("hi".into())).unwrap(); diff --git a/llm-router/build.rs b/llm-router/build.rs index cf090b655..3ddb4b182 100644 --- a/llm-router/build.rs +++ b/llm-router/build.rs @@ -36,6 +36,10 @@ fn main() { // project links @iii-dev/console-ui from packages/console-ui). println!("cargo:rerun-if-changed=../pnpm-lock.yaml"); println!("cargo:rerun-if-changed=ui/tsconfig.json"); + // The workspace-linked host-API package the bundle compiles against. + println!("cargo:rerun-if-changed=../packages/console-ui/src"); + // Toggling the skip must re-evaluate the embed, not reuse a cached one. + println!("cargo:rerun-if-env-changed=SKIP_UI_BUILD"); let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); let ui_dir = manifest_dir.join("ui"); @@ -61,6 +65,14 @@ fn main() { ); } } + // Reaching here means dist_is_fresh said NO: the embedded bundle is + // older than its sources. Say so loudly — a stale embed ships an old + // config UI to every console. + println!( + "cargo:warning=SKIP_UI_BUILD set but ui/dist is STALE (older than \ + ui sources); embedding the old bundle — rebuild with \ + `cd ui && pnpm install && pnpm build`" + ); return; } @@ -115,6 +127,20 @@ fn dist_is_fresh(dist_asset: &Path, ui_dir: &Path) -> bool { ui_dir.join("../../pnpm-lock.yaml"), ui_dir.join("tsconfig.json"), ]; + // The linked @iii-dev/console-ui package: its types gate `tsc --noEmit` + // and its API shape is what the bundle runs against. Absent is fine + // (published-tarball builds); any OTHER metadata error means we cannot + // inspect the tree, and this file's policy is conservative — rebuild. + let linked_pkg_src = ui_dir.join("../../packages/console-ui/src"); + match std::fs::metadata(&linked_pkg_src) { + Ok(_) => { + if !subtree_older_than(&linked_pkg_src, dist_mtime) { + return false; + } + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(_) => return false, + } for f in watched_files.iter() { if !f.exists() { continue; diff --git a/llm-router/src/catalog/store.rs b/llm-router/src/catalog/store.rs index 7bd1cf491..c1bc680ea 100644 --- a/llm-router/src/catalog/store.rs +++ b/llm-router/src/catalog/store.rs @@ -82,4 +82,21 @@ impl CatalogStore { let value = serde_json::to_value(&*slices).unwrap_or_default(); state_set(&self.iii, CATALOG_KEY, value).await } + + /// Drop a provider's slice (unregister path) so stale models never route. + /// Returns whether a slice existed. Persist-then-swap: memory only + /// changes after the durable write succeeds, so a failed persist leaves + /// both in agreement and the retry starts clean. + pub async fn remove_slice(&self, provider: &str) -> Result { + let mut slices = self.slices.lock().await; // serialized writer + if !slices.contains_key(provider) { + return Ok(false); + } + let mut next = slices.clone(); + next.remove(provider); + let value = serde_json::to_value(&next).unwrap_or_default(); + state_set(&self.iii, CATALOG_KEY, value).await?; + *slices = next; + Ok(true) + } } diff --git a/llm-router/src/main.rs b/llm-router/src/main.rs index 82d37f360..f4ffaf8ef 100644 --- a/llm-router/src/main.rs +++ b/llm-router/src/main.rs @@ -85,9 +85,13 @@ async fn main() -> Result<(), Box> { }, ); - register_router(iii.clone()).await?; + // UI assets first: register_router awaits engine round-trips (state + // restore, config reconcile), and every millisecond before the + // console:script announcement is a window where an open console paints + // the generic schema form instead of this worker's config UI. #[cfg(feature = "console-ui")] llm_router::ui::register(&iii); + register_router(iii.clone()).await?; tracing::info!(url = %cli.url, "llm-router registered"); tokio::signal::ctrl_c().await?; diff --git a/llm-router/src/register.rs b/llm-router/src/register.rs index 4738fcf62..966f5cc9c 100644 --- a/llm-router/src/register.rs +++ b/llm-router/src/register.rs @@ -30,7 +30,7 @@ use crate::config::on_changed::make_on_config_changed; use crate::config::schema::provider_entry_schema; use crate::config::state::{new_config_cell, ConfigCell}; use crate::registry::availability::make_provider_list; -use crate::registry::register::make_provider_register; +use crate::registry::register::{make_provider_register, make_provider_unregister}; use crate::registry::resolve::{make_provider_resolve, make_update_credential}; use crate::registry::store::RegistryStore; use crate::surface; @@ -205,6 +205,21 @@ pub async fn register_router(iii: IIIClient) -> Result { .description(surface::PROVIDER_REGISTER_DESC) .metadata(internal_meta()), ); + iii.register_function( + surface::PROVIDER_UNREGISTER_ID, + RegisterFunction::new_async_with_bad_request( + make_provider_unregister( + iii.clone(), + registry.clone(), + catalog.clone(), + entry_lock.clone(), + events.clone(), + ), + invalid_request_from_serde, + ) + .description(surface::PROVIDER_UNREGISTER_DESC) + .metadata(internal_meta()), + ); iii.register_function( surface::PROVIDER_RESOLVE_ID, RegisterFunction::new_async(make_provider_resolve(config.clone(), registry.clone())) @@ -316,7 +331,8 @@ pub async fn register_router(iii: IIIClient) -> Result { // returning provider is resolvable in seconds instead of minutes. { let iii_handler = iii.clone(); - let sweep = crate::registry::rediscover::spawn_debounced_sweep(iii_handler); + let sweep = + crate::registry::rediscover::spawn_debounced_sweep(iii_handler, registry.clone()); iii.register_function( surface::ON_FUNCTIONS_CHANGED_ID, RegisterFunction::new_async(move |_event: FunctionsChangedEvent| { diff --git a/llm-router/src/registry/rediscover.rs b/llm-router/src/registry/rediscover.rs index 364ac762e..e94a3d769 100644 --- a/llm-router/src/registry/rediscover.rs +++ b/llm-router/src/registry/rediscover.rs @@ -27,10 +27,11 @@ //! deterministic per-provider handler (the same one the fan-out targets), which //! makes it both the discovery key and the thing to call. +use crate::registry::store::RegistryStore; use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; use serde_json::{json, Value}; -use std::collections::HashSet; +use std::sync::Arc; /// The handler suffix that identifies a provider and receives the nudge. const READY_HANDLER_SUFFIX: &str = "::on_router_ready"; @@ -86,14 +87,24 @@ fn provider_id_of(function_id: &str) -> Option<&str> { (!id.is_empty() && !id.contains("::")).then_some(id) } -fn newly_live_provider_ids(known: &mut HashSet, live: Vec) -> Vec { - let added = live - .iter() - .filter(|id| !known.contains(*id)) - .cloned() - .collect(); - *known = live.into_iter().collect(); - added +/// Live providers the router does not currently hold as registered AND +/// available. A newly live provider has no record; one that missed the ready +/// fan-out (or whose registry was boot-reset) has a record with +/// `available=false`. Both need the nudge. Registered-and-up providers are +/// skipped so routine registration churn never spams healthy providers. +async fn stale_provider_ids(registry: &RegistryStore, live: Vec) -> Vec { + let mut stale = Vec::new(); + for id in live { + let up = registry + .get(&id) + .await + .map(|r| r.available) + .unwrap_or(false); + if !up { + stale.push(id); + } + } + stale } async fn nudge_provider_ids(iii: &IIIClient, ids: &[String]) { @@ -120,8 +131,8 @@ pub async fn nudge_live_providers(iii: &IIIClient) -> usize { /// Coalescing handle for the re-discovery sweep. `engine::functions-available` /// fires for every function registration change, including unrelated console -/// subscriptions. The sweep tracks provider membership and nudges only newly -/// live providers after the burst goes quiet. +/// subscriptions. After the burst goes quiet the sweep nudges every live +/// provider the registry does not hold as available. #[derive(Clone)] pub struct SweepHandle { tx: tokio::sync::mpsc::Sender<()>, @@ -139,10 +150,9 @@ impl SweepHandle { const SWEEP_DEBOUNCE: std::time::Duration = std::time::Duration::from_secs(3); /// Spawn the sweep task and return its handle. -pub fn spawn_debounced_sweep(iii: IIIClient) -> SweepHandle { +pub fn spawn_debounced_sweep(iii: IIIClient, registry: Arc) -> SweepHandle { let (tx, mut rx) = tokio::sync::mpsc::channel::<()>(1); tokio::spawn(async move { - let mut known = HashSet::new(); while rx.recv().await.is_some() { // Drain the rest of the burst, extending the quiet period each // time another request arrives. @@ -151,12 +161,12 @@ pub fn spawn_debounced_sweep(iii: IIIClient) -> SweepHandle { .is_ok_and(|received| received.is_some()) {} let live = live_provider_ids(&iii).await; - let added = newly_live_provider_ids(&mut known, live); - let count = added.len(); - nudge_provider_ids(&iii, &added).await; + let stale = stale_provider_ids(®istry, live).await; + let count = stale.len(); + nudge_provider_ids(&iii, &stale).await; tracing::debug!( providers = count, - "registration change: nudged newly live providers to re-declare" + "registration change: nudged live-but-unavailable providers to re-declare" ); } }); @@ -189,21 +199,4 @@ mod tests { // A nested id would make `provider::::on_router_ready` ambiguous. assert_eq!(provider_id_of("provider::a::b::on_router_ready"), None); } - - #[test] - fn provider_set_diff_only_returns_new_ids() { - let mut known = std::collections::HashSet::from(["anthropic".to_string()]); - - assert!(newly_live_provider_ids(&mut known, vec!["anthropic".into()]).is_empty()); - assert_eq!( - newly_live_provider_ids(&mut known, vec!["anthropic".into(), "openai-codex".into()]), - vec!["openai-codex"] - ); - - assert!(newly_live_provider_ids(&mut known, vec!["openai-codex".into()]).is_empty()); - assert_eq!( - newly_live_provider_ids(&mut known, vec!["anthropic".into(), "openai-codex".into()]), - vec!["anthropic"] - ); - } } diff --git a/llm-router/src/registry/register.rs b/llm-router/src/registry/register.rs index efaf0a5b7..c427eec85 100644 --- a/llm-router/src/registry/register.rs +++ b/llm-router/src/registry/register.rs @@ -8,7 +8,10 @@ use std::collections::BTreeMap; use std::sync::Arc; use crate::types::errors::{RouterCode, RouterError}; -use crate::types::router::{ProviderRegisterRequest, ProviderRegisterResponse}; +use crate::types::router::{ + ProviderRegisterRequest, ProviderRegisterResponse, ProviderUnregisterRequest, + ProviderUnregisterResponse, +}; use futures::future::BoxFuture; use iii_sdk::{errors::Error, IIIClient}; use serde_json::{json, Value}; @@ -76,6 +79,14 @@ pub fn make_provider_register( let worker_id = declaration.worker_id.clone(); let static_models = declaration.models.clone(); let id = declaration.id.clone(); + + // The guard spans registry upsert, schema recomposition, AND the + // static catalog write below, so register and unregister can + // never interleave their registry/catalog/schema mutations — + // an unregister racing this handler cannot end up with a + // resurrected catalog slice for a removed record. + let _guard = entry_lock.lock().await; + let upserted = registry .upsert(declaration, worker_id, input.token) .await @@ -83,10 +94,8 @@ pub fn make_provider_register( let token = upserted.token; let availability_recovered = upserted.availability_recovered; - // Re-compose the entry schema from every registered declaration — - // under the entry write lock so concurrent boots compose. + // Re-compose the entry schema from every registered declaration. { - let _guard = entry_lock.lock().await; let mut provider_schemas = BTreeMap::new(); for rec in registry.list().await { let schema = provider_entry_schema( @@ -141,3 +150,89 @@ pub fn make_provider_register( }) } } + +/// The `router::provider::unregister` iii function — operator escape hatch +/// for a token lock-out: without it, a provider whose state diverged from the +/// router's persisted registry ("bound to another worker") could never come +/// back. Drops the record and its catalog slice, re-composes the entry schema, +/// and emits provider/model change events so open consoles converge. +pub fn make_provider_unregister( + iii: IIIClient, + registry: Arc, + catalog: Arc, + entry_lock: EntryWriteLock, + events: Arc, +) -> impl Fn( + ProviderUnregisterRequest, +) -> BoxFuture<'static, Result> + + Send + + Sync + + 'static { + move |input: ProviderUnregisterRequest| { + let (iii, registry, catalog, entry_lock, events) = ( + iii.clone(), + registry.clone(), + catalog.clone(), + entry_lock.clone(), + events.clone(), + ); + Box::pin(async move { + let id = input.id; + if !valid_id(&id) { + return Err(RouterError::new( + RouterCode::InvalidRequest, + format!("invalid provider id: {id}"), + ) + .into()); + } + // One guard across registry removal, catalog pruning, and schema + // recomposition so a concurrent register cannot interleave and + // resurrect state mid-teardown. + let _guard = entry_lock.lock().await; + + let removed = registry.remove(&id).await.map_err(|e| { + Error::from(RouterError::new( + RouterCode::InvalidRequest, + format!("registry persist failed: {e}"), + )) + })?; + // Cleanup runs regardless of `removed`: an earlier unregister + // that failed after the registry write left the slice/schema + // behind, and the retry (removed=false) must still converge. + let slice_removed = catalog.remove_slice(&id).await?; + + if removed || slice_removed { + let mut provider_schemas = BTreeMap::new(); + for rec in registry.list().await { + let schema = provider_entry_schema( + rec.declaration.config_schema.as_ref(), + &serde_json::to_value(rec.declaration.defaults.clone()) + .unwrap_or(Value::Null), + rec.declaration.system_prompt.as_deref(), + ); + provider_schemas.insert(rec.declaration.id.clone(), schema); + } + register_entry(&iii, &provider_schemas).await?; + } + + if removed { + events + .emit( + triggers::PROVIDER_CHANGED, + json!({ "provider": id, "op": "unregister" }), + ) + .await; + } + if slice_removed { + events + .emit( + triggers::MODELS_CHANGED, + json!({ "provider": id, "count": 0 }), + ) + .await; + } + + Ok(ProviderUnregisterResponse { ok: true, removed }) + }) + } +} diff --git a/llm-router/src/registry/store.rs b/llm-router/src/registry/store.rs index 0f38b4322..79247369b 100644 --- a/llm-router/src/registry/store.rs +++ b/llm-router/src/registry/store.rs @@ -197,6 +197,23 @@ impl RegistryStore { } } + /// Operator escape hatch (`router::provider::unregister`): drop a record + /// so a provider that lost its registration token can register fresh. + /// Returns whether a record existed. Persist-then-swap: the live map only + /// changes after the durable write succeeds, so a failed persist leaves + /// memory and state in agreement and the retry starts clean. + pub async fn remove(&self, id: &str) -> Result { + let mut records = self.records.lock().await; // serialized writer + if !records.contains_key(id) { + return Ok(false); + } + let mut next = records.clone(); + next.remove(id); + self.persist(&next).await?; + *records = next; + Ok(true) + } + /// Returns true when the flag actually changed (callers emit on change only). pub async fn set_availability(&self, id: &str, available: bool) -> bool { let mut records = self.records.lock().await; diff --git a/llm-router/src/routing.rs b/llm-router/src/routing.rs index 6db153e31..e2f6e0529 100644 --- a/llm-router/src/routing.rs +++ b/llm-router/src/routing.rs @@ -46,11 +46,15 @@ pub fn decide(input: &DecideInput) -> Result, RouterError> { } // 2. Unique catalog owner; 2+ owners → ambiguous (the router never guesses). + // Catalog slices outlive registrations (they persist across provider + // departures), so an owner counts only while it is registered — otherwise + // a stale slice routes to a provider that `router::chat` will reject. let mut owners: Vec<&str> = input .catalog .iter() .filter(|(_, ids)| ids.iter().any(|m| m == &input.model)) .map(|(p, _)| p.as_str()) + .filter(|p| registered(p)) .collect(); owners.sort_unstable(); match owners.len() { @@ -185,6 +189,21 @@ mod tests { ); } + #[test] + fn step2_stale_catalog_owner_is_skipped_when_unregistered() { + let mut input = base(); + input.model = "local-llama".into(); + // lmstudio's slice persisted but the provider is gone: never route to it. + input.registered_providers = vec!["anthropic".into(), "openai".into()]; + assert_eq!( + decide(&input).unwrap_err().code, + RouterCode::NoProviderForModel + ); + // A shared model with one surviving owner routes to that owner. + input.model = "shared-model".into(); + assert_eq!(decide(&input).unwrap(), vec!["openai"]); + } + #[test] fn step3_heuristics_first_match_registered_only_invalid_regex_skipped() { let mut input = base(); diff --git a/llm-router/src/surface.rs b/llm-router/src/surface.rs index f539fa878..1c254df61 100644 --- a/llm-router/src/surface.rs +++ b/llm-router/src/surface.rs @@ -17,9 +17,9 @@ use crate::types::router::{ ModelGetResponse, ModelsListRequest, ModelsListResponse, ModelsReconcileRequest, ModelsReconcileResponse, ModelsSupportsRequest, ModelsSupportsResponse, ProviderListRequest, ProviderListResponse, ProviderRegisterRequest, ProviderRegisterResponse, - ProviderResolveRequest, ProviderResolveResponse, RouteRequest, RouteResponse, RouterAck, - SystemPromptGetRequest, SystemPromptGetResponse, UpdateCredentialRequest, - UpdateCredentialResponse, + ProviderResolveRequest, ProviderResolveResponse, ProviderUnregisterRequest, + ProviderUnregisterResponse, RouteRequest, RouteResponse, RouterAck, SystemPromptGetRequest, + SystemPromptGetResponse, UpdateCredentialRequest, UpdateCredentialResponse, }; // ── function id + description constants — consumed by both register_router and @@ -83,6 +83,11 @@ pub const PROVIDER_REGISTER_ID: &str = "router::provider::register"; pub const PROVIDER_REGISTER_DESC: &str = "Provider self-declaration at attach time (token-gated \ upsert); composes the configuration entry schema and reconciles static models."; +pub const PROVIDER_UNREGISTER_ID: &str = "router::provider::unregister"; +pub const PROVIDER_UNREGISTER_DESC: &str = + "Operator escape hatch: drop a provider's registration record and catalog \ + slice so a provider that lost its registration token can register fresh."; + pub const PROVIDER_RESOLVE_ID: &str = "router::provider::resolve"; pub const PROVIDER_RESOLVE_DESC: &str = "Resolve a provider's effective credential + api_url + max_tokens (token-gated)."; @@ -168,6 +173,10 @@ pub fn catalog() -> Vec { PROVIDER_REGISTER_ID, PROVIDER_REGISTER_DESC, ), + spec::( + PROVIDER_UNREGISTER_ID, + PROVIDER_UNREGISTER_DESC, + ), spec::( PROVIDER_RESOLVE_ID, PROVIDER_RESOLVE_DESC, diff --git a/llm-router/src/types/router.rs b/llm-router/src/types/router.rs index 7d52f65b3..ddb685f8b 100644 --- a/llm-router/src/types/router.rs +++ b/llm-router/src/types/router.rs @@ -399,6 +399,23 @@ pub struct ProviderRegisterRequest { pub token: Option, } +/// Input of `router::provider::unregister` — the operator escape hatch for a +/// provider whose registration token was lost (state wiped on one side): drop +/// the record and its catalog slice so the provider can register fresh. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct ProviderUnregisterRequest { + /// Provider id to unbind. + pub id: String, +} + +/// Output of `router::provider::unregister`. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct ProviderUnregisterResponse { + pub ok: bool, + /// Whether a record existed and was removed. + pub removed: bool, +} + /// Input of `router::provider::resolve`. #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct ProviderResolveRequest { diff --git a/llm-router/src/ui.rs b/llm-router/src/ui.rs index 87c301135..39bf0f1c4 100644 --- a/llm-router/src/ui.rs +++ b/llm-router/src/ui.rs @@ -36,9 +36,11 @@ fn console_ui() -> ConsoleUi { .style(STYLES_PATH, STYLES_CSS) } -/// Register the llm-router console UI. Call after the router's functions -/// are registered. Takes the bare client (what `register_worker` hands -/// `main`); the shared crate wants an `Arc` for its spawned tasks. +/// Register the llm-router console UI. Independent of the router's function +/// surface — call it FIRST at boot so an open console swaps in the custom +/// config form before the router's slower engine round-trips finish. Takes +/// the bare client (what `register_worker` hands `main`); the shared crate +/// wants an `Arc` for its spawned tasks. pub fn register(iii: &IIIClient) { console_ui().register(&Arc::new(iii.clone())); } diff --git a/llm-router/tests/golden/schemas/router.provider.unregister.json b/llm-router/tests/golden/schemas/router.provider.unregister.json new file mode 100644 index 000000000..4961475e4 --- /dev/null +++ b/llm-router/tests/golden/schemas/router.provider.unregister.json @@ -0,0 +1,38 @@ +{ + "description": "Operator escape hatch: drop a provider's registration record and catalog slice so a provider that lost its registration token can register fresh.", + "function_id": "router::provider::unregister", + "request_schema": { + "$schema": "http://json-schema.org/draft-07/schema#", + "description": "Input of `router::provider::unregister` — the operator escape hatch for a provider whose registration token was lost (state wiped on one side): drop the record and its catalog slice so the provider can register fresh.", + "properties": { + "id": { + "description": "Provider id to unbind.", + "type": "string" + } + }, + "required": [ + "id" + ], + "title": "ProviderUnregisterRequest", + "type": "object" + }, + "response_schema": { + "$schema": "http://json-schema.org/draft-07/schema#", + "description": "Output of `router::provider::unregister`.", + "properties": { + "ok": { + "type": "boolean" + }, + "removed": { + "description": "Whether a record existed and was removed.", + "type": "boolean" + } + }, + "required": [ + "ok", + "removed" + ], + "title": "ProviderUnregisterResponse", + "type": "object" + } +} diff --git a/llm-router/tests/integration.rs b/llm-router/tests/integration.rs index ed8a010b1..6ff5d83f5 100644 --- a/llm-router/tests/integration.rs +++ b/llm-router/tests/integration.rs @@ -767,6 +767,88 @@ async fn registration_token_gates_takeover_resolve_and_reconcile() { router_iii.shutdown(); } +/// The operator escape hatch for a token lock-out: a provider that lost its +/// registration token is rejected forever — until `provider::unregister` +/// drops the record (and its catalog slice), after which a fresh register +/// binds a new token. +#[tokio::test(flavor = "multi_thread")] +async fn unregister_frees_a_token_locked_provider_and_prunes_its_catalog() { + let engine = engine_or_skip!(); + + let router_iii = register_worker(&engine.url, InitOptions::default()); + register_router(router_iii.clone()) + .await + .expect("router boots"); + let provider = start_live_provider(&engine.url, ProviderOptions::default()).await; + + // Lost token: re-register without it is the lock-out. + let err = call( + &provider.iii, + "router::provider::register", + json!({ "id": "real" }), + ) + .await + .unwrap_err(); + assert_eq!(remote_code(&err), "router/registration_rejected"); + + // Unknown id: ok, nothing removed. + let res = call( + &provider.iii, + "router::provider::unregister", + json!({ "id": "missing" }), + ) + .await + .expect("unregister answers"); + assert_eq!(res["removed"], json!(false)); + + // The escape hatch: drop the record. + let res = call( + &provider.iii, + "router::provider::unregister", + json!({ "id": "real" }), + ) + .await + .expect("unregister succeeds"); + assert_eq!(res["removed"], json!(true)); + + // Gone from the provider list, and the static catalog slice is pruned. + let list = call(&provider.iii, "router::provider::list", json!({})) + .await + .expect("list"); + assert!( + !list["providers"] + .as_array() + .unwrap() + .iter() + .any(|p| p["id"] == "real"), + "unregistered provider must leave the list" + ); + let model = call( + &provider.iii, + "router::models::get", + json!({ "provider": "real", "id": "live-1" }), + ) + .await + .expect("models::get answers"); + assert!( + model.is_null(), + "catalog slice must be pruned on unregister, got {model}" + ); + + // Fresh register without a token now binds cleanly with a NEW token. + let res = call( + &provider.iii, + "router::provider::register", + json!({ "id": "real" }), + ) + .await + .expect("fresh register succeeds after unregister"); + let new_token = res["registration_token"].as_str().expect("new token"); + assert_ne!(new_token, provider.token, "a new token must be minted"); + + router_iii.shutdown(); +} + #[tokio::test(flavor = "multi_thread")] async fn resolve_precedence_config_over_env_over_none() { let engine = engine_or_skip!(); diff --git a/llm-router/tests/schemas.rs b/llm-router/tests/schemas.rs index bfcc5e34d..762346e98 100644 --- a/llm-router/tests/schemas.rs +++ b/llm-router/tests/schemas.rs @@ -32,7 +32,7 @@ fn spec_to_pretty_json(spec: &FunctionSpec) -> String { pretty } -/// The catalog must cover exactly the 18 registered functions, in registration +/// The catalog must cover exactly the 19 registered functions, in registration /// order (kept in lockstep with `register::register_router`). #[test] fn catalog_lists_all_functions_in_registration_order() { @@ -53,6 +53,7 @@ fn catalog_lists_all_functions_in_registration_order() { "router::system_prompt::get", "router::route", "router::provider::register", + "router::provider::unregister", "router::provider::resolve", "router::provider::update_credential", "router::models::reconcile", diff --git a/provider-anthropic/Cargo.lock b/provider-anthropic/Cargo.lock index 30b8d6ae2..e850ee3a3 100644 --- a/provider-anthropic/Cargo.lock +++ b/provider-anthropic/Cargo.lock @@ -996,7 +996,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.4.2" +version = "1.4.7" dependencies = [ "async-trait", "clap",