From 71bd1bf1bbf2f3c823123e8d3951e22a39b4410a Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Fri, 14 Aug 2026 23:35:12 -0300 Subject: [PATCH 1/5] (MOT-4412) fix(llm-router): harden provider lifecycle --- .github/scripts/discover_changed_workers.py | 22 + .../tests/test_discover_changed_workers.py | 48 + .../scripts/tests/test_rust_ci_workflows.py | 6 +- .github/workflows/ci.yml | 55 ++ .../provider-integration-testkit/Cargo.lock | 2 +- llm-router/Cargo.lock | 2 +- llm-router/Cargo.toml | 2 +- llm-router/src/catalog/store.rs | 159 +++- llm-router/src/chat/abort.rs | 2 +- llm-router/src/chat/chat.rs | 540 ++++++++++- llm-router/src/chat/inflight.rs | 124 ++- llm-router/src/count_tokens.rs | 11 +- llm-router/src/registry/register.rs | 153 ++- llm-router/src/registry/store.rs | 297 +++++- llm-router/src/routing.rs | 133 ++- llm-router/src/types/router.rs | 5 +- .../schemas/router.provider.register.json | 2 +- llm-router/tests/integration.rs | 872 +++++++++++++++++- 18 files changed, 2268 insertions(+), 167 deletions(-) diff --git a/.github/scripts/discover_changed_workers.py b/.github/scripts/discover_changed_workers.py index 2d7a93654..b5dbe90c8 100644 --- a/.github/scripts/discover_changed_workers.py +++ b/.github/scripts/discover_changed_workers.py @@ -6,6 +6,7 @@ source_changed : workers whose change wasn't only metadata rust / node / python : language buckets (subset of changed_workers) integration_changed : bool, did an integration-stack input change + llm_router_integration : bool, must the live llm-router suite run provider_contract : providers whose hermetic contract must run crates : shared crates/ dirs with source changes any : bool, any worker or crate change @@ -65,6 +66,15 @@ "harness/tests/quickstart/", ) +# The llm-router owns a separate real-engine lifecycle suite. Keep this gate +# independent from Harness Integration: that stack intentionally substitutes a +# ScriptedRouter and therefore cannot validate router transport/lifecycle bugs. +LLM_ROUTER_INTEGRATION_WORKERS = {"llm-router"} +LLM_ROUTER_INTEGRATION_INFRA_PATHS = { + ".github/scripts/discover_changed_workers.py", + ".github/workflows/ci.yml", +} + # Hermetic provider contracts run the real engine, llm-router, and selected # provider against a loopback HTTP/SSE upstream. Direct provider changes stay # narrow; shared router/testkit/CI changes fan out to every supported provider. @@ -275,6 +285,13 @@ def main(argv: list[str] | None = None) -> int: INTEGRATION_INFRA_PATHS, INTEGRATION_EXCLUDED_PREFIXES, ) + llm_router_integration = suite_changed( + files, + forced, + LLM_ROUTER_INTEGRATION_WORKERS, + LLM_ROUTER_INTEGRATION_INFRA_PATHS, + (), + ) provider_contract = provider_contract_selection(files, workers) by_language: dict[str, list[str]] = {"rust": [], "node": [], "python": []} for w in changed: @@ -289,6 +306,7 @@ def main(argv: list[str] | None = None) -> int: "source_changed": source_changed, "by_language": by_language, "integration_changed": integration_changed, + "llm_router_integration": llm_router_integration, "provider_contract": provider_contract, "crates": changed_crates, } @@ -305,6 +323,10 @@ def main(argv: list[str] | None = None) -> int: f.write( f"integration_changed={'true' if integration_changed else 'false'}\n" ) + f.write( + "llm_router_integration=" + f"{'true' if llm_router_integration else 'false'}\n" + ) f.write(f"provider_contract={json.dumps(provider_contract)}\n") f.write(f"crates={json.dumps(changed_crates)}\n") f.write(f"any={'true' if any_change else 'false'}\n") diff --git a/.github/scripts/tests/test_discover_changed_workers.py b/.github/scripts/tests/test_discover_changed_workers.py index 824ac12fd..dfda5e350 100644 --- a/.github/scripts/tests/test_discover_changed_workers.py +++ b/.github/scripts/tests/test_discover_changed_workers.py @@ -283,8 +283,56 @@ def test_provider_change_stays_out_of_integration(self, tmp_path): data = json.loads(r.stdout) assert data["changed_workers"] == ["provider-anthropic"] assert data["integration_changed"] is False + assert data["llm_router_integration"] is False assert data["provider_contract"] == ["provider-anthropic"] + @pytest.mark.parametrize( + "changed_path", + [ + "llm-router/src/lib.rs", + "llm-router/tests/integration.rs", + ".github/workflows/ci.yml", + ".github/scripts/discover_changed_workers.py", + ], + ) + def test_llm_router_runtime_inputs_run_live_router_integration( + self, tmp_path, changed_path + ): + repo = make_repo_with_harness(tmp_path) + path = repo / changed_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("changed\n") + subprocess.run( + ["git", "add", "."], cwd=repo, check=True, env=GIT_HERMETIC_ENV + ) + subprocess.run( + ["git", "commit", "-q", "-m", "router integration input"], + cwd=repo, + check=True, + env=GIT_HERMETIC_ENV, + ) + r = run_script(repo, "main~1") + assert r.returncode == 0, r.stderr + assert json.loads(r.stdout)["llm_router_integration"] is True + + def test_llm_router_docs_do_not_run_live_router_integration(self, tmp_path): + repo = make_repo_with_harness(tmp_path) + path = repo / "llm-router" / "README.md" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("# docs\n") + subprocess.run( + ["git", "add", "."], cwd=repo, check=True, env=GIT_HERMETIC_ENV + ) + subprocess.run( + ["git", "commit", "-q", "-m", "router docs"], + cwd=repo, + check=True, + env=GIT_HERMETIC_ENV, + ) + r = run_script(repo, "main~1") + assert r.returncode == 0, r.stderr + assert json.loads(r.stdout)["llm_router_integration"] is False + def test_database_change_stays_out_of_integration(self, tmp_path): repo = make_repo_with_harness(tmp_path) (repo / "database" / "lib.rs").write_text("// change\n") diff --git a/.github/scripts/tests/test_rust_ci_workflows.py b/.github/scripts/tests/test_rust_ci_workflows.py index 3c6fdc9eb..c00025212 100644 --- a/.github/scripts/tests/test_rust_ci_workflows.py +++ b/.github/scripts/tests/test_rust_ci_workflows.py @@ -1,6 +1,7 @@ from __future__ import annotations from pathlib import Path +import re import tomllib import yaml @@ -26,8 +27,9 @@ def test_rust_toolchain_is_pinned_to_the_last_verified_stable() -> None: assert toolchain["toolchain"]["channel"] == "1.97.1" bodies = "\n".join(path.read_text() for path in WORKFLOWS.glob("*.yml")) - assert "dtolnay/rust-toolchain@stable" not in bodies - assert bodies.count("dtolnay/rust-toolchain@1.97.1") == 13 + workflow_toolchains = re.findall(r"dtolnay/rust-toolchain@([^\s]+)", bodies) + assert workflow_toolchains + assert set(workflow_toolchains) == {"1.97.1"} def test_prs_restore_rust_caches_and_main_pushes_publish_them() -> None: diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index af9ce3ae7..1902eb78c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -53,6 +53,7 @@ jobs: changed_workers: ${{ steps.bucket.outputs.changed_workers }} source_changed: ${{ steps.bucket.outputs.source_changed }} integration_changed: ${{ steps.bucket.outputs.integration_changed }} + llm_router_integration: ${{ steps.bucket.outputs.llm_router_integration }} provider_contract: ${{ steps.bucket.outputs.provider_contract }} any: ${{ steps.bucket.outputs.any }} steps: @@ -267,6 +268,60 @@ jobs: - name: Run tests run: cargo test --locked --all-features + # ────────────────────────────────────────────────────────────── + # llm-router lifecycle contract: unlike the regular Rust job, this always + # supplies the pinned engine and therefore cannot silently self-skip the + # real-bus registration, streaming, cancellation, and restart scenarios. + # ────────────────────────────────────────────────────────────── + llm-router-integration: + name: "llm-router: live engine integration" + needs: discover + if: needs.discover.outputs.llm_router_integration == 'true' + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - uses: actions/checkout@v5 + + - name: Rewrite SSH to HTTPS for public deps + run: git config --global url."https://github.com/".insteadOf "ssh://git@github.com/" + + - uses: dtolnay/rust-toolchain@1.97.1 + + - uses: Swatinem/rust-cache@v2 + with: + shared-key: llm-router-integration + save-if: false + workspaces: llm-router -> target + + - name: Install pinned iii engine + env: + VERSION: '0.22.1' + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -euo pipefail + curl -fsSL https://install.iii.dev/iii/main/install.sh -o /tmp/install-iii.sh + sh /tmp/install-iii.sh + { + echo "$HOME/.local/bin" + echo "$HOME/.iii/bin" + } >> "$GITHUB_PATH" + export PATH="$HOME/.local/bin:$HOME/.iii/bin:$PATH" + engine_bin=$(command -v iii) + [[ -x "$engine_bin" ]] || { echo "::error::iii engine is not executable"; exit 3; } + echo "III_ENGINE_BIN=$engine_bin" >> "$GITHUB_ENV" + iii --version + + - name: Run real-engine router lifecycle suite + run: | + set -euo pipefail + [[ -x "$III_ENGINE_BIN" ]] || { echo "::error::III_ENGINE_BIN is unavailable"; exit 3; } + cargo test \ + --locked \ + --manifest-path llm-router/Cargo.toml \ + --no-default-features \ + --test integration \ + -- --nocapture --test-threads=1 + # ────────────────────────────────────────────────────────────── # Provider contracts: real engine + real router + selected provider, # with the vendor HTTP/SSE boundary replaced by a loopback stub. No live diff --git a/crates/provider-integration-testkit/Cargo.lock b/crates/provider-integration-testkit/Cargo.lock index 03d021656..9c547ca45 100644 --- a/crates/provider-integration-testkit/Cargo.lock +++ b/crates/provider-integration-testkit/Cargo.lock @@ -1036,7 +1036,7 @@ checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae" [[package]] name = "llm-router" -version = "1.4.7" +version = "1.4.8" dependencies = [ "async-trait", "clap", diff --git a/llm-router/Cargo.lock b/llm-router/Cargo.lock index 8ba4a06e9..5aec7d8c9 100644 --- a/llm-router/Cargo.lock +++ b/llm-router/Cargo.lock @@ -1008,7 +1008,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.4.7" +version = "1.4.8" dependencies = [ "async-trait", "clap", diff --git a/llm-router/Cargo.toml b/llm-router/Cargo.toml index bee3491a3..287b9b5d3 100644 --- a/llm-router/Cargo.toml +++ b/llm-router/Cargo.toml @@ -4,7 +4,7 @@ # this lineage starts above them so Create Tag never collides. [package] name = "llm-router" -version = "1.4.7" +version = "1.4.8" edition = "2021" publish = false license = "Apache-2.0" diff --git a/llm-router/src/catalog/store.rs b/llm-router/src/catalog/store.rs index 7bd1cf491..e81882c96 100644 --- a/llm-router/src/catalog/store.rs +++ b/llm-router/src/catalog/store.rs @@ -1,6 +1,8 @@ //! Durable model catalog: one slice per provider, replaced wholesale on -//! reconcile. Same serialized-writer pattern as the registry — persisted via -//! the engine's `state::get`/`state::set` iii functions (src/state.rs). +//! reconcile. Same serialized-writer pattern as the registry: prepare a +//! snapshot, persist it, then publish it in memory while holding the writer +//! lock. Persistence uses the engine's `state::get`/`state::set` iii functions +//! (src/state.rs). //! //! Engine-backed coverage: tests/integration.rs (reconcile, restart restore). use std::collections::HashMap; @@ -9,13 +11,44 @@ use crate::state::{state_get, state_set}; use crate::types::errors::is_function_not_found; use crate::types::model::Model; use iii_sdk::{errors::Error, IIIClient}; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, MutexGuard}; const CATALOG_KEY: &str = "catalog"; pub struct CatalogStore { iii: IIIClient, slices: Mutex>>, + #[cfg(test)] + persist_result: Option>, +} + +/// A catalog snapshot that is durable but intentionally not visible in +/// memory yet. Provider registration commits it only after the matching +/// registry record is durable; dropping it without `commit` keeps readers on +/// the previous snapshot. Because iii state has no multi-key transaction, a +/// process crash between this durable write and registry commit can still +/// leave a catalog-only snapshot for restart recovery to reconcile. +#[must_use = "a prepared catalog write must be committed or rolled back"] +pub struct PreparedSlice<'a> { + store: &'a CatalogStore, + slices: MutexGuard<'a, HashMap>>, + previous: HashMap>, + next: HashMap>, +} + +impl PreparedSlice<'_> { + pub fn commit(mut self) { + *self.slices = self.next; + } + + /// Restore the durable snapshot while leaving the already-visible memory + /// snapshot untouched. If this fails, routing still excludes the staged + /// catalog owner because no registry record was published, but a restart + /// can reload the orphaned catalog slice until a later reconcile repairs + /// state. + pub async fn rollback(self) -> Result<(), Error> { + self.store.persist(&self.previous).await + } } impl CatalogStore { @@ -23,6 +56,8 @@ impl CatalogStore { Self { iii, slices: Mutex::new(HashMap::new()), + #[cfg(test)] + persist_result: None, } } @@ -76,10 +111,120 @@ impl CatalogStore { .cloned() } - pub async fn set_slice(&self, provider: &str, models: Vec) -> Result<(), Error> { - let mut slices = self.slices.lock().await; // serialized writer - slices.insert(provider.to_string(), models); - let value = serde_json::to_value(&*slices).unwrap_or_default(); + async fn persist(&self, slices: &HashMap>) -> Result<(), Error> { + #[cfg(test)] + if let Some(result) = &self.persist_result { + return result.clone(); + } + let value = serde_json::to_value(slices).unwrap_or_default(); state_set(&self.iii, CATALOG_KEY, value).await } + + /// Persist a candidate slice while blocking catalog readers and writers, + /// but defer its in-memory publication until `PreparedSlice::commit`. + pub async fn prepare_slice( + &self, + provider: &str, + models: Vec, + ) -> Result, Error> { + let slices = self.slices.lock().await; // serialized writer + let previous = slices.clone(); + let mut next = previous.clone(); + next.insert(provider.to_string(), models); + self.persist(&next).await?; + Ok(PreparedSlice { + store: self, + slices, + previous, + next, + }) + } + + pub async fn set_slice(&self, provider: &str, models: Vec) -> Result<(), Error> { + self.prepare_slice(provider, models).await?.commit(); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn model(id: &str) -> Model { + Model { + id: id.into(), + provider: "anthropic".into(), + display_name: None, + context_window: 200_000, + max_output_tokens: 64_000, + input_limit: None, + supports_thinking: None, + supports_xhigh: None, + reasoning_efforts: None, + supports_tools: None, + supports_vision: None, + supports_cache: None, + supports_structured_output: None, + thinking_budgets: None, + pricing: None, + } + } + + fn store_with_persistence(result: Result<(), Error>) -> CatalogStore { + CatalogStore { + iii: IIIClient::new("ws://unused.invalid"), + slices: Mutex::new(HashMap::new()), + persist_result: Some(result), + } + } + + #[tokio::test] + async fn prepared_slice_is_invisible_until_commit() { + let store = store_with_persistence(Ok(())); + store + .slices + .lock() + .await + .insert("anthropic".into(), vec![model("old")]); + + let prepared = store + .prepare_slice("anthropic", vec![model("new")]) + .await + .expect("candidate persists"); + assert_eq!( + prepared.slices["anthropic"][0].id, "old", + "persistence alone must not publish the candidate" + ); + prepared.commit(); + assert_eq!(store.slice("anthropic").await, vec![model("new")]); + } + + #[tokio::test(start_paused = true)] + async fn failed_persist_keeps_previous_slice() { + // An unconnected client makes state::set time out. Paused Tokio time + // skips the SDK timeout without requiring a live engine. + let store = CatalogStore::new(IIIClient::new("ws://unconnected.invalid")); + store + .slices + .lock() + .await + .insert("anthropic".into(), vec![model("old")]); + + let err = store + .set_slice("anthropic", vec![model("new")]) + .await + .expect_err("persistence must fail"); + assert!(matches!(err, Error::Timeout)); + assert_eq!(store.slice("anthropic").await, vec![model("old")]); + + let err = store + .set_slice("openai", vec![model("gpt")]) + .await + .expect_err("persistence must fail"); + assert!(matches!(err, Error::Timeout)); + assert!( + store.slice("openai").await.is_empty(), + "a failed write must not publish a new provider slice" + ); + } } diff --git a/llm-router/src/chat/abort.rs b/llm-router/src/chat/abort.rs index 888a79bbc..e14444da0 100644 --- a/llm-router/src/chat/abort.rs +++ b/llm-router/src/chat/abort.rs @@ -36,7 +36,7 @@ mod tests { #[tokio::test] async fn aborts_known_requests_once_and_ignores_unknown() { let inflight = Arc::new(InflightMap::default()); - inflight.insert("r1"); + let _reservation = inflight.reserve("r1").unwrap(); let abort = make_abort(inflight); // First abort of a known request succeeds, the second is a no-op // (idempotent), and an unknown id reports `aborted: false`. diff --git a/llm-router/src/chat/chat.rs b/llm-router/src/chat/chat.rs index 371fccbc8..5b73a2b5d 100644 --- a/llm-router/src/chat/chat.rs +++ b/llm-router/src/chat/chat.rs @@ -105,6 +105,171 @@ fn is_router_coded(err: &Error) -> bool { matches!(err, Error::Remote { code, .. } if code.starts_with("router/")) } +type ProviderCallOutcome = Result; +const MAX_PROVIDER_CALLS: usize = 64; + +fn provider_call_slots() -> Arc { + static SLOTS: std::sync::OnceLock> = std::sync::OnceLock::new(); + Arc::clone(SLOTS.get_or_init(|| Arc::new(tokio::sync::Semaphore::new(MAX_PROVIDER_CALLS)))) +} + +enum ProviderCallAdmission { + Admitted(tokio::sync::OwnedSemaphorePermit), + Aborted, + Saturated, +} + +fn try_acquire_provider_call_slot( + slots: Arc, + aborted: &std::sync::atomic::AtomicBool, +) -> ProviderCallAdmission { + if aborted.load(Ordering::SeqCst) { + return ProviderCallAdmission::Aborted; + } + + match slots.try_acquire_owned() { + Ok(permit) => { + if aborted.load(Ordering::SeqCst) { + ProviderCallAdmission::Aborted + } else { + ProviderCallAdmission::Admitted(permit) + } + } + Err(tokio::sync::TryAcquireError::NoPermits) if aborted.load(Ordering::SeqCst) => { + ProviderCallAdmission::Aborted + } + Err(tokio::sync::TryAcquireError::NoPermits) => ProviderCallAdmission::Saturated, + Err(tokio::sync::TryAcquireError::Closed) => { + panic!("provider call semaphore is never closed") + } + } +} + +fn pre_stream_error_kind(code: RouterCode) -> ErrorKind { + match code { + RouterCode::ProviderUnavailable => ErrorKind::Transient, + _ => ErrorKind::Permanent, + } +} + +fn take_provider_outcome( + outcome_rx: &mut tokio::sync::oneshot::Receiver, +) -> Option { + outcome_rx.try_recv().ok() +} + +/// Release foreground ownership of a provider RPC once the relay has reached a +/// definitive result. +/// +/// The RPC outcome is published before its error path closes the relay reader, +/// so a pre-stream dispatch failure remains available for classification. If +/// the RPC is still running, a reaper awaits it in the background instead of +/// aborting its `IIIClient::trigger` future: that future owns the SDK timeout +/// which removes the invocation from the SDK pending map. Thus Done, Error, +/// Idle, and Abort return immediately while cleanup remains bounded by the +/// `timeout_ms` passed to the trigger. The task itself owns an admission permit +/// acquired before the trigger is spawned, so the same process-wide bound +/// covers both foreground RPCs and background reapers. +/// +/// This cannot repair the SDK's separate synchronous `send_message` failure +/// path, where `trigger` returns after inserting a pending invocation. That +/// cleanup belongs in `iii-sdk` itself. +async fn finish_provider_call( + call_task: tokio::task::JoinHandle<()>, + mut outcome_rx: tokio::sync::oneshot::Receiver, +) -> Option { + let published = take_provider_outcome(&mut outcome_rx); + if call_task.is_finished() { + let joined = call_task.await; + return published + .or_else(|| take_provider_outcome(&mut outcome_rx)) + .or_else(|| { + joined + .err() + .filter(|err| err.is_panic()) + .map(|_| Err(Error::Handler("provider task panicked".into()))) + }); + } + + tokio::spawn(async move { + if call_task.await.as_ref().is_err_and(|err| err.is_panic()) { + eprintln!("[llm-router] provider stream task panicked after relay completion"); + } + }); + published +} + +#[allow(clippy::too_many_arguments)] +fn send_error_terminal( + sink: &dyn FrameSink, + partial: Option<&crate::types::messages::AssistantMessage>, + model: &str, + provider: &str, + message: &str, + error_kind: ErrorKind, + usage: Option, +) -> AssistantMessageEvent { + let frame = synthesize_error( + partial, + model, + provider, + message, + error_kind, + usage, + now_ms(), + ); + let _ = sink.send(&serde_json::to_string(&frame).expect("serializable frame")); + frame +} + +#[allow(clippy::too_many_arguments)] +fn fail_with_terminal( + sink: &dyn FrameSink, + partial: Option<&crate::types::messages::AssistantMessage>, + model: &str, + provider: &str, + code: RouterCode, + message: String, + error_kind: ErrorKind, + usage: Option, +) -> Error { + send_error_terminal(sink, partial, model, provider, &message, error_kind, usage); + RouterError::new(code, message).into() +} + +fn provider_call_saturated_response( + sink: &dyn FrameSink, + partial: Option<&crate::types::messages::AssistantMessage>, + model: &str, + provider: &str, + usage: Option, +) -> ChatResponse { + let message = format!("provider {provider} call capacity is saturated; retry later"); + let terminal = send_error_terminal( + sink, + partial, + model, + provider, + &message, + ErrorKind::Transient, + usage, + ); + let AssistantMessageEvent::Error { error } = terminal else { + unreachable!() + }; + ChatResponse { + ok: false, + provider: provider.to_string(), + model: model.to_string(), + stop_reason: Some(StopReason::Error), + usage: error.usage, + error: Some(ErrorShape { + code: "transient".into(), + message, + }), + } +} + impl ChatPipeline { pub async fn run( &self, @@ -116,20 +281,18 @@ impl ChatPipeline { // reader budget and `router::chat` consumers never see a terminal. // (Regression: pre_stream_routing_failure_emits_one_error_terminal_frame.) let fail_pre_stream = |provider: &str, code: RouterCode, message: String| -> Error { - // Pre-stream failures are permanent: a bad model or an unrouted - // request won't succeed on retry. Mark the frame Permanent so a - // streaming consumer inspecting error_kind doesn't retry it. - let frame = synthesize_error( + // Most routing/validation decisions are permanent. Availability is + // topology state and may recover without changing the request. + fail_with_terminal( + sink.as_ref(), None, &call.model, provider, - &message, - ErrorKind::Permanent, + code, + message, + pre_stream_error_kind(code), None, - now_ms(), - ); - let _ = sink.send(&serde_json::to_string(&frame).expect("serializable frame")); - RouterError::new(code, message).into() + ) }; // ── validate (pre-stream typed throws) ── @@ -150,17 +313,29 @@ impl ChatPipeline { let config = snapshot(&self.config); let settings = config.settings().clone(); + let provider_records = self.registry.list().await; let candidates = decide(&DecideInput { model: call.model.clone(), provider: call.provider.clone(), - registered_providers: self.registry.ids().await, + registered_providers: provider_records + .iter() + .map(|record| record.declaration.id.clone()) + .collect(), + available_providers: provider_records + .iter() + .filter(|record| record.available) + .map(|record| record.declaration.id.clone()) + .collect(), catalog: self.catalog.model_ids().await, heuristics: settings.routing_heuristics.clone(), default_provider: settings.default_provider.clone(), }) .map_err(|e| fail_pre_stream("", e.code, e.message))?; let provider = candidates[0].clone(); // MVP consumes candidates[0] - let record = match self.registry.get(&provider).await { + let record = match provider_records + .into_iter() + .find(|record| record.declaration.id == provider) + { Some(record) => record, None => { return Err(fail_pre_stream( @@ -201,16 +376,24 @@ impl ChatPipeline { .unwrap_or(8192), settings.output_token_max, ); + let provider_generation = record.generation; let request_id = call .request_id .clone() .unwrap_or_else(|| Uuid::new_v4().to_string()); - let entry_handle = self.inflight.insert(&request_id); + let Some(entry_handle) = self.inflight.reserve(&request_id) else { + return Err(fail_pre_stream( + &provider, + RouterCode::InvalidRequest, + format!("request_id {request_id} is already in flight"), + )); + }; let result = self .attempts( &call, &provider, + provider_generation, model_meta.as_ref(), max_output_tokens, &settings, @@ -219,7 +402,6 @@ impl ChatPipeline { sink, ) .await; - self.inflight.remove(&request_id); // Stamp token usage + cost onto the active `execute router::chat` // span. Every terminal outcome of the attempt loop funnels into the // ChatResponse, which carries the cost-filled usage. @@ -240,6 +422,7 @@ impl ChatPipeline { &self, call: &ChatCall, provider: &str, + provider_generation: u64, model_meta: Option<&crate::types::model::Model>, max_output_tokens: u64, settings: &RouterSettings, @@ -271,8 +454,46 @@ impl ChatPipeline { if inflight.aborted.load(Ordering::SeqCst) { break; } - let channel = create_router_channel(&self.iii).await?; + // Admission is synchronous and happens before even allocating the + // provider channel. Saturation cannot grow a hidden waiter queue or + // create another SDK invocation/in-flight stream. + let call_slot = match try_acquire_provider_call_slot( + provider_call_slots(), + inflight.aborted.as_ref(), + ) { + ProviderCallAdmission::Admitted(permit) => permit, + ProviderCallAdmission::Aborted => break, + ProviderCallAdmission::Saturated => { + return Ok(provider_call_saturated_response( + sink.as_ref(), + last_partial.as_ref(), + &call.model, + provider, + last_usage, + )); + } + }; + let channel = match create_router_channel(&self.iii).await { + Ok(channel) => channel, + Err(error) => { + let message = format!("router channel creation failed: {error}"); + send_error_terminal( + sink.as_ref(), + last_partial.as_ref(), + &call.model, + provider, + &message, + ErrorKind::Transient, + last_usage, + ); + return Err(error); + } + }; let mut reader = channel.reader; + if inflight.aborted.load(Ordering::SeqCst) { + drop(call_slot); + break; + } // router::abort closes the relay AND actively cancels the // provider's upstream via `provider::::abort` — without it the // provider only notices the closed channel on its next (ping) @@ -317,13 +538,22 @@ impl ChatPipeline { let function_id = format!("provider::{provider}::stream"); let closer = reader.closer(); let timeout = settings.stream_timeout_ms; + // Holding the admission permit inside the call task makes the cap + // a real bound on SDK pending stream invocations, including tasks + // transferred to the post-terminal reaper. + if inflight.aborted.load(Ordering::SeqCst) { + drop(call_slot); + break; + } // Carry the caller's OTel context into the spawned task so the // provider stream nests under `router::chat` instead of rooting a // detached trace. `with_context` (not `cx.attach()`) because the // `ContextGuard` is `!Send` and can't cross the `.await` below. let parent_cx = iii_helpers::observability::opentelemetry::Context::current(); + let (outcome_tx, outcome_rx) = tokio::sync::oneshot::channel(); let call_task = tokio::spawn( async move { + let _call_slot = call_slot; let out = iii .trigger(TriggerRequest { function_id, @@ -332,10 +562,14 @@ impl ChatPipeline { timeout_ms: Some(timeout), }) .await; - if out.is_err() { + let failed = out.is_err(); + // Publish before closing: once the reader wakes on EOF the + // orchestration can classify function_not_found without + // waiting for (or racing) the provider task. + let _ = outcome_tx.send(out); + if failed { closer(); } - out } .with_context(parent_cx), ); @@ -350,9 +584,12 @@ impl ChatPipeline { }, ) .await; - let call_outcome = call_task - .await - .unwrap_or(Err(Error::Handler("provider task panicked".into()))); + // Relay completion owns the foreground attempt lifecycle. Closing + // the reader propagates cancellation to a cooperative provider; + // the reaper lets a non-cooperative trigger run only until its SDK + // timeout without delaying Done/Error/Idle/Abort. + reader.close(); + let call_outcome = finish_provider_call(call_task, outcome_rx).await; match relay { RelayResult::Done { terminal, .. } => { @@ -363,7 +600,11 @@ impl ChatPipeline { // serving — heal a stale "down" flag (e.g. the boot-time // reset in `RegistryStore::load`, or a past transient // function_not_found) without waiting for a re-register. - if self.registry.set_availability(provider, true).await { + if self + .registry + .set_availability_if_current(provider, provider_generation, true) + .await + { self.events .emit( triggers::PROVIDER_CHANGED, @@ -443,13 +684,55 @@ impl ChatPipeline { } => { last_partial = partial; last_usage = usage.clone(); - if let Err(err) = call_outcome { + if let Some(Err(err)) = call_outcome { if !forwarded { if is_router_coded(&err) { + let (message, kind) = match &err { + Error::Remote { code, message, .. } + if code == RouterCode::ProviderUnavailable.as_str() => + { + (message.clone(), ErrorKind::Transient) + } + Error::Remote { message, .. } => { + (message.clone(), ErrorKind::Permanent) + } + _ => unreachable!("is_router_coded only matches remote errors"), + }; + send_error_terminal( + sink.as_ref(), + last_partial.as_ref(), + &call.model, + provider, + &message, + kind, + usage, + ); return Err(err); } if is_function_not_found(&err) { - if self.registry.set_availability(provider, false).await { + let message = format!("provider {provider} unavailable"); + let terminal_error = fail_with_terminal( + sink.as_ref(), + last_partial.as_ref(), + &call.model, + provider, + RouterCode::ProviderUnavailable, + message, + ErrorKind::Transient, + usage, + ); + // Emit before availability persistence/event + // fan-out: neither side effect may hold the + // consumer's terminal frame hostage. + if self + .registry + .set_availability_if_current( + provider, + provider_generation, + false, + ) + .await + { self.events .emit( triggers::PROVIDER_CHANGED, @@ -457,11 +740,7 @@ impl ChatPipeline { ) .await; } - return Err(RouterError::new( - RouterCode::ProviderUnavailable, - format!("provider {provider} unavailable"), - ) - .into()); + return Err(terminal_error); } } } @@ -600,6 +879,16 @@ fn insert_present(map: &mut serde_json::Map, key: &str, value: Op mod tests { use super::*; + struct DropSignal(Option>); + + impl Drop for DropSignal { + fn drop(&mut self) { + if let Some(tx) = self.0.take() { + let _ = tx.send(()); + } + } + } + fn remote(code: &str) -> Error { Error::Remote { code: code.into(), @@ -619,6 +908,201 @@ mod tests { assert!(!is_router_coded(&remote("function_not_found"))); } + #[test] + fn only_provider_unavailable_is_a_transient_pre_stream_decision() { + assert_eq!( + pre_stream_error_kind(RouterCode::ProviderUnavailable), + ErrorKind::Transient + ); + for code in [ + RouterCode::InvalidRequest, + RouterCode::UnknownProvider, + RouterCode::NoProviderForModel, + RouterCode::AmbiguousModel, + RouterCode::NotConfigured, + RouterCode::StructuredOutputUnsupported, + RouterCode::RegistrationRejected, + ] { + assert_eq!( + pre_stream_error_kind(code), + ErrorKind::Permanent, + "{} must remain permanent", + code.as_str() + ); + } + } + + #[tokio::test] + async fn provider_call_admission_slot_stays_held_while_the_rpc_is_reaped() { + let slots = Arc::new(tokio::sync::Semaphore::new(1)); + let aborted = std::sync::atomic::AtomicBool::new(false); + let ProviderCallAdmission::Admitted(call_slot) = + try_acquire_provider_call_slot(slots.clone(), &aborted) + else { + panic!("first provider call must be admitted"); + }; + let (outcome_tx, outcome_rx) = tokio::sync::oneshot::channel::(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let (dropped_tx, mut dropped_rx) = tokio::sync::oneshot::channel(); + let call_task = tokio::spawn(async move { + let _call_slot = call_slot; + let _drop_signal = DropSignal(Some(dropped_tx)); + let _outcome_tx = outcome_tx; + let _ = started_tx.send(()); + let _ = release_rx.await; + }); + started_rx.await.unwrap(); + + let outcome = tokio::time::timeout( + std::time::Duration::from_secs(1), + finish_provider_call(call_task, outcome_rx), + ) + .await + .expect("relay completion must not wait for the provider RPC timeout"); + + assert!(outcome.is_none()); + assert!( + slots.clone().try_acquire_owned().is_err(), + "a second SDK invocation must not be admitted while the reaper owns the first" + ); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(25), &mut dropped_rx) + .await + .is_err(), + "cleanup must not abort the SDK trigger future" + ); + + release_tx.send(()).unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), dropped_rx) + .await + .expect("the background reaper must observe bounded task completion") + .expect("drop signal sender must remain alive until task completion"); + let permit = tokio::time::timeout( + std::time::Duration::from_secs(1), + slots.clone().acquire_owned(), + ) + .await + .expect("the admission slot must be released after SDK task completion") + .expect("test semaphore remains open"); + drop(permit); + } + + #[tokio::test] + async fn saturated_provider_call_is_rejected_transiently_then_admitted_after_release() { + use crate::chat::relay::{ReadEvent, RelayRead}; + use crate::testkit::fake_channels::FakeChannel; + use std::time::Duration; + + let slots = Arc::new(tokio::sync::Semaphore::new(1)); + let aborted = std::sync::atomic::AtomicBool::new(false); + let ProviderCallAdmission::Admitted(first_slot) = + try_acquire_provider_call_slot(slots.clone(), &aborted) + else { + panic!("first provider call must be admitted"); + }; + assert!( + matches!( + try_acquire_provider_call_slot(slots.clone(), &aborted), + ProviderCallAdmission::Saturated + ), + "a full limiter must reject synchronously rather than queue" + ); + + let ch = FakeChannel::new(); + let response = + provider_call_saturated_response(&ch.writer, None, "busy-model", "busy-provider", None); + assert!(!response.ok); + assert_eq!(response.stop_reason, Some(StopReason::Error)); + assert_eq!(response.error.as_ref().unwrap().code, "transient"); + + let mut reader = ch.reader; + let ReadEvent::Msg(frame) = reader.next(Duration::from_millis(50)).await else { + panic!("saturation must emit a terminal error frame"); + }; + let AssistantMessageEvent::Error { error } = + serde_json::from_str::(&frame).unwrap() + else { + panic!("saturation terminal must be an error"); + }; + assert_eq!(error.error_kind, Some(ErrorKind::Transient)); + assert!(matches!( + reader.next(Duration::from_millis(25)).await, + ReadEvent::Timeout + )); + + aborted.store(true, Ordering::SeqCst); + assert!(matches!( + try_acquire_provider_call_slot(slots.clone(), &aborted), + ProviderCallAdmission::Aborted + )); + aborted.store(false, Ordering::SeqCst); + + drop(first_slot); + let ProviderCallAdmission::Admitted(_next_slot) = + try_acquire_provider_call_slot(slots, &aborted) + else { + panic!("released capacity must admit the next provider call"); + }; + } + + #[tokio::test] + async fn dispatch_error_published_before_reader_close_survives_task_cleanup() { + let (outcome_tx, outcome_rx) = tokio::sync::oneshot::channel::(); + let (published_tx, published_rx) = tokio::sync::oneshot::channel(); + let call_task = tokio::spawn(async move { + let _ = outcome_tx.send(Err(remote("function_not_found"))); + let _ = published_tx.send(()); + }); + published_rx.await.unwrap(); + + let outcome = finish_provider_call(call_task, outcome_rx) + .await + .expect("published outcome must be retained"); + assert!(is_function_not_found(&outcome.unwrap_err())); + } + + #[tokio::test] + async fn provider_unavailable_failure_emits_exactly_one_terminal_error() { + use crate::chat::relay::{ReadEvent, RelayRead}; + use crate::testkit::fake_channels::FakeChannel; + use std::time::Duration; + + let ch = FakeChannel::new(); + let err = fail_with_terminal( + &ch.writer, + None, + "ghost-model", + "ghost", + RouterCode::ProviderUnavailable, + "provider ghost unavailable".into(), + ErrorKind::Transient, + None, + ); + assert!(matches!( + err, + Error::Remote { ref code, .. } if code == "router/provider_unavailable" + )); + + let mut reader = ch.reader; + let ReadEvent::Msg(frame) = reader.next(Duration::from_millis(50)).await else { + panic!("expected terminal error frame"); + }; + let event: AssistantMessageEvent = serde_json::from_str(&frame).unwrap(); + let AssistantMessageEvent::Error { error } = event else { + panic!("provider_unavailable must emit an error terminal"); + }; + assert_eq!(error.error_kind, Some(ErrorKind::Transient)); + assert_eq!( + error.error_message.as_deref(), + Some("provider ghost unavailable") + ); + assert!(matches!( + reader.next(Duration::from_millis(25)).await, + ReadEvent::Timeout + )); + } + /// A pre-stream failure (here: a model that routes to no provider) must /// still emit exactly one terminal error frame to the sink. Without it, /// `router::complete`'s drain blocks for the full reader budget and diff --git a/llm-router/src/chat/inflight.rs b/llm-router/src/chat/inflight.rs index af0c830bd..c025c7754 100644 --- a/llm-router/src/chat/inflight.rs +++ b/llm-router/src/chat/inflight.rs @@ -2,6 +2,7 @@ //! an abort must reach the router instance holding the stream //! (single-instance worker — design doc § router::abort). use std::collections::HashMap; +use std::ops::Deref; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; @@ -16,13 +17,22 @@ pub struct InflightEntry { impl InflightEntry { /// chat.rs points this at the current attempt's provider-channel closer. pub fn set_closer(&self, closer: Arc) { - *self.closer.lock().unwrap() = Some(closer); + let mut slot = self.closer.lock().unwrap(); + if self.aborted.load(Ordering::SeqCst) { + drop(slot); + closer(); + } else { + *slot = Some(closer); + } } - pub fn abort(&self) { - self.aborted.store(true, Ordering::SeqCst); + pub fn abort(&self) -> bool { + if self.aborted.swap(true, Ordering::SeqCst) { + return false; + } if let Some(close) = self.closer.lock().unwrap().clone() { close(); } + true } } @@ -31,28 +41,60 @@ pub struct InflightMap { entries: Mutex>, } +/// Owns one request-id reservation. Dropping the chat future releases only +/// this exact entry, so cancellation cannot leak an id or remove a newer call. +pub struct InflightReservation { + map: Arc, + request_id: String, + entry: InflightEntry, +} + +impl Deref for InflightReservation { + type Target = InflightEntry; + + fn deref(&self) -> &Self::Target { + &self.entry + } +} + +impl Drop for InflightReservation { + fn drop(&mut self) { + self.map.remove_if_same(&self.request_id, &self.entry); + } +} + impl InflightMap { - pub fn insert(&self, request_id: &str) -> InflightEntry { + /// Atomically reserves a request id. A live request keeps the reservation + /// after `router::abort` until its chat future actually terminates. + pub fn reserve(self: &Arc, request_id: &str) -> Option { let entry = InflightEntry::default(); - self.entries - .lock() - .unwrap() - .insert(request_id.to_string(), entry.clone()); - entry + let mut entries = self.entries.lock().unwrap(); + if entries.contains_key(request_id) { + return None; + } + entries.insert(request_id.to_string(), entry.clone()); + Some(InflightReservation { + map: self.clone(), + request_id: request_id.to_string(), + entry, + }) } - pub fn remove(&self, request_id: &str) { - self.entries.lock().unwrap().remove(request_id); + + fn remove_if_same(&self, request_id: &str, expected: &InflightEntry) { + let mut entries = self.entries.lock().unwrap(); + if entries + .get(request_id) + .is_some_and(|current| Arc::ptr_eq(¤t.aborted, &expected.aborted)) + { + entries.remove(request_id); + } } - /// false when unknown or already terminal (spec: AbortResponse.aborted). + + /// False when unknown, already aborted, or terminal + /// (spec: AbortResponse.aborted). pub fn abort(&self, request_id: &str) -> bool { - let entry = self.entries.lock().unwrap().remove(request_id); - match entry { - Some(e) => { - e.abort(); - true - } - None => false, - } + let entry = self.entries.lock().unwrap().get(request_id).cloned(); + entry.is_some_and(|entry| entry.abort()) } } @@ -63,8 +105,8 @@ mod tests { #[test] fn abort_fires_once_sets_flag_and_runs_the_closer() { - let map = InflightMap::default(); - let entry = map.insert("r1"); + let map = Arc::new(InflightMap::default()); + let entry = map.reserve("r1").unwrap(); let closed = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); let c2 = closed.clone(); entry.set_closer(std::sync::Arc::new(move || { @@ -73,9 +115,41 @@ mod tests { assert!(map.abort("r1")); assert!(entry.aborted.load(Ordering::SeqCst)); assert_eq!(closed.load(Ordering::SeqCst), 1); - assert!(!map.abort("r1")); // already terminal - map.insert("r2"); - map.remove("r2"); + assert!(!map.abort("r1")); // already aborted + let r2 = map.reserve("r2").unwrap(); + drop(r2); // terminal chat releases the reservation assert!(!map.abort("r2")); } + + #[test] + fn duplicate_reservation_is_rejected_without_orphaning_the_first() { + let map = Arc::new(InflightMap::default()); + let first = map.reserve("same-id").unwrap(); + + assert!(map.reserve("same-id").is_none()); + assert!(map.abort("same-id"), "the original remains abortable"); + assert!(first.aborted.load(Ordering::SeqCst)); + assert!( + map.reserve("same-id").is_none(), + "abort does not release an id while its chat is still winding down" + ); + + drop(first); + assert!(map.reserve("same-id").is_some()); + } + + #[test] + fn closer_installed_after_abort_is_closed_immediately() { + let map = Arc::new(InflightMap::default()); + let entry = map.reserve("r1").unwrap(); + assert!(map.abort("r1")); + + let closed = Arc::new(std::sync::atomic::AtomicU32::new(0)); + let observed = closed.clone(); + entry.set_closer(Arc::new(move || { + observed.fetch_add(1, Ordering::SeqCst); + })); + + assert_eq!(closed.load(Ordering::SeqCst), 1); + } } diff --git a/llm-router/src/count_tokens.rs b/llm-router/src/count_tokens.rs index bdec99be7..b54c9b571 100644 --- a/llm-router/src/count_tokens.rs +++ b/llm-router/src/count_tokens.rs @@ -106,10 +106,19 @@ pub fn make_count_tokens( let config = snapshot(&config); let heuristics = config.settings().routing_heuristics.clone(); let default_provider = config.settings().default_provider.clone(); + let providers = registry.list().await; let candidates = decide(&DecideInput { model: model.clone(), provider: req.provider, - registered_providers: registry.ids().await, + registered_providers: providers + .iter() + .map(|record| record.declaration.id.clone()) + .collect(), + available_providers: providers + .iter() + .filter(|record| record.available) + .map(|record| record.declaration.id.clone()) + .collect(), catalog: catalog.model_ids().await, heuristics, default_provider, diff --git a/llm-router/src/registry/register.rs b/llm-router/src/registry/register.rs index efaf0a5b7..94ecee995 100644 --- a/llm-router/src/registry/register.rs +++ b/llm-router/src/registry/register.rs @@ -1,6 +1,8 @@ //! The `router::provider::register` iii function (spec § register, -//! § Registration lifecycle): validate → token-gated upsert → entry-schema -//! re-compose (under the entry write lock) → static models reconcile → emits. +//! § Registration lifecycle): validate → token-gated prepare → entry-schema +//! re-compose → durable static models → registry commit → publish + emits. +//! The entry write lock spans that transaction so concurrent provider boots +//! cannot invalidate one another's staged schemas or tokens. //! //! Engine-backed coverage: tests/integration.rs (declare, token gate, //! takeover rejection, schema composition). @@ -16,7 +18,7 @@ use serde_json::{json, Value}; use crate::catalog::store::CatalogStore; use crate::config::entry::{register_entry, EntryWriteLock}; use crate::config::schema::{provider_entry_schema, validate_custom_schema}; -use crate::registry::store::RegistryStore; +use crate::registry::store::{ProviderRecord, RegistryStore}; use crate::triggers::{self, RouterEvents}; fn valid_id(id: &str) -> bool { @@ -27,6 +29,31 @@ fn valid_id(id: &str) -> bool { .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_' || c == '-') } +fn schemas_for(records: &[ProviderRecord]) -> BTreeMap { + records + .iter() + .map(|rec| { + 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(), + ); + (rec.declaration.id.clone(), schema) + }) + .collect() +} + +fn failure_with_rollbacks(original: Error, failures: Vec) -> Error { + if failures.is_empty() { + original + } else { + Error::Handler(format!( + "{original}; registration rollback incomplete: {}", + failures.join("; ") + )) + } +} + pub fn make_provider_register( iii: IIIClient, registry: Arc, @@ -72,33 +99,89 @@ pub fn make_provider_register( } } } - let worker_id = declaration.worker_id.clone(); let static_models = declaration.models.clone(); let id = declaration.id.clone(); - let upserted = registry - .upsert(declaration, worker_id, input.token) + // This lock is the provider-registration transaction boundary. It + // also serializes against configuration writes/reloads. + let entry_guard = entry_lock.lock().await; + let prepared = registry + .prepare_upsert(declaration, worker_id, input.token) .await .map_err(Error::from)?; - let token = upserted.token; - let availability_recovered = upserted.availability_recovered; + let current_records = registry.list().await; + let previous_schemas = schemas_for(¤t_records); + let mut candidate_records = current_records; + candidate_records.retain(|record| record.declaration.id != id); + candidate_records.push(prepared.record().clone()); + let candidate_schemas = schemas_for(&candidate_records); - // Re-compose the entry schema from every registered declaration — - // under the entry write lock so concurrent boots compose. - { - let _guard = entry_lock.lock().await; - 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); + // Configuration is the first fallible external step. If the call + // reports failure after a partial remote apply, restore the prior + // schema while the transaction lock is still held. + if let Err(error) = register_entry(&iii, &candidate_schemas).await { + let mut rollback_failures = Vec::new(); + if let Err(rollback) = register_entry(&iii, &previous_schemas).await { + rollback_failures.push(format!("configuration: {rollback}")); + } + drop(entry_guard); + return Err(failure_with_rollbacks(error, rollback_failures)); + } + + // Persist static models without exposing them in memory. A failed + // catalog write leaves the registry untouched, so a fresh token + // can be minted again on retry. + let mut prepared_catalog = if let Some(models) = static_models { + if models.is_empty() { + None + } else { + match catalog.prepare_slice(&id, models).await { + Ok(prepared) => Some(prepared), + Err(error) => { + let mut rollback_failures = Vec::new(); + if let Err(rollback) = register_entry(&iii, &previous_schemas).await { + rollback_failures.push(format!("configuration: {rollback}")); + } + drop(entry_guard); + return Err(failure_with_rollbacks(error, rollback_failures)); + } + } } - register_entry(&iii, &provider_schemas).await?; + } else { + None + }; + + // The token hash becomes durable only after configuration and the + // optional catalog slice succeeded. If registry persistence fails, + // keep the staged catalog invisible and compensate both durable + // side effects before returning the original failure. + let upserted = match registry.commit_upsert(prepared).await { + Ok(upserted) => upserted, + Err(error) => { + let mut rollback_failures = Vec::new(); + if let Some(prepared) = prepared_catalog.take() { + if let Err(rollback) = prepared.rollback().await { + // No registry record was published, so routing + // excludes this owner even if its durable catalog + // rollback also fails. A restart may temporarily + // reload that orphan until the next reconcile. + rollback_failures.push(format!("catalog: {rollback}")); + } + } + if let Err(rollback) = register_entry(&iii, &previous_schemas).await { + rollback_failures.push(format!("configuration: {rollback}")); + } + drop(entry_guard); + return Err(failure_with_rollbacks(error.into(), rollback_failures)); + } + }; + if let Some(prepared) = prepared_catalog { + prepared.commit(); } + drop(entry_guard); + + let token = upserted.token; + let availability_recovered = upserted.availability_recovered; events .emit( @@ -119,18 +202,20 @@ pub fn make_provider_register( .await; } - // Static catalog slice: reconciled at registration (spec § register). - if let Some(models) = static_models { - if !models.is_empty() { - let count = models.len(); - catalog.set_slice(&id, models).await?; - events - .emit( - triggers::MODELS_CHANGED, - json!({ "provider": id, "count": count }), - ) - .await; - } + if let Some(count) = upserted + .record + .declaration + .models + .as_ref() + .filter(|models| !models.is_empty()) + .map(Vec::len) + { + events + .emit( + triggers::MODELS_CHANGED, + json!({ "provider": id, "count": count }), + ) + .await; } Ok(ProviderRegisterResponse { diff --git a/llm-router/src/registry/store.rs b/llm-router/src/registry/store.rs index 0f38b4322..394bc41bd 100644 --- a/llm-router/src/registry/store.rs +++ b/llm-router/src/registry/store.rs @@ -1,5 +1,5 @@ //! Durable provider registry. Single writer: the records Mutex is held across -//! mutate + state::set (spec § Registration lifecycle, "Serialized merges"). +//! prepare + state::set + publish (spec § Registration lifecycle, "Serialized merges"). //! Persistence is iii state (`state::get`/`state::set` engine functions via //! src/state.rs) under scope "llm-router"; the router is the only writer of //! its state keys (single-instance worker). @@ -37,11 +37,18 @@ pub struct ProviderRecord { pub worker_id: Option, pub available: bool, pub registered_at: i64, + /// Monotonic identity for one successful registration of this provider. + /// Legacy snapshots deserialize as generation zero; their first + /// re-registration advances to one. + #[serde(default)] + pub generation: u64, } pub struct RegistryStore { iii: IIIClient, records: Mutex>, + #[cfg(test)] + persist_result: Option>, } /// Outcome of `upsert`: the stored record, the raw registration token (it @@ -53,6 +60,27 @@ pub struct Upserted { pub availability_recovered: bool, } +/// Token-gated registration assembled without changing durable or in-memory +/// state. The register handler may safely perform its other fallible setup +/// before passing this value to `commit_upsert`. +pub struct PreparedUpsert { + record: ProviderRecord, + token: String, + expected: Option, +} + +#[derive(Clone)] +struct RegistrationRevision { + token_hash: String, + generation: u64, +} + +impl PreparedUpsert { + pub fn record(&self) -> &ProviderRecord { + &self.record + } +} + /// Pure record assembly for upsert (no I/O — persistence is the caller's job). /// Token hash and registered_at are preserved across a re-register; the rest /// is recomposed from the new declaration. @@ -72,10 +100,33 @@ fn build_record( // flips this back down (chat.rs). available: true, registered_at: existing.map(|e| e.registered_at).unwrap_or_else(now_ms), + generation: existing + .map(|e| e.generation.wrapping_add(1).max(1)) + .unwrap_or(1), declaration, } } +fn revision(record: &ProviderRecord) -> RegistrationRevision { + RegistrationRevision { + token_hash: record.token_hash.clone(), + generation: record.generation, + } +} + +fn revision_matches( + current: Option<&ProviderRecord>, + expected: Option<&RegistrationRevision>, +) -> bool { + match (current, expected) { + (None, None) => true, + (Some(current), Some(expected)) => { + current.token_hash == expected.token_hash && current.generation == expected.generation + } + _ => false, + } +} + /// Whether this registration brings a previously-down provider back up. A fresh /// register is not a recovery (the `op:"register"` event already signals /// presence); only a known provider whose `available` flag was false @@ -90,6 +141,8 @@ impl RegistryStore { Self { iii, records: Mutex::new(HashMap::new()), + #[cfg(test)] + persist_result: None, } } @@ -122,6 +175,10 @@ impl RegistryStore { } async fn persist(&self, records: &HashMap) -> Result<(), Error> { + #[cfg(test)] + if let Some(result) = &self.persist_result { + return result.clone(); + } let value = serde_json::to_value(records).unwrap_or_default(); state_set(&self.iii, REGISTRY_KEY, value).await } @@ -138,13 +195,13 @@ impl RegistryStore { /// First register binds (mints a token, persists its hash); later /// registers must present the raw token. Returns the raw token — it /// exists nowhere else. - pub async fn upsert( + pub async fn prepare_upsert( &self, declaration: ProviderDeclaration, worker_id: Option, token: Option, - ) -> Result { - let mut records = self.records.lock().await; // serialized writer + ) -> Result { + let records = self.records.lock().await; let existing = records.get(&declaration.id); if let Some(existing) = existing { let presented = token.as_deref().map(hash_token); @@ -159,22 +216,57 @@ impl RegistryStore { } } let raw_token = token.unwrap_or_else(|| Uuid::new_v4().to_string()); - let recovered = availability_recovered(existing); let record = build_record(existing, declaration, worker_id, &raw_token); - records.insert(record.declaration.id.clone(), record.clone()); - self.persist(&records).await.map_err(|e| { + Ok(PreparedUpsert { + record, + token: raw_token, + expected: existing.map(revision), + }) + } + + /// Persist and publish a prepared registration if its predecessor is + /// still current. The revision check makes the two-phase API safe even if + /// a caller fails to hold the higher-level registration lock. + pub async fn commit_upsert(&self, prepared: PreparedUpsert) -> Result { + let mut records = self.records.lock().await; // serialized writer + let id = prepared.record.declaration.id.clone(); + let existing = records.get(&id); + if !revision_matches(existing, prepared.expected.as_ref()) { + return Err(RouterError::new( + RouterCode::RegistrationRejected, + format!("provider {id} changed while registration was being prepared; retry"), + )); + } + let recovered = availability_recovered(existing); + let mut next_records = records.clone(); + next_records.insert(id, prepared.record.clone()); + self.persist(&next_records).await.map_err(|e| { RouterError::new( RouterCode::InvalidRequest, format!("registry persist failed: {e}"), ) })?; + // Publish only after the durable write succeeds. In particular, a + // failed first registration must not retain a hash for the raw token + // that the caller never received. + *records = next_records; Ok(Upserted { - record, - token: raw_token, + record: prepared.record, + token: prepared.token, availability_recovered: recovered, }) } + pub async fn upsert( + &self, + declaration: ProviderDeclaration, + worker_id: Option, + token: Option, + ) -> Result { + let prepared = self.prepare_upsert(declaration, worker_id, token).await?; + self.commit_upsert(prepared).await + } + /// Token gate for resolve / reconcile / update_credential (and re-register). pub async fn verify_token( &self, @@ -197,19 +289,30 @@ impl RegistryStore { } } - /// Returns true when the flag actually changed (callers emit on change only). - pub async fn set_availability(&self, id: &str, available: bool) -> bool { + /// Change availability only for the registration that originated an + /// observed dispatch result. Availability is deliberately memory-only: + /// persisted flags are stale after a restart and `load` always restores + /// them as down. Avoiding a full-registry state write here also keeps a + /// state-worker outage out of the dispatch hot path and prevents an old + /// availability snapshot from overwriting a newer registration. + /// Returns true when the in-memory flag changed. + pub async fn set_availability_if_current( + &self, + id: &str, + generation: u64, + available: bool, + ) -> bool { let mut records = self.records.lock().await; - let Some(rec) = records.get_mut(id) else { + let Some(rec) = records.get(id) else { return false; }; - if rec.available == available { + if rec.generation != generation || rec.available == available { return false; } - rec.available = available; - let snapshot = records.clone(); - drop(records); - let _ = self.persist(&snapshot).await; // best-effort persist of a flag flip + records + .get_mut(id) + .expect("record came from the same snapshot") + .available = available; true } } @@ -223,6 +326,14 @@ mod tests { serde_json::from_value(json!({ "id": id })).expect("minimal declaration") } + fn store_with_persistence(result: Result<(), Error>) -> RegistryStore { + RegistryStore { + iii: IIIClient::new("ws://unused.invalid"), + records: Mutex::new(HashMap::new()), + persist_result: Some(result), + } + } + // F9: the engine keys topology events by a per-connection UUID while the // registry stores the provider's self-declared name, so the availability // handler can never resolve an event to a provider. A provider that just @@ -236,6 +347,7 @@ mod tests { record.available, "a provider that just registered must be marked available" ); + assert_eq!(record.generation, 1); } #[test] @@ -250,6 +362,7 @@ mod tests { back.available, "re-registering brings a downed provider back up" ); + assert_eq!(back.generation, down.generation + 1); } // A down→up transition on re-register must be reported so the register @@ -268,9 +381,159 @@ mod tests { assert!(!availability_recovered(None)); } + #[test] + fn legacy_record_without_generation_deserializes_as_zero() { + let record: ProviderRecord = serde_json::from_value(json!({ + "declaration": { "id": "anthropic" }, + "token_hash": "hash", + "worker_id": "w-legacy", + "available": true, + "registered_at": 1 + })) + .expect("legacy registry snapshot remains readable"); + assert_eq!(record.generation, 0); + let next = build_record( + Some(&record), + decl("anthropic"), + Some("w-new".into()), + "ignored", + ); + assert_eq!(next.generation, 1); + } + #[test] fn reregistering_an_up_provider_is_not_a_recovery() { let up = build_record(None, decl("anthropic"), Some("w-1".into()), "tok"); // available: true assert!(!availability_recovered(Some(&up))); } + + #[tokio::test(start_paused = true)] + async fn failed_persist_does_not_publish_registration_or_orphan_token() { + // An unconnected client makes state::set time out. Tokio's paused + // clock advances directly to the SDK timeout, so this stays a fast + // unit test without a live engine. + let store = RegistryStore::new(IIIClient::new("ws://unconnected.invalid")); + + for _ in 0..2 { + let err = store + .upsert(decl("anthropic"), Some("w-1".into()), None) + .await + .err() + .expect("persistence must fail"); + assert_eq!(err.code, RouterCode::InvalidRequest); + assert!(err.message.contains("registry persist failed")); + assert!( + store.get("anthropic").await.is_none(), + "a failed write must not publish a record or its token hash" + ); + } + } + + #[tokio::test(start_paused = true)] + async fn failed_persist_keeps_previous_registration_unchanged() { + let store = RegistryStore::new(IIIClient::new("ws://unconnected.invalid")); + let original = ProviderRecord { + available: false, + ..build_record(None, decl("anthropic"), Some("w-old".into()), "tok") + }; + store + .records + .lock() + .await + .insert("anthropic".into(), original.clone()); + + let mut replacement = decl("anthropic"); + replacement.display_name = Some("replacement".into()); + let err = store + .upsert(replacement, Some("w-new".into()), Some("tok".into())) + .await + .err() + .expect("persistence must fail"); + assert_eq!(err.code, RouterCode::InvalidRequest); + + let current = store + .get("anthropic") + .await + .expect("the previous registration must remain"); + assert_eq!(current.token_hash, original.token_hash); + assert_eq!(current.worker_id, original.worker_id); + assert_eq!(current.registered_at, original.registered_at); + assert_eq!(current.available, original.available); + assert_eq!(current.declaration, original.declaration); + } + + #[tokio::test] + async fn stale_prepared_registration_is_rejected_after_another_commit() { + let store = store_with_persistence(Ok(())); + let stale = store + .prepare_upsert(decl("anthropic"), Some("w-stale".into()), None) + .await + .expect("first prepare"); + let winner = store + .prepare_upsert(decl("anthropic"), Some("w-winner".into()), None) + .await + .expect("concurrent prepare"); + let winner = store.commit_upsert(winner).await.expect("winner commits"); + + let err = store + .commit_upsert(stale) + .await + .err() + .expect("stale candidate must lose the compare-and-set"); + assert_eq!(err.code, RouterCode::RegistrationRejected); + let current = store.get("anthropic").await.expect("winner remains"); + assert_eq!(current.worker_id.as_deref(), Some("w-winner")); + assert_eq!(current.generation, winner.record.generation); + } + + #[tokio::test] + async fn old_generation_cannot_change_reregistered_provider_availability() { + let store = store_with_persistence(Ok(())); + let first = store + .upsert(decl("anthropic"), Some("w-1".into()), None) + .await + .expect("first registration"); + let old_generation = first.record.generation; + let second = store + .upsert(decl("anthropic"), Some("w-2".into()), Some(first.token)) + .await + .expect("re-registration"); + assert!(second.record.generation > old_generation); + + assert!( + !store + .set_availability_if_current("anthropic", old_generation, false) + .await, + "a late function_not_found from the old registration must be ignored" + ); + let current = store.get("anthropic").await.expect("provider remains"); + assert_eq!(current.generation, second.record.generation); + assert!(current.available); + } + + #[tokio::test] + async fn availability_change_does_not_depend_on_state_persistence() { + let store = store_with_persistence(Err(Error::Timeout)); + let record = build_record(None, decl("anthropic"), Some("w-1".into()), "tok"); + let generation = record.generation; + store + .records + .lock() + .await + .insert("anthropic".into(), record); + + assert!( + store + .set_availability_if_current("anthropic", generation, false) + .await + ); + assert!( + !store + .get("anthropic") + .await + .expect("provider remains") + .available, + "availability must remain operational while durable state is unavailable" + ); + } } diff --git a/llm-router/src/routing.rs b/llm-router/src/routing.rs index 6db153e31..015ff1654 100644 --- a/llm-router/src/routing.rs +++ b/llm-router/src/routing.rs @@ -26,6 +26,7 @@ pub struct DecideInput { pub model: String, pub provider: Option, pub registered_providers: Vec, + pub available_providers: Vec, pub catalog: Vec<(String, Vec)>, // provider id -> model ids pub heuristics: Vec, pub default_provider: Option, @@ -33,6 +34,7 @@ pub struct DecideInput { pub fn decide(input: &DecideInput) -> Result, RouterError> { let registered = |id: &str| input.registered_providers.iter().any(|p| p == id); + let available = |id: &str| input.available_providers.iter().any(|p| p == id); // 1. Explicit provider — sole candidate; cold-catalog tolerant; typos loud. if let Some(provider) = &input.provider { @@ -42,33 +44,52 @@ pub fn decide(input: &DecideInput) -> Result, RouterError> { format!("unknown provider {provider}"), )); } + if !available(provider) { + return Err(RouterError::new( + RouterCode::ProviderUnavailable, + format!("provider {provider} unavailable"), + )); + } return Ok(vec![provider.clone()]); } - // 2. Unique catalog owner; 2+ owners → ambiguous (the router never guesses). - let mut owners: Vec<&str> = input + // 2. Unique available catalog owner; 2+ available owners → ambiguous (the + // router never guesses). Ignore stale catalog slices belonging to a + // missing/down provider so route previews cannot select a known-dead + // dispatch target. + let owners: Vec<&str> = input .catalog .iter() - .filter(|(_, ids)| ids.iter().any(|m| m == &input.model)) + .filter(|(provider, ids)| registered(provider) && ids.iter().any(|m| m == &input.model)) .map(|(p, _)| p.as_str()) .collect(); - owners.sort_unstable(); - match owners.len() { - 1 => return Ok(vec![owners[0].to_string()]), + let mut available_owners: Vec<&str> = owners + .iter() + .copied() + .filter(|provider| available(provider)) + .collect(); + available_owners.sort_unstable(); + match available_owners.len() { + 1 => return Ok(vec![available_owners[0].to_string()]), n if n > 1 => { return Err(RouterError::new( RouterCode::AmbiguousModel, format!( "ambiguous model {} (providers: {})", input.model, - owners.join(", ") + available_owners.join(", ") ), )) } _ => {} } + let mut unavailable_matches: Vec<&str> = owners + .iter() + .copied() + .filter(|provider| !available(provider)) + .collect(); - // 3. Operator heuristics from the llm-router entry; first match wins. + // 3. Operator heuristics from the llm-router entry; first available match wins. for h in &input.heuristics { if !registered(&h.provider) { continue; @@ -77,18 +98,39 @@ pub fn decide(input: &DecideInput) -> Result, RouterError> { continue; // an invalid operator regex never takes the router down }; if re.is_match(&input.model) { - return Ok(vec![h.provider.clone()]); + if available(&h.provider) { + return Ok(vec![h.provider.clone()]); + } + unavailable_matches.push(&h.provider); } } // 4. Configured default provider makes routing a total function. if let Some(default) = &input.default_provider { - if registered(default) { + if registered(default) && available(default) { return Ok(vec![default.clone()]); } + if registered(default) { + unavailable_matches.push(default); + } + } + + // 5. Preserve an actionable distinction between an unknown model and a + // model whose catalog, heuristic, or default candidates are all down. + unavailable_matches.sort_unstable(); + unavailable_matches.dedup(); + if !unavailable_matches.is_empty() { + return Err(RouterError::new( + RouterCode::ProviderUnavailable, + format!( + "no available provider for model {} (unavailable: {})", + input.model, + unavailable_matches.join(", ") + ), + )); } - // 5. Loud failure. + // 6. Loud failure. Err(RouterError::new( RouterCode::NoProviderForModel, format!("no provider registered for model {}", input.model), @@ -115,10 +157,19 @@ pub fn make_route( let config = snapshot(&config); let heuristics = config.settings().routing_heuristics.clone(); let default_provider = config.settings().default_provider.clone(); + let providers = registry.list().await; let candidates = decide(&DecideInput { model: req.model, provider: req.provider, - registered_providers: registry.ids().await, + registered_providers: providers + .iter() + .map(|record| record.declaration.id.clone()) + .collect(), + available_providers: providers + .iter() + .filter(|record| record.available) + .map(|record| record.declaration.id.clone()) + .collect(), catalog: catalog.model_ids().await, heuristics, default_provider, @@ -142,6 +193,7 @@ mod tests { model: String::new(), provider: None, registered_providers: vec!["anthropic".into(), "openai".into(), "lmstudio".into()], + available_providers: vec!["anthropic".into(), "openai".into(), "lmstudio".into()], catalog: vec![ ("anthropic".into(), vec!["claude-sonnet-4".into()]), ("openai".into(), vec!["gpt-5".into(), "shared-model".into()]), @@ -218,4 +270,61 @@ mod tests { input.default_provider = Some("anthropic".into()); assert_eq!(decide(&input).unwrap(), vec!["anthropic"]); } + + #[test] + fn explicit_unavailable_provider_is_distinct_from_unknown_provider() { + let mut input = base(); + input.model = "gpt-5".into(); + input.provider = Some("openai".into()); + input.available_providers.retain(|id| id != "openai"); + + let err = decide(&input).unwrap_err(); + assert_eq!(err.code, RouterCode::ProviderUnavailable); + assert_eq!(err.message, "provider openai unavailable"); + } + + #[test] + fn stale_catalog_owner_is_excluded_from_routing() { + let mut input = base(); + input.model = "shared-model".into(); + input.available_providers.retain(|id| id != "lmstudio"); + + assert_eq!(decide(&input).unwrap(), vec!["openai"]); + + input.available_providers.retain(|id| id != "openai"); + let err = decide(&input).unwrap_err(); + assert_eq!(err.code, RouterCode::ProviderUnavailable); + assert_eq!( + err.message, + "no available provider for model shared-model (unavailable: lmstudio, openai)" + ); + } + + #[test] + fn stale_owner_does_not_block_an_available_default() { + let mut input = base(); + input.model = "local-llama".into(); + input.available_providers.retain(|id| id != "lmstudio"); + input.default_provider = Some("anthropic".into()); + + assert_eq!(decide(&input).unwrap(), vec!["anthropic"]); + } + + #[test] + fn unavailable_heuristic_or_default_reports_provider_unavailable() { + let mut input = base(); + input.model = "gpt-future".into(); + input.available_providers.retain(|id| id != "openai"); + + let err = decide(&input).unwrap_err(); + assert_eq!(err.code, RouterCode::ProviderUnavailable); + assert!(err.message.contains("unavailable: openai")); + + input.model = "mystery".into(); + input.heuristics.clear(); + input.default_provider = Some("openai".into()); + let err = decide(&input).unwrap_err(); + assert_eq!(err.code, RouterCode::ProviderUnavailable); + assert!(err.message.contains("unavailable: openai")); + } } diff --git a/llm-router/src/types/router.rs b/llm-router/src/types/router.rs index 7d52f65b3..95cac957b 100644 --- a/llm-router/src/types/router.rs +++ b/llm-router/src/types/router.rs @@ -161,8 +161,9 @@ pub struct SystemPromptGetResponse { pub system_prompt: Option, } -/// registration_token: spec adaptation — the engine exposes no caller identity, -/// so identity binding is a bearer token; only its sha256 hash is persisted. +/// `registration_token` is the provider ownership credential; only its sha256 +/// hash is persisted. Engine caller metadata is not an authorization identity +/// because worker names are self-reported. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct ProviderRegisterResponse { pub ok: bool, diff --git a/llm-router/tests/golden/schemas/router.provider.register.json b/llm-router/tests/golden/schemas/router.provider.register.json index b49de1b5c..ba2f6f57e 100644 --- a/llm-router/tests/golden/schemas/router.provider.register.json +++ b/llm-router/tests/golden/schemas/router.provider.register.json @@ -255,7 +255,7 @@ }, "response_schema": { "$schema": "http://json-schema.org/draft-07/schema#", - "description": "registration_token: spec adaptation — the engine exposes no caller identity, so identity binding is a bearer token; only its sha256 hash is persisted.", + "description": "`registration_token` is the provider ownership credential; only its sha256 hash is persisted. Engine caller metadata is not an authorization identity because worker names are self-reported.", "properties": { "id": { "type": "string" diff --git a/llm-router/tests/integration.rs b/llm-router/tests/integration.rs index ed8a010b1..5dd5a7b93 100644 --- a/llm-router/tests/integration.rs +++ b/llm-router/tests/integration.rs @@ -4,6 +4,7 @@ //! //! **Self-skips** when no engine is available (storage-worker pattern): //! set `III_ENGINE_BIN=/path/to/iii` or have `iii` on PATH. +use std::collections::HashMap; use std::io::Write as _; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; @@ -216,6 +217,60 @@ async fn call_until( } } +/// Register the minimal state surface on a bare engine and fail exactly one +/// numbered write. Used to prove router stores do not publish uncommitted +/// in-memory state when persistence rejects a mutation. +async fn start_flaky_state(url: &str, fail_on_set_call: u64) -> (IIIClient, Arc) { + let iii = register_worker(url, InitOptions::default()); + let values = Arc::new(std::sync::Mutex::new(HashMap::::new())); + let get_values = values.clone(); + iii.register_function( + "state::get", + RegisterFunction::new_async(move |input: Value| { + let values = get_values.clone(); + async move { + let key = input["key"].as_str().unwrap_or_default(); + Ok::(values.lock().unwrap().get(key).cloned().unwrap_or_default()) + } + }), + ); + let set_calls = Arc::new(AtomicU64::new(0)); + let calls = set_calls.clone(); + let set_values = values; + iii.register_function( + "state::set", + RegisterFunction::new_async(move |input: Value| { + let call = calls.fetch_add(1, Ordering::SeqCst) + 1; + let values = set_values.clone(); + async move { + if call == fail_on_set_call { + Err(Error::Handler("injected state write failure".into())) + } else { + let key = input["key"].as_str().unwrap_or_default().to_string(); + values.lock().unwrap().insert(key, input["value"].clone()); + Ok::(json!({ "ok": true })) + } + } + }), + ); + call_until( + &iii, + "engine::functions::list", + json!({ "include_internal": true }), + |value| { + value["functions"].as_array().is_some_and(|functions| { + ["state::get", "state::set"].iter().all(|id| { + functions + .iter() + .any(|function| function["function_id"] == *id) + }) + }) + }, + ) + .await; + (iii, set_calls) +} + fn remote_code(err: &Error) -> &str { match err { Error::Remote { code, .. } => code, @@ -225,9 +280,12 @@ fn remote_code(err: &Error) -> &str { // ── live provider helper ──────────────────────────────────────────────────── -#[derive(Default)] +#[derive(Clone, Default)] struct ProviderOptions { ping_forever: bool, + /// Keep the function handler alive after the terminal frame. A compliant + /// router must complete from the stream terminal, not this RPC lifetime. + done_linger_ms: Option, credential_env_var: Option, supports_model_listing: bool, /// model ids returned by provider::real::refresh_models @@ -237,28 +295,48 @@ struct ProviderOptions { struct LiveProvider { iii: IIIClient, token: String, + stream_calls: Arc, write_failed: Arc, fail_at_ms: Arc, } +fn live_provider_declaration(opts: &ProviderOptions, token: Option) -> Value { + let mut payload = json!({ + "id": "real", + "credential_env_var": opts.credential_env_var, + "defaults": { "api_url": "https://api.example.test/v1", "max_tokens": 8192 }, + "supports_model_listing": opts.supports_model_listing, + "models": [{ "id": "live-1", "provider": "real", "context_window": 100000, "max_output_tokens": 8192 }] + }); + if let Some(token) = token { + payload["token"] = json!(token); + } + payload +} + /// A live provider worker on its own connection: registers /// provider::real::stream (+ refresh_models) and declares itself. async fn start_live_provider(url: &str, opts: ProviderOptions) -> LiveProvider { let iii = register_worker(url, InitOptions::default()); let write_failed = Arc::new(AtomicBool::new(false)); let fail_at_ms = Arc::new(AtomicU64::new(0)); + let stream_calls = Arc::new(AtomicU64::new(0)); let token_cell: Arc>> = Arc::default(); let address = url.to_string(); let wf = write_failed.clone(); let fam = fail_at_ms.clone(); + let calls = stream_calls.clone(); let ping_forever = opts.ping_forever; + let done_linger_ms = opts.done_linger_ms; iii.register_function( "provider::real::stream", RegisterFunction::new_async(move |input: Value| { let address = address.clone(); let (wf, fam) = (wf.clone(), fam.clone()); + let calls = calls.clone(); async move { + calls.fetch_add(1, Ordering::SeqCst); let r: StreamChannelRef = serde_json::from_value(input["writer_ref"].clone()) .map_err(|e| Error::Serde(e.to_string()))?; @@ -307,6 +385,9 @@ async fn start_live_provider(url: &str, opts: ProviderOptions) -> LiveProvider { .map_err(|e| Error::Handler(e.to_string()))?; } let _ = writer.close().await; + if let Some(delay_ms) = done_linger_ms { + tokio::time::sleep(Duration::from_millis(delay_ms)).await; + } Ok(json!({ "ok": true })) } }), @@ -339,19 +420,47 @@ async fn start_live_provider(url: &str, opts: ProviderOptions) -> LiveProvider { ); } + // Match the production provider lifecycle: keep a deterministic direct + // ready handler so a restarted router can rediscover this still-live + // provider even if the one-shot trigger binding was replayed too late. + let iii_ready = iii.clone(); + let token_for_ready = token_cell.clone(); + let opts_for_ready = opts.clone(); + iii.register_function( + "provider::real::on_router_ready", + RegisterFunction::new_async(move |_input: Value| { + let iii = iii_ready.clone(); + let token = token_for_ready.lock().unwrap().clone(); + let opts = opts_for_ready.clone(); + async move { + let token = token.ok_or_else(|| { + Error::Handler("provider ready handler has no registration token".into()) + })?; + iii.trigger(TriggerRequest { + function_id: "router::provider::register".into(), + payload: live_provider_declaration(&opts, Some(token)), + action: None, + timeout_ms: Some(5_000), + }) + .await?; + Ok::(json!({ "ok": true })) + } + }), + ); + let _ = iii.register_trigger(RegisterTriggerInput { + trigger_type: "router::ready".into(), + function_id: "provider::real::on_router_ready".into(), + config: json!({}), + metadata: None, + }); + // declare (with a short retry in case the router is still booting) let mut token = None; for _ in 0..50 { let res = call( &iii, "router::provider::register", - json!({ - "id": "real", - "credential_env_var": opts.credential_env_var, - "defaults": { "api_url": "https://api.example.test/v1", "max_tokens": 8192 }, - "supports_model_listing": opts.supports_model_listing, - "models": [{ "id": "live-1", "provider": "real", "context_window": 100000, "max_output_tokens": 8192 }] - }), + live_provider_declaration(&opts, None), ) .await; if let Ok(res) = res { @@ -366,6 +475,7 @@ async fn start_live_provider(url: &str, opts: ProviderOptions) -> LiveProvider { LiveProvider { iii, token, + stream_calls, write_failed, fail_at_ms, } @@ -595,20 +705,21 @@ async fn registry_survives_a_router_restart_and_token_stays_bound() { .await .expect("router reboots"); - let list = call(&second, "router::provider::list", json!({})) - .await - .expect("provider list"); + // Availability is pessimistically reset on load. The router must recover + // this still-live provider itself through its direct engine-function nudge; + // relying only on the one-shot router::ready replay has a boot race. + let list = call_until(&second, "router::provider::list", json!({}), |value| { + value["providers"].as_array().is_some_and(|providers| { + providers + .iter() + .any(|provider| provider["id"] == "real" && provider["available"] == true) + }) + }) + .await; assert_eq!(list["providers"][0]["id"], "real", "list: {list}"); - // Availability is NOT trusted across a restart: the persisted "up" flag - // could belong to a provider that died while the router was down. It is - // restored DOWN and only re-declaration / a successful dispatch flip it up. - assert_eq!( - list["providers"][0]["available"], - json!(false), - "restored availability must be pessimistic: {list}" - ); - // re-declare with the original token: idempotent, same token accepted + // The automatic re-declare used the original bearer token; it remains + // accepted for an explicit idempotent declaration too. let again = call( &provider.iii, "router::provider::register", @@ -618,16 +729,7 @@ async fn registry_survives_a_router_restart_and_token_stays_bound() { .expect("re-declare accepted"); assert_eq!(again["registration_token"], json!(provider.token.clone())); - // …and the re-declaration is what brings the provider back up. - let list = call(&second, "router::provider::list", json!({})) - .await - .expect("provider list after re-declare"); - assert_eq!( - list["providers"][0]["available"], - json!(true), - "re-declare restores availability: {list}" - ); - + provider.iii.shutdown(); second.shutdown(); } @@ -985,6 +1087,116 @@ async fn abort_stops_the_stream_and_terminates_with_done_aborted() { router_iii.shutdown(); } +#[tokio::test(flavor = "multi_thread")] +async fn duplicate_request_id_is_rejected_without_orphaning_the_original() { + let engine = engine_or_skip!(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + let provider = start_live_provider( + &engine.url, + ProviderOptions { + ping_forever: true, + ..Default::default() + }, + ) + .await; + + let first_consumer = register_worker(&engine.url, InitOptions::default()); + let (first_ref, first_frames, first_pump) = consumer_channel(&first_consumer).await; + let first_chat = { + let consumer = first_consumer.clone(); + tokio::spawn(async move { + consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ + "writer_ref": first_ref, + "request_id": "same-request", + "model": "live-1", + "messages": [] + }), + action: None, + timeout_ms: Some(10_000), + }) + .await + }) + }; + let deadline = Instant::now() + Duration::from_secs(3); + while first_frames.lock().unwrap().is_empty() { + assert!(Instant::now() < deadline, "first stream did not start"); + tokio::time::sleep(Duration::from_millis(25)).await; + } + assert_eq!(provider.stream_calls.load(Ordering::SeqCst), 1); + + let second_consumer = register_worker(&engine.url, InitOptions::default()); + let (second_ref, second_frames, second_pump) = consumer_channel(&second_consumer).await; + let duplicate = second_consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ + "writer_ref": second_ref, + "request_id": "same-request", + "model": "live-1", + "messages": [] + }), + action: None, + timeout_ms: Some(3_000), + }) + .await + .expect_err("duplicate live request id must be rejected"); + assert_eq!( + remote_code(&duplicate), + "router/invalid_request", + "{duplicate:?}" + ); + tokio::time::timeout(Duration::from_secs(2), second_pump) + .await + .expect("duplicate channel EOF") + .expect("duplicate channel pump joins"); + let duplicate_frames: Vec = second_frames + .lock() + .unwrap() + .iter() + .map(|frame| serde_json::from_str(frame).expect("valid frame")) + .collect(); + let duplicate_terminals = duplicate_frames + .iter() + .filter(|frame| matches!(frame["type"].as_str(), Some("done" | "error"))) + .count(); + assert_eq!(duplicate_terminals, 1, "frames: {duplicate_frames:?}"); + assert_eq!(duplicate_frames.last().unwrap()["type"], "error"); + assert_eq!( + provider.stream_calls.load(Ordering::SeqCst), + 1, + "duplicate request reached the provider" + ); + assert!(!first_chat.is_finished(), "original stream was disturbed"); + + let aborted = call( + &router, + "router::abort", + json!({ "request_id": "same-request" }), + ) + .await + .expect("abort original"); + assert_eq!(aborted["aborted"], true, "{aborted}"); + let first = tokio::time::timeout(Duration::from_secs(3), first_chat) + .await + .expect("original chat resolves") + .expect("original task joins") + .expect("original chat response"); + assert_eq!(first["stop_reason"], "aborted", "{first}"); + tokio::time::timeout(Duration::from_secs(2), first_pump) + .await + .expect("original channel EOF") + .expect("original channel pump joins"); + + first_consumer.shutdown(); + second_consumer.shutdown(); + provider.iii.shutdown(); + router.shutdown(); +} + #[tokio::test(flavor = "multi_thread")] async fn update_credential_persists_and_resolves_back() { let engine = engine_or_skip!(); @@ -1132,12 +1344,604 @@ async fn router_boots_its_interface_against_a_bare_engine() { router_iii.shutdown(); } +#[tokio::test(flavor = "multi_thread")] +async fn failed_registry_write_rolls_back_and_retry_can_bind_without_a_token() { + let engine = bare_engine_or_skip!(); + let (state, set_calls) = start_flaky_state(&engine.url, 1).await; + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + + let first = call( + &router, + "router::provider::register", + json!({ "id": "recoverable" }), + ) + .await + .expect_err("the injected first state write must fail"); + assert_eq!(remote_code(&first), "router/invalid_request", "{first:?}"); + + let after_failure = call(&router, "router::provider::list", json!({})) + .await + .expect("provider list"); + assert_eq!( + after_failure["providers"], + json!([]), + "failed persistence must not publish a ghost provider: {after_failure}" + ); + + let retry = call( + &router, + "router::provider::register", + json!({ "id": "recoverable" }), + ) + .await + .expect("retry without an undelivered token must bind cleanly"); + assert_eq!(retry["ok"], true, "{retry}"); + assert!( + retry["registration_token"] + .as_str() + .is_some_and(|token| !token.is_empty()), + "registration token missing: {retry}" + ); + assert_eq!(set_calls.load(Ordering::SeqCst), 2); + + router.shutdown(); + state.shutdown(); +} + +async fn assert_static_registration_failure_is_atomic( + fail_on_set_call: u64, + expected_set_calls_after_retry: u64, +) { + let engine = bare_engine_or_skip!(); + let (state, set_calls) = start_flaky_state(&engine.url, fail_on_set_call).await; + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + + let provider_id = format!("atomic-static-{fail_on_set_call}"); + let model_id = format!("atomic-model-{fail_on_set_call}"); + let declaration = json!({ + "id": provider_id, + "models": [{ + "id": model_id, + "provider": provider_id, + "context_window": 1_000, + "max_output_tokens": 100 + }] + }); + call(&router, "router::provider::register", declaration.clone()) + .await + .expect_err("the selected state write must fail registration"); + + let providers = call(&router, "router::provider::list", json!({})) + .await + .expect("provider list after failure"); + assert!( + providers["providers"] + .as_array() + .is_some_and(|rows| rows.iter().all(|row| row["id"] != provider_id)), + "failed registration published a provider ghost: {providers}" + ); + let models = call( + &router, + "router::models::list", + json!({ "provider": provider_id }), + ) + .await + .expect("model list after failure"); + assert_eq!( + models["models"], + json!([]), + "failed registration published a catalog ghost: {models}" + ); + + // Rebuild both stores from the fake state's successful writes. In the + // fail-on-2 case this specifically proves the durable catalog candidate + // was compensated after the registry commit failed. + router.shutdown(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()) + .await + .expect("router reboots"); + let providers = call(&router, "router::provider::list", json!({})) + .await + .expect("provider list after restart"); + assert!( + providers["providers"] + .as_array() + .is_some_and(|rows| rows.iter().all(|row| row["id"] != provider_id)), + "failed registration survived restart as a provider ghost: {providers}" + ); + let models = call( + &router, + "router::models::list", + json!({ "provider": provider_id }), + ) + .await + .expect("model list after restart"); + assert_eq!( + models["models"], + json!([]), + "failed registration survived restart as a catalog ghost: {models}" + ); + + let retry = call(&router, "router::provider::register", declaration) + .await + .expect("retry without an undelivered token must bind cleanly"); + assert_eq!(retry["ok"], true, "{retry}"); + assert!( + retry["registration_token"] + .as_str() + .is_some_and(|token| !token.is_empty()), + "registration token missing after retry: {retry}" + ); + let models = call( + &router, + "router::models::list", + json!({ "provider": provider_id }), + ) + .await + .expect("model list after retry"); + assert_eq!(models["models"][0]["id"], model_id, "{models}"); + assert_eq!( + set_calls.load(Ordering::SeqCst), + expected_set_calls_after_retry + ); + + router.shutdown(); + state.shutdown(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn failed_static_catalog_write_keeps_provider_and_models_invisible() { + // catalog write fails before the registry is attempted; retry performs + // one catalog and one registry write. + assert_static_registration_failure_is_atomic(1, 3).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn failed_registry_after_static_catalog_rolls_back_before_retry() { + // catalog persists, registry fails, catalog rollback persists the prior + // snapshot, then retry performs catalog + registry writes. + assert_static_registration_failure_is_atomic(2, 5).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn late_failure_from_an_old_registration_cannot_mark_the_new_generation_down() { + let engine = engine_or_skip!(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + + let entered = Arc::new(tokio::sync::Semaphore::new(0)); + let release = Arc::new(tokio::sync::Semaphore::new(0)); + let provider = register_worker(&engine.url, InitOptions::default()); + let handler_entered = entered.clone(); + let handler_release = release.clone(); + provider.register_function( + "provider::generation-race::stream", + RegisterFunction::new_async(move |_input: Value| { + let entered = handler_entered.clone(); + let release = handler_release.clone(); + async move { + entered.add_permits(1); + release + .acquire() + .await + .expect("release semaphore stays open") + .forget(); + Err::(Error::Remote { + code: "function_not_found".into(), + message: "delayed failure from the old registration".into(), + stacktrace: None, + }) + } + }), + ); + + let declaration = json!({ + "id": "generation-race", + "models": [{ + "id": "generation-race-model", + "provider": "generation-race", + "context_window": 1_000, + "max_output_tokens": 100 + }] + }); + let first = call(&provider, "router::provider::register", declaration.clone()) + .await + .expect("first registration"); + let token = first["registration_token"] + .as_str() + .expect("registration token") + .to_string(); + + let consumer = register_worker(&engine.url, InitOptions::default()); + let (writer_ref, frames, pump) = consumer_channel(&consumer).await; + let chat = { + let consumer = consumer.clone(); + tokio::spawn(async move { + consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ + "writer_ref": writer_ref, + "model": "generation-race-model", + "messages": [] + }), + action: None, + timeout_ms: Some(10_000), + }) + .await + }) + }; + tokio::time::timeout(Duration::from_secs(3), entered.acquire()) + .await + .expect("old provider handler was not entered") + .expect("entered semaphore stays open") + .forget(); + + let mut redeclaration = declaration; + redeclaration["token"] = json!(token); + call(&provider, "router::provider::register", redeclaration) + .await + .expect("new generation registers before the old failure lands"); + release.add_permits(1); + + let error = tokio::time::timeout(Duration::from_secs(3), chat) + .await + .expect("chat resolves after releasing old handler") + .expect("chat task joins") + .expect_err("old dispatch still failed"); + assert_eq!( + remote_code(&error), + "router/provider_unavailable", + "{error:?}" + ); + tokio::time::timeout(Duration::from_secs(2), pump) + .await + .expect("consumer channel EOF") + .expect("consumer pump joins"); + let terminal_count = frames + .lock() + .unwrap() + .iter() + .filter_map(|frame| serde_json::from_str::(frame).ok()) + .filter(|frame| matches!(frame["type"].as_str(), Some("done" | "error"))) + .count(); + assert_eq!(terminal_count, 1, "frames: {:?}", frames.lock().unwrap()); + + let providers = call(&router, "router::provider::list", json!({})) + .await + .expect("provider list"); + let current = providers["providers"] + .as_array() + .and_then(|rows| rows.iter().find(|row| row["id"] == "generation-race")) + .expect("generation-race provider remains registered"); + assert_eq!( + current["available"], true, + "late old-generation failure marked the new registration down: {providers}" + ); + let route = call( + &router, + "router::route", + json!({ "model": "generation-race-model" }), + ) + .await + .expect("new registration remains routable"); + assert_eq!(route["provider"], "generation-race", "{route}"); + + consumer.shutdown(); + provider.shutdown(); + router.shutdown(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn router_coded_provider_failure_still_emits_one_terminal_and_eof() { + let engine = engine_or_skip!(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + + let provider = register_worker(&engine.url, InitOptions::default()); + provider.register_function( + "provider::coded-error::stream", + RegisterFunction::new_async(|_input: Value| async move { + Err::(Error::Remote { + code: "router/not_configured".into(), + message: "provider credentials are not configured".into(), + stacktrace: None, + }) + }), + ); + call( + &provider, + "router::provider::register", + json!({ + "id": "coded-error", + "models": [{ + "id": "coded-error-model", + "provider": "coded-error", + "context_window": 1_000, + "max_output_tokens": 100 + }] + }), + ) + .await + .expect("provider registers"); + + let consumer = register_worker(&engine.url, InitOptions::default()); + let (writer_ref, frames, pump) = consumer_channel(&consumer).await; + let error = consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ + "writer_ref": writer_ref, + "model": "coded-error-model", + "messages": [] + }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .expect_err("router-coded provider failure remains a typed bus error"); + assert_eq!(remote_code(&error), "router/not_configured", "{error:?}"); + tokio::time::timeout(Duration::from_secs(2), pump) + .await + .expect("caller channel reaches EOF") + .expect("channel pump joins"); + + let parsed: Vec = frames + .lock() + .unwrap() + .iter() + .map(|frame| serde_json::from_str(frame).expect("valid frame")) + .collect(); + let terminals: Vec<&Value> = parsed + .iter() + .filter(|frame| matches!(frame["type"].as_str(), Some("done" | "error"))) + .collect(); + assert_eq!(terminals.len(), 1, "frames: {parsed:?}"); + assert_eq!(terminals[0]["type"], "error", "frames: {parsed:?}"); + assert_eq!( + terminals[0]["error"]["error_kind"], "permanent", + "frames: {parsed:?}" + ); + + consumer.shutdown(); + provider.shutdown(); + router.shutdown(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn internal_channel_creation_failure_still_emits_one_terminal_and_eof() { + let engine = engine_or_skip!(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + let provider = start_live_provider(&engine.url, ProviderOptions::default()).await; + + // Mint the consumer channel before adversarially shadowing the engine's + // channel factory. The chat pipeline must turn the later internal factory + // failure into a terminal frame on this already-valid caller channel. + let consumer = register_worker(&engine.url, InitOptions::default()); + let (writer_ref, frames, pump) = consumer_channel(&consumer).await; + let saboteur = register_worker(&engine.url, InitOptions::default()); + saboteur.register_function( + "engine::channels::create", + RegisterFunction::new_async(|_input: Value| async move { + Err::(Error::Remote { + code: "injected/channel_unavailable".into(), + message: "injected channel factory failure".into(), + stacktrace: None, + }) + }), + ); + let deadline = Instant::now() + Duration::from_secs(3); + loop { + match call(&saboteur, "engine::channels::create", json!({})).await { + Err(error) if remote_code(&error) == "injected/channel_unavailable" => break, + _ => { + assert!( + Instant::now() < deadline, + "adversarial channel factory did not become active" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + } + } + } + + let error = consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ + "writer_ref": writer_ref, + "model": "live-1", + "messages": [] + }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .expect_err("internal channel creation failure remains a bus error"); + assert_eq!( + remote_code(&error), + "injected/channel_unavailable", + "{error:?}" + ); + tokio::time::timeout(Duration::from_secs(2), pump) + .await + .expect("caller channel reaches EOF") + .expect("channel pump joins"); + + let parsed: Vec = frames + .lock() + .unwrap() + .iter() + .map(|frame| serde_json::from_str(frame).expect("valid frame")) + .collect(); + let terminals: Vec<&Value> = parsed + .iter() + .filter(|frame| matches!(frame["type"].as_str(), Some("done" | "error"))) + .collect(); + assert_eq!(terminals.len(), 1, "frames: {parsed:?}"); + assert_eq!(terminals[0]["type"], "error", "frames: {parsed:?}"); + assert_eq!( + terminals[0]["error"]["error_kind"], "transient", + "frames: {parsed:?}" + ); + + saboteur.shutdown(); + consumer.shutdown(); + provider.iii.shutdown(); + router.shutdown(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn chat_provider_disappears_before_dispatch_emits_one_terminal_and_eof() { + let engine = engine_or_skip!(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + call( + &router, + "router::provider::register", + json!({ + "id": "ghost", + "models": [{ + "id": "ghost-model", + "provider": "ghost", + "context_window": 1000, + "max_output_tokens": 100 + }] + }), + ) + .await + .expect("declaration accepted"); + + let consumer = register_worker(&engine.url, InitOptions::default()); + let (writer_ref, frames, pump) = consumer_channel(&consumer).await; + let err = consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ + "writer_ref": writer_ref, + "model": "ghost-model", + "messages": [] + }), + action: None, + timeout_ms: Some(5_000), + }) + .await + .expect_err("missing provider function must be a typed error"); + assert_eq!(remote_code(&err), "router/provider_unavailable", "{err:?}"); + tokio::time::timeout(Duration::from_secs(2), pump) + .await + .expect("caller channel must reach EOF") + .expect("channel pump joins"); + + let parsed: Vec = frames + .lock() + .unwrap() + .iter() + .map(|frame| serde_json::from_str(frame).expect("valid frame")) + .collect(); + let terminals: Vec<&Value> = parsed + .iter() + .filter(|frame| matches!(frame["type"].as_str(), Some("done" | "error"))) + .collect(); + assert_eq!(terminals.len(), 1, "frames: {parsed:?}"); + assert_eq!(terminals[0]["type"], "error", "frames: {parsed:?}"); + assert_eq!( + terminals[0]["error"]["error_kind"], "transient", + "frames: {parsed:?}" + ); + + let listed = call(&router, "router::provider::list", json!({})) + .await + .expect("provider list"); + assert_eq!(listed["providers"][0]["available"], false, "{listed}"); + let route = call(&router, "router::route", json!({ "model": "ghost-model" })) + .await + .expect_err("an unavailable catalog owner must not remain routable"); + assert_eq!( + remote_code(&route), + "router/provider_unavailable", + "unexpected routing error: {route:?}" + ); + + consumer.shutdown(); + router.shutdown(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn terminal_frame_completes_even_when_the_provider_handler_lingers() { + let engine = engine_or_skip!(); + let router = register_worker(&engine.url, InitOptions::default()); + register_router(router.clone()).await.expect("router boots"); + let provider = start_live_provider( + &engine.url, + ProviderOptions { + done_linger_ms: Some(5_000), + ..Default::default() + }, + ) + .await; + call( + &router, + "configuration::set", + json!({ + "id": "llm-router", + "value": { + "settings": { + "stream_timeout_ms": 1_000, + "idle_timeout_ms": 500, + "retry_max": 0 + } + } + }), + ) + .await + .expect("settings update"); + tokio::time::sleep(Duration::from_millis(250)).await; + + let consumer = register_worker(&engine.url, InitOptions::default()); + let (writer_ref, frames, pump) = consumer_channel(&consumer).await; + let started = Instant::now(); + let response = consumer + .trigger(TriggerRequest { + function_id: "router::chat".into(), + payload: json!({ "writer_ref": writer_ref, "model": "live-1", "messages": [] }), + action: None, + timeout_ms: Some(2_000), + }) + .await + .expect("done frame completes chat"); + let elapsed = started.elapsed(); + assert_eq!(response["ok"], true, "{response}"); + assert!( + elapsed < Duration::from_millis(500), + "chat waited for the provider RPC after done: {elapsed:?}" + ); + tokio::time::timeout(Duration::from_secs(2), pump) + .await + .expect("caller channel EOF") + .expect("channel pump joins"); + let terminal_count = frames + .lock() + .unwrap() + .iter() + .filter_map(|frame| serde_json::from_str::(frame).ok()) + .filter(|frame| matches!(frame["type"].as_str(), Some("done" | "error"))) + .count(); + assert_eq!(terminal_count, 1, "frames: {:?}", frames.lock().unwrap()); + + consumer.shutdown(); + provider.iii.shutdown(); + router.shutdown(); +} + #[tokio::test(flavor = "multi_thread")] async fn complete_fails_fast_when_the_provider_worker_is_gone() { - // Regression for the boundary hang: a registered provider whose worker is - // gone makes the dispatch fail with function_not_found; `run` returns a - // typed Err without ever writing a frame, and `router::complete`'s drain - // must not block on a channel that will never EOF. + // Regression for the completion boundary: a registered provider whose + // worker is gone emits a terminal error and returns a typed Err. The + // internal channel drain must still race the pipeline and answer promptly. let engine = engine_or_skip!(); let router_iii = register_worker(&engine.url, InitOptions::default()); register_router(router_iii.clone()) From eafac6207fbb40415a9ac5a8f1d150d8e07df5ae Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Sat, 15 Aug 2026 00:03:23 -0300 Subject: [PATCH 2/5] (MOT-4412) test(harness): cover adversarial router failure --- harness/tests/integration/README.md | 1 + .../tests/integration/src/fixtures/loading.rs | 10 +- .../tests/integration/src/fixtures/tests.rs | 4 +- .../tests/integration/src/scenarios/dsl.rs | 193 +++++++++++++++++- .../tests/integration/src/scenarios/mod.rs | 4 +- .../router_midstream_terminal_error.rs | 148 ++++++++++++++ 6 files changed, 341 insertions(+), 19 deletions(-) create mode 100644 harness/tests/integration/src/scenarios/router_midstream_terminal_error.rs diff --git a/harness/tests/integration/README.md b/harness/tests/integration/README.md index f347e6237..7d151c330 100644 --- a/harness/tests/integration/README.md +++ b/harness/tests/integration/README.md @@ -28,6 +28,7 @@ No provider key or network access is required. | INT-018 | `spawn-reuse-guard` | direct | an in-turn spawn into an existing session owned by another parent is refused naming the owner (no hijack turn ever starts); re-spawning its own child appends the new task to the retained transcript and reports `reused: true` | | INT-019 | `condition-failure-notice` | direct | a binding whose condition ERRORS on a fire wakes its owner with an actionable `[notification]` (once per binding) instead of starving silently; the skip record still lands and the binding stays armed | | INT-020 | `child-discovery-granted` | direct | a child narrowed to its work functions can still dispatch the mandatory `engine::functions::list`/`::info` round (the discovery union); its native toolset stays the work functions only | +| INT-021 | `router-midstream-terminal-error` | direct | partial content and keepalive noise followed by one permanent router error preserve the partial, fail exactly once, and leave no pending work | | UI-001 | `console-streamed-text` | playground | a message sent by the Console streams to durable completion | | UI-002 | `multi-turn-traces` | playground | a native function turn and a Console turn expose distinct traces and function-call events | diff --git a/harness/tests/integration/src/fixtures/loading.rs b/harness/tests/integration/src/fixtures/loading.rs index 799fb2d16..42a3b8ed9 100644 --- a/harness/tests/integration/src/fixtures/loading.rs +++ b/harness/tests/integration/src/fixtures/loading.rs @@ -16,8 +16,8 @@ pub struct ScenarioFixture { /// first in the statuses list. pub expected_terminal_turns: usize, /// Each completion's lifecycle status, in completion order — parked - /// completions first, then terminal turns. The last must be `completed` — - /// the floor's durable-status check binds to it. + /// completions first, then terminal turns. The last status is also the + /// durable outcome that the floor requires from `harness::status`. pub expected_turn_statuses: Vec, pub scenario: CompiledScenarioV1, pub script: RouterScriptV1, @@ -155,12 +155,6 @@ impl ScenarioFixture { ); } } - if self.intervention.is_none() { - anyhow::ensure!( - self.expected_turn_statuses.last().map(String::as_str) == Some("completed"), - "the last terminal turn must be completed" - ); - } if let Some(intervention) = &self.intervention { match intervention { ScenarioIntervention::StopCancelCascade { diff --git a/harness/tests/integration/src/fixtures/tests.rs b/harness/tests/integration/src/fixtures/tests.rs index 523a706ea..7a8834317 100644 --- a/harness/tests/integration/src/fixtures/tests.rs +++ b/harness/tests/integration/src/fixtures/tests.rs @@ -12,7 +12,7 @@ fn all_selection_returns_the_checked_in_fixtures() { std::collections::BTreeSet::from([ "INT-001", "INT-002", "INT-003", "INT-005", "INT-006", "INT-010", "INT-011", "INT-012", "INT-013", "INT-014", "INT-015", "INT-016", "INT-017", "INT-018", "INT-019", "INT-020", - "UI-001", "UI-002" + "INT-021", "UI-001", "UI-002" ]) ); assert_eq!( @@ -20,7 +20,7 @@ fn all_selection_returns_the_checked_in_fixtures() { .iter() .filter(|fixture| fixture.driver == crate::scenarios::ScenarioDriver::Direct) .count(), - 16 + 17 ); } diff --git a/harness/tests/integration/src/scenarios/dsl.rs b/harness/tests/integration/src/scenarios/dsl.rs index 23c853271..f29fb6227 100644 --- a/harness/tests/integration/src/scenarios/dsl.rs +++ b/harness/tests/integration/src/scenarios/dsl.rs @@ -9,7 +9,7 @@ use serde_json::{json, Value}; use super::{ScenarioDriver, VerifyFn}; use crate::fixtures::{ScenarioFixture, ScenarioIntervention}; use crate::types::frames::{ - AssistantMessage, AssistantMessageEvent, AssistantRoleTag, ContentBlock, ErrorShape, + AssistantMessage, AssistantMessageEvent, AssistantRoleTag, ContentBlock, ErrorKind, ErrorShape, RouterChatResponse, StopReason, Usage, }; use crate::types::probe::ControlledTargetV1; @@ -239,10 +239,8 @@ impl Scenario { self } - /// Declare each terminal turn's status, in completion order, when not - /// every turn completes (e.g. a failed turn whose finalize drain reseeds a - /// completing follow-on turn). The last turn must complete: the floor's - /// durable-status check has no meaning for a run that ends failed. + /// Declare each terminal turn's status, in completion order, including a + /// run whose final durable outcome is failed or cancelled. pub(super) fn terminal_turn_statuses<'a>( mut self, statuses: impl IntoIterator, @@ -744,6 +742,12 @@ enum ResponseKind { FunctionCalls { calls: Vec<(String, String, Value)>, }, + TerminalError { + text: String, + chunks: Vec, + message: String, + kind: ErrorKind, + }, } impl Response { @@ -773,6 +777,31 @@ impl Response { } } + /// Stream a partial text response through keepalive noise, then terminate + /// with the router's authoritative error frame and failed RPC response. + pub(super) fn terminal_error_after_text( + text: &str, + chunks: I, + message: &str, + kind: ErrorKind, + input_tokens: u64, + output_tokens: u64, + ) -> Self + where + I: IntoIterator, + S: Into, + { + Self { + kind: ResponseKind::TerminalError { + text: text.to_string(), + chunks: chunks.into_iter().map(Into::into).collect(), + message: message.to_string(), + kind, + }, + usage: usage(input_tokens, output_tokens), + } + } + pub(super) fn function_call( call_id: &str, function: &ControlledFunction, @@ -837,10 +866,12 @@ impl Response { ) -> (Vec, RouterChatResponse) { let usage = self.usage; let timestamp = i64::try_from(ordinal).expect("generation ordinal fits i64"); - let (frames, stop_reason) = match self.kind { + let (frames, stop_reason, ok, error) = match self.kind { ResponseKind::StreamedText { text, chunks } => ( streamed_text_frames(&text, &chunks, &usage, model, timestamp), StopReason::End, + true, + None, ), ResponseKind::Text(text) => ( vec![AssistantMessageEvent::Done { @@ -853,6 +884,8 @@ impl Response { ), }], StopReason::End, + true, + None, ), ResponseKind::FunctionCall { call_id, @@ -873,6 +906,8 @@ impl Response { ), }], StopReason::FunctionCall, + true, + None, ), ResponseKind::FunctionCalls { calls } => ( vec![AssistantMessageEvent::Done { @@ -892,15 +927,31 @@ impl Response { ), }], StopReason::FunctionCall, + true, + None, + ), + ResponseKind::TerminalError { + text, + chunks, + message, + kind, + } => ( + streamed_error_frames(&text, &chunks, &message, kind, &usage, model, timestamp), + StopReason::Error, + false, + Some(ErrorShape { + code: error_kind_code(kind).to_string(), + message, + }), ), }; let response = RouterChatResponse { - ok: true, + ok, provider: model.provider.clone(), model: model.id.clone(), stop_reason: Some(stop_reason), usage: Some(usage), - error: None, + error, }; (frames, response) } @@ -1118,6 +1169,88 @@ fn streamed_text_frames( frames } +fn streamed_error_frames( + text: &str, + chunks: &[String], + message: &str, + kind: ErrorKind, + usage: &Usage, + model: &ModelFixtureV1, + timestamp: i64, +) -> Vec { + let mut error = assistant_message( + vec![ContentBlock::Text { + text: text.to_string(), + }], + StopReason::Error, + Some(usage.clone()), + model, + timestamp, + ); + error.error_message = Some(message.to_string()); + error.error_kind = Some(kind); + + let mut frames = vec![ + AssistantMessageEvent::Start { + partial: assistant_message(Vec::new(), StopReason::End, None, model, timestamp), + }, + AssistantMessageEvent::TextStart { + partial: assistant_message( + vec![ContentBlock::Text { + text: String::new(), + }], + StopReason::End, + None, + model, + timestamp, + ), + }, + ]; + for (index, delta) in chunks.iter().cloned().enumerate() { + frames.push(AssistantMessageEvent::TextDelta { + partial: None, + delta, + }); + if index + 1 < chunks.len() { + frames.push(AssistantMessageEvent::Ping); + } + } + frames.extend([ + AssistantMessageEvent::TextEnd { + partial: assistant_message( + vec![ContentBlock::Text { + text: text.to_string(), + }], + StopReason::End, + None, + model, + timestamp, + ), + }, + AssistantMessageEvent::Usage { + usage: usage.clone(), + }, + AssistantMessageEvent::Stop { + stop_reason: StopReason::Error, + error_message: Some(message.to_string()), + error_kind: Some(kind), + }, + AssistantMessageEvent::Ping, + AssistantMessageEvent::Error { error }, + ]); + frames +} + +fn error_kind_code(kind: ErrorKind) -> &'static str { + match kind { + ErrorKind::AuthExpired => "auth_expired", + ErrorKind::RateLimited => "rate_limited", + ErrorKind::ContextOverflow => "context_overflow", + ErrorKind::Transient => "transient", + ErrorKind::Permanent => "permanent", + } +} + fn system_prompt(allowed_functions: &[String]) -> String { let base = DEFAULT_SYSTEM_PROMPT .strip_suffix('\n') @@ -1152,6 +1285,50 @@ mod tests { assert_eq!(response.stop_reason, Some(StopReason::End)); } + #[test] + fn adversarial_terminal_error_keeps_partial_and_has_one_terminal_frame() { + let model = Model::scripted("fixture-model"); + let (frames, response) = Response::terminal_error_after_text( + "partial answer", + ["partial ", "answer"], + "provider disappeared after content", + ErrorKind::Permanent, + 5, + 2, + ) + .compile(&model, 1); + + assert_eq!(frames.iter().filter(|frame| frame.is_terminal()).count(), 1); + assert!( + frames + .iter() + .filter(|frame| matches!(frame, AssistantMessageEvent::Ping)) + .count() + >= 2 + ); + let Some(AssistantMessageEvent::Error { error }) = frames.last() else { + panic!("terminal frame must be error") + }; + assert_eq!(error.stop_reason, StopReason::Error); + assert_eq!(error.error_kind, Some(ErrorKind::Permanent)); + assert_eq!( + error.error_message.as_deref(), + Some("provider disappeared after content") + ); + assert_eq!( + error.content, + vec![ContentBlock::Text { + text: "partial answer".to_string() + }] + ); + assert!(!response.ok); + assert_eq!(response.stop_reason, Some(StopReason::Error)); + assert_eq!( + response.error.as_ref().map(|error| error.code.as_str()), + Some("permanent") + ); + } + #[test] fn controlled_function_uses_one_contract_for_tool_and_target() { let function = ControlledFunction::new("{{run_id}}::record", "Record value") diff --git a/harness/tests/integration/src/scenarios/mod.rs b/harness/tests/integration/src/scenarios/mod.rs index 7b942c200..b14b2cfe2 100644 --- a/harness/tests/integration/src/scenarios/mod.rs +++ b/harness/tests/integration/src/scenarios/mod.rs @@ -12,6 +12,7 @@ mod leaf_denied_control_plane; mod multi_turn_traces; mod queued_message_edit_unqueue; mod reseed_parked_message; +mod router_midstream_terminal_error; mod spawn_reuse_guard; mod standing_wake_delivery; mod state_worker_sidecar; @@ -47,6 +48,7 @@ pub fn all() -> Vec { standing_wake_delivery::scenario(), state_worker_sidecar::scenario(), reseed_parked_message::scenario(), + router_midstream_terminal_error::scenario(), spawn_reuse_guard::scenario(), stop_cancel_cascade::scenario(), queued_message_edit_unqueue::scenario(), @@ -63,7 +65,7 @@ mod tests { #[test] fn every_fixture_is_unique_and_valid() { let fixtures = all(); - assert_eq!(fixtures.len(), 18); + assert_eq!(fixtures.len(), 19); let mut slugs = std::collections::BTreeSet::new(); let mut ids = std::collections::BTreeSet::new(); for fixture in fixtures { diff --git a/harness/tests/integration/src/scenarios/router_midstream_terminal_error.rs b/harness/tests/integration/src/scenarios/router_midstream_terminal_error.rs new file mode 100644 index 000000000..98e5b4d7a --- /dev/null +++ b/harness/tests/integration/src/scenarios/router_midstream_terminal_error.rs @@ -0,0 +1,148 @@ +//! INT-021 — partial provider output ends in one authoritative router error. +//! +//! The scripted router sends content in multiple deltas with keepalive noise, +//! reports a non-terminal stop, then closes with one permanent error frame and +//! a failed RPC response. The Harness must preserve the useful partial, refuse +//! to resume a permanent failure, and reach one durable terminal failure with +//! no queued work or pending spans. + +use serde_json::Value; + +use super::dsl::{Generation, Message, Model, Request, Response, Scenario, Send}; +use super::ScenarioDriver; +use crate::fixtures::ScenarioFixture; +use crate::types::frames::ErrorKind; + +const ID: &str = "INT-021"; +const SLUG: &str = "router-midstream-terminal-error"; +const MESSAGE: &str = "Return an answer that exercises terminal failure handling."; +const PARTIAL: &str = "useful partial answer"; +const ERROR: &str = "provider disappeared after streaming content"; + +pub(super) fn scenario() -> ScenarioFixture { + Scenario::new( + ID, + SLUG, + "A permanent router error after partial output terminates the turn without losing the partial or leaving work pending.", + ScenarioDriver::Direct, + Model::scripted("fixture-model"), + ) + .send( + Send::message(MESSAGE) + .idempotency_key("{{run_id}}:integration-021") + .without_functions(), + ) + .terminal_turn_statuses(["failed"]) + .generation( + Generation::new(1) + .expect( + Request::new() + .turn_request() + .system_prompt_sha256("{{system_prompt_sha256}}") + .messages_exact([Message::user(MESSAGE)]) + .without_tools(), + ) + .respond(Response::terminal_error_after_text( + PARTIAL, + ["useful ", "partial ", "answer"], + ERROR, + ErrorKind::Permanent, + 12, + 3, + )), + ) + .verify(|run| { + run.expect_assistant_texts([PARTIAL])?; + run.expect_message_counts(1, 1, 0)?; + run.expect_no_duplicate_messages()?; + + anyhow::ensure!( + run.status.get("result_error").and_then(Value::as_str) == Some(ERROR), + "terminal error reason was not preserved: {}", + run.status + ); + anyhow::ensure!( + run.status + .get("partial_result_available") + .and_then(Value::as_bool) + == Some(true), + "failed turn did not retain its partial result: {}", + run.status + ); + anyhow::ensure!( + run.status.get("transient_resumes").and_then(Value::as_u64) == Some(0), + "permanent terminal failure must not resume: {}", + run.status + ); + anyhow::ensure!( + run.router_evidence + .pointer("/calls/0/outcome") + .and_then(Value::as_str) + == Some("matched") + && run + .router_evidence + .get("calls") + .and_then(Value::as_array) + .is_some_and(|calls| calls.len() == 1), + "adversarial generation was not served exactly once: {}", + run.router_evidence + ); + Ok(()) + }) + .scenario_timeout_ms(60_000) + .build() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::frames::{AssistantMessageEvent, ContentBlock, StopReason}; + + #[test] + fn fixture_streams_partial_keepalives_then_one_permanent_terminal() { + let fixture = scenario(); + fixture.validate().unwrap(); + assert_eq!(fixture.expected_terminal_turns, 1); + assert_eq!(fixture.expected_turn_statuses, ["failed"]); + + let generation = &fixture.script.generations[0]; + assert_eq!( + generation + .frames + .iter() + .filter(|frame| frame.is_terminal()) + .count(), + 1 + ); + assert!( + generation + .frames + .iter() + .filter(|frame| matches!(frame, AssistantMessageEvent::Ping)) + .count() + >= 2 + ); + let Some(AssistantMessageEvent::Error { error }) = generation.frames.last() else { + panic!("adversarial stream must end in an error frame") + }; + assert_eq!(error.stop_reason, StopReason::Error); + assert_eq!(error.error_kind, Some(ErrorKind::Permanent)); + assert_eq!(error.error_message.as_deref(), Some(ERROR)); + assert_eq!( + error.content, + vec![ContentBlock::Text { + text: PARTIAL.to_string() + }] + ); + assert!(!generation.response.ok); + assert_eq!(generation.response.stop_reason, Some(StopReason::Error)); + assert_eq!( + generation + .response + .error + .as_ref() + .map(|error| error.code.as_str()), + Some("permanent") + ); + } +} From a0d96f6e8526d519c14d3ecf96a14f9948fda0bf Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Sat, 15 Aug 2026 06:47:19 -0300 Subject: [PATCH 3/5] (MOT-4412) chore(ci): update cargo-audit installer --- .github/scripts/tests/test_rust_ci_workflows.py | 2 +- .github/workflows/rust-security-audit.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/scripts/tests/test_rust_ci_workflows.py b/.github/scripts/tests/test_rust_ci_workflows.py index c00025212..4ced61217 100644 --- a/.github/scripts/tests/test_rust_ci_workflows.py +++ b/.github/scripts/tests/test_rust_ci_workflows.py @@ -110,7 +110,7 @@ def test_rust_security_audit_is_narrow_on_prs_and_complete_on_schedule() -> None steps = audit["jobs"]["audit"]["steps"] install = named_step(steps, "Install cargo-audit") run = named_step(steps, "Audit Rust lockfiles")["run"] - assert install["uses"] == "taiki-e/install-action@v2.79.11" + assert install["uses"] == "taiki-e/install-action@v2.85.13" assert install["with"]["tool"] == "cargo-audit@0.22.2" assert "git diff --name-only -z" in run assert "find . -name Cargo.lock" in run diff --git a/.github/workflows/rust-security-audit.yml b/.github/workflows/rust-security-audit.yml index ed5cad130..e2669b53e 100644 --- a/.github/workflows/rust-security-audit.yml +++ b/.github/workflows/rust-security-audit.yml @@ -24,7 +24,7 @@ jobs: fetch-depth: 0 - name: Install cargo-audit - uses: taiki-e/install-action@v2.79.11 + uses: taiki-e/install-action@v2.85.13 with: tool: cargo-audit@0.22.2 fallback: none From a6e1a6e786beb127e4a48a280e30701733186594 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Sat, 15 Aug 2026 06:50:03 -0300 Subject: [PATCH 4/5] (MOT-4412) fix(deps): update quinn-proto --- llm-router/Cargo.lock | 74 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 61 insertions(+), 13 deletions(-) diff --git a/llm-router/Cargo.lock b/llm-router/Cargo.lock index 5aec7d8c9..8cc702e15 100644 --- a/llm-router/Cargo.lock +++ b/llm-router/Cargo.lock @@ -194,6 +194,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "clap" version = "4.6.1" @@ -280,6 +291,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crossbeam-deque" version = "0.8.7" @@ -613,11 +633,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", - "js-sys", "libc", "r-efi 5.3.0", "wasip2", - "wasm-bindgen", ] [[package]] @@ -627,10 +645,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", + "wasm-bindgen", ] [[package]] @@ -1216,7 +1237,7 @@ dependencies = [ "futures-util", "opentelemetry", "percent-encoding", - "rand", + "rand 0.9.4", "thiserror", "tokio", "tokio-stream", @@ -1305,14 +1326,15 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "bytes", - "getrandom 0.3.4", + "getrandom 0.4.2", "lru-slab", - "rand", + "rand 0.10.2", + "rand_pcg", "ring", "rustc-hash 2.1.2", "rustls", @@ -1366,7 +1388,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ "rand_chacha", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", ] [[package]] @@ -1376,7 +1409,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -1388,6 +1421,21 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + [[package]] name = "rayon" version = "1.12.0" @@ -1706,7 +1754,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -1717,7 +1765,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -1937,7 +1985,7 @@ dependencies = [ "macro_rules_attribute", "monostate", "paste", - "rand", + "rand 0.9.4", "rayon", "rayon-cond", "regex", @@ -2138,7 +2186,7 @@ dependencies = [ "http", "httparse", "log", - "rand", + "rand 0.9.4", "rustls", "rustls-pki-types", "sha1", From 8604b01f30c93d2e360f45a746b41fe067ffbe90 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Sat, 15 Aug 2026 06:54:31 -0300 Subject: [PATCH 5/5] (MOT-4412) fix(ci): handle Vercel link challenges --- .github/scripts/tests/test_check_links.py | 62 +++++++++++++++++++++++ scripts/check-links.sh | 9 +++- 2 files changed, 70 insertions(+), 1 deletion(-) create mode 100644 .github/scripts/tests/test_check_links.py diff --git a/.github/scripts/tests/test_check_links.py b/.github/scripts/tests/test_check_links.py new file mode 100644 index 000000000..a7233a3ff --- /dev/null +++ b/.github/scripts/tests/test_check_links.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[3] +CHECK_LINKS = ROOT / "scripts" / "check-links.sh" + + +def fake_curl(tmp_path: Path, mitigation: str) -> Path: + curl = tmp_path / "curl" + curl.write_text( + """#!/usr/bin/env bash +set -euo pipefail +headers_file= +while (( $# )); do + if [[ "$1" == "-D" ]]; then + headers_file=$2 + shift 2 + else + shift + fi +done +printf 'HTTP/2 403\\r\\nx-vercel-mitigated: %s\\r\\n\\r\\n' "$FAKE_VERCEL_MITIGATION" > "$headers_file" +printf '403' +""", + encoding="utf-8", + ) + curl.chmod(0o755) + return curl + + +def run_check(tmp_path: Path, mitigation: str) -> subprocess.CompletedProcess[str]: + fake_curl(tmp_path, mitigation) + env = os.environ.copy() + env["FAKE_VERCEL_MITIGATION"] = mitigation + env["PATH"] = f"{tmp_path}:{env['PATH']}" + return subprocess.run( + [str(CHECK_LINKS)], + cwd=ROOT, + env=env, + text=True, + capture_output=True, + timeout=30, + check=False, + ) + + +def test_vercel_security_challenge_is_reachable(tmp_path: Path) -> None: + result = run_check(tmp_path, "challenge") + + assert result.returncode == 0, result.stdout + result.stderr + assert "Vercel security challenge" in result.stdout + + +def test_vercel_deny_remains_a_failure(tmp_path: Path) -> None: + result = run_check(tmp_path, "deny") + + assert result.returncode == 1 + assert "FAIL 403" in result.stdout diff --git a/scripts/check-links.sh b/scripts/check-links.sh index d3b5d4048..ef5918e19 100755 --- a/scripts/check-links.sh +++ b/scripts/check-links.sh @@ -3,6 +3,9 @@ set -uo pipefail cd "$(dirname "$0")/.." || exit +headers_file=$(mktemp) +trap 'rm -f "$headers_file"' EXIT + # Known-good URLs that answer 404 to a plain GET; not broken doc links. ignore=( 'https://api.workers.iii.dev' # API base, only POST routes @@ -19,12 +22,16 @@ fail=0 for url in $urls; do for skip in "${ignore[@]}"; do [[ "$url" == "$skip" ]] && continue 2; done for attempt in 1 2 3; do - code=$(curl -sS -o /dev/null -w '%{http_code}' -L --max-time 20 "$url") + : > "$headers_file" + code=$(curl -sS -D "$headers_file" -o /dev/null -w '%{http_code}' -L --max-time 20 "$url") [[ "$code" =~ ^(429|5..|000)$ ]] || break (( attempt < 3 )) && sleep $((attempt * attempt * 2)) # backoff: 2s, 8s done if [[ "$code" =~ ^[23] ]]; then printf 'OK %s %s\n' "$code" "$url" + elif [[ "$code" == "403" ]] && + tr -d '\r' < "$headers_file" | grep -Eiq '^x-vercel-mitigated:[[:space:]]*challenge[[:space:]]*$'; then + printf 'OK %s %s (Vercel security challenge)\n' "$code" "$url" else printf 'FAIL %s %s\n' "$code" "$url" grep -rl --fixed-strings "${url%%#*}" . \