diff --git a/console/web/src/lib/backend/real.ts b/console/web/src/lib/backend/real.ts index 1ae70b1d5..e19f351a4 100644 --- a/console/web/src/lib/backend/real.ts +++ b/console/web/src/lib/backend/real.ts @@ -616,12 +616,20 @@ async function realCompactSession( } } - const resp = await client.trigger('context::compact', { - messages, - model: modelInput, - // Serialise concurrent compactions of the same conversation. - options: { lease_key: sessionId }, - }) + const resp = await client.trigger( + 'context::compact', + { + messages, + model: modelInput, + // Serialise concurrent compactions of the same conversation. + options: { lease_key: sessionId }, + }, + // Compaction makes a summariser LLM call budgeted up to 320s + // (context-manager `summarizer_timeout_ms`); the SDK's default 30s + // invocation timeout kills it mid-summary. Give the invocation + // headroom over the summariser's own budget. + { timeoutMs: 330_000 }, + ) if (resp?.status === 'ok') { const tailStartEntryId = diff --git a/console/web/src/lib/iii-client.ts b/console/web/src/lib/iii-client.ts index 11e27d5d1..1f4937400 100644 --- a/console/web/src/lib/iii-client.ts +++ b/console/web/src/lib/iii-client.ts @@ -35,10 +35,15 @@ export interface IiiClient { * Invoke an iii bus function and await its result. In the iii ecosystem * every bus invocation is a *trigger* — there is no separate "call" * concept. Thin wrapper over the SDK's `trigger({ function_id, payload })`. + * + * The SDK's default invocation timeout is 30s; pass `timeoutMs` for + * functions that legitimately run longer (e.g. `context::compact`, whose + * summariser call is budgeted up to 320s). */ trigger( functionId: string, payload?: Record, + options?: { timeoutMs?: number }, ): Promise on

( functionId: string, @@ -130,13 +135,17 @@ export function wrapSdk(sdk: ISdk, browserId: string): IiiClient { // unregister still releases the engine-side binding on dispose(). const triggerUnregisters = new Set<() => void>() + const DEFAULT_TRIGGER_TIMEOUT_MS = 5 * 60 * 1000 + function trigger( functionId: string, payload: Record = {}, + options?: { timeoutMs?: number }, ): Promise { return sdk.trigger({ function_id: functionId, payload, + timeoutMs: options?.timeoutMs ?? DEFAULT_TRIGGER_TIMEOUT_MS, }) } diff --git a/context-manager/src/adapters/cache.rs b/context-manager/src/adapters/cache.rs new file mode 100644 index 000000000..6bbee1f5f --- /dev/null +++ b/context-manager/src/adapters/cache.rs @@ -0,0 +1,277 @@ +//! TTL cache over a [`ModelResolver`]. Model budgets change only when a +//! provider reconciles its catalog slice or an operator edits limits, yet +//! the resolver was hitting `router::models::budget` on every non-inline +//! assemble/compact/count call. The decorator bounds that to one bus round +//! trip per (provider, id) per TTL window, and [`CachingModelResolver::flush`] +//! empties it eagerly when the router announces `router::models::changed`. +//! +//! TTLs: known budgets are safe to hold for a minute (staleness only skews +//! an estimate reserve); an unknown model (`Ok(None)`) is held briefly so a +//! just-reconciled catalog is picked up quickly; errors are never cached — +//! an absent router must not linger after it comes up. + +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +use async_trait::async_trait; +use iii_sdk::protocol::RegisterTriggerInput; +use iii_sdk::{IIIClient, RegisterFunction}; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +use crate::ports::{Clock, ModelBudget, ModelResolver}; + +const MODELS_CHANGED_FN_ID: &str = "context::on-models-changed"; + +const HIT_TTL_MS: i64 = 60_000; +const NEGATIVE_TTL_MS: i64 = 5_000; + +struct CacheEntry { + budget: Option, + stored_at_ms: i64, +} + +pub struct CachingModelResolver { + inner: Arc, + clock: Arc, + entries: RwLock, String), CacheEntry>>, +} + +impl CachingModelResolver { + pub fn new(inner: Arc, clock: Arc) -> Self { + Self { + inner, + clock, + entries: RwLock::new(HashMap::new()), + } + } + + /// Drop every cached budget (bound to `router::models::changed`). + pub fn flush(&self) { + self.entries + .write() + .expect("model-budget cache lock poisoned") + .clear(); + } + + fn fresh(&self, key: &(Option, String)) -> Option> { + let entries = self + .entries + .read() + .expect("model-budget cache lock poisoned"); + let entry = entries.get(key)?; + let ttl = if entry.budget.is_some() { + HIT_TTL_MS + } else { + NEGATIVE_TTL_MS + }; + (self.clock.now_ms().saturating_sub(entry.stored_at_ms) <= ttl) + .then(|| entry.budget.clone()) + } +} + +#[async_trait] +impl ModelResolver for CachingModelResolver { + async fn get_model_budget( + &self, + provider: Option<&str>, + id: &str, + ) -> Result, String> { + let key = (provider.map(str::to_string), id.to_string()); + if let Some(budget) = self.fresh(&key) { + return Ok(budget); + } + // The lock is never held across this await; concurrent misses may + // duplicate one resolve, which is harmless and self-heals on store. + let budget = self.inner.get_model_budget(provider, id).await?; + self.entries + .write() + .expect("model-budget cache lock poisoned") + .insert( + key, + CacheEntry { + budget: budget.clone(), + stored_at_ms: self.clock.now_ms(), + }, + ); + Ok(budget) + } +} + +/// `router::models::changed` payload — `{ provider, count }`. The handler +/// flushes everything regardless, so the fields are advisory only. +#[derive(Debug, Default, Deserialize, schemars::JsonSchema)] +struct OnModelsChangedEvent { + #[serde(default)] + #[allow(dead_code)] + provider: Option, +} + +/// Ack returned by the internal `context::on-models-changed` handler. +#[derive(Debug, Serialize, schemars::JsonSchema)] +struct OnModelsChangedResponse { + ok: bool, +} + +/// Best-effort flush binding: when the router announces a catalog change +/// (`router::models::changed`), drop every cached budget so the next call +/// re-resolves. A failed registration (router absent at boot, older router +/// without the trigger type) only costs TTL-bounded staleness — never boot. +pub fn register_models_changed_flush(iii: &IIIClient, cache: Arc) { + iii.register_function( + MODELS_CHANGED_FN_ID, + RegisterFunction::new_async(move |_event: OnModelsChangedEvent| { + let cache = cache.clone(); + async move { + cache.flush(); + Ok::(OnModelsChangedResponse { + ok: true, + }) + } + }) + .description("Internal: flush the model-budget cache when the router's catalog changes.") + .metadata(json!({ "internal": true, "trace_hidden": true })), + ); + + if let Err(e) = iii.register_trigger(RegisterTriggerInput { + trigger_type: "router::models::changed".to_string(), + function_id: MODELS_CHANGED_FN_ID.to_string(), + config: json!({}), + metadata: None, + }) { + tracing::warn!( + error = %e, + "could not bind router::models::changed; model-budget cache degrades to TTL-only invalidation" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::Model; + use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering}; + + struct FakeClock(AtomicI64); + impl Clock for FakeClock { + fn now_ms(&self) -> i64 { + self.0.load(Ordering::SeqCst) + } + } + + struct CountingResolver { + calls: AtomicUsize, + response: Result, String>, + } + + #[async_trait] + impl ModelResolver for CountingResolver { + async fn get_model_budget( + &self, + _provider: Option<&str>, + _id: &str, + ) -> Result, String> { + self.calls.fetch_add(1, Ordering::SeqCst); + self.response.clone() + } + } + + fn budget() -> ModelBudget { + let model: Model = serde_json::from_value(serde_json::json!({ + "id": "m", + "provider": "p", + "context_window": 100_000, + "max_output_tokens": 8_000, + })) + .unwrap(); + ModelBudget { + effective_max_output_tokens: model.max_output_tokens, + model, + } + } + + fn harness( + response: Result, String>, + ) -> (Arc, Arc, CachingModelResolver) { + let inner = Arc::new(CountingResolver { + calls: AtomicUsize::new(0), + response, + }); + let clock = Arc::new(FakeClock(AtomicI64::new(0))); + let cache = CachingModelResolver::new(inner.clone(), clock.clone()); + (inner, clock, cache) + } + + #[tokio::test] + async fn hit_within_ttl_avoids_the_inner_call() { + let (inner, clock, cache) = harness(Ok(Some(budget()))); + assert!(cache + .get_model_budget(Some("p"), "m") + .await + .unwrap() + .is_some()); + clock.0.store(HIT_TTL_MS, Ordering::SeqCst); // exactly at the TTL edge + assert!(cache + .get_model_budget(Some("p"), "m") + .await + .unwrap() + .is_some()); + assert_eq!(inner.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn expiry_re_resolves() { + let (inner, clock, cache) = harness(Ok(Some(budget()))); + cache.get_model_budget(Some("p"), "m").await.unwrap(); + clock.0.store(HIT_TTL_MS + 1, Ordering::SeqCst); + cache.get_model_budget(Some("p"), "m").await.unwrap(); + assert_eq!(inner.calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn keys_are_per_provider_and_id() { + let (inner, _clock, cache) = harness(Ok(Some(budget()))); + cache.get_model_budget(Some("p"), "m").await.unwrap(); + cache.get_model_budget(Some("q"), "m").await.unwrap(); + cache.get_model_budget(None, "m").await.unwrap(); + assert_eq!(inner.calls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn unknown_model_uses_the_short_negative_ttl() { + let (inner, clock, cache) = harness(Ok(None)); + assert!(cache + .get_model_budget(Some("p"), "m") + .await + .unwrap() + .is_none()); + clock.0.store(NEGATIVE_TTL_MS + 1, Ordering::SeqCst); + assert!(cache + .get_model_budget(Some("p"), "m") + .await + .unwrap() + .is_none()); + assert_eq!( + inner.calls.load(Ordering::SeqCst), + 2, + "a just-reconciled model must be seen within seconds" + ); + } + + #[tokio::test] + async fn errors_are_never_cached() { + let (inner, _clock, cache) = harness(Err("router unreachable".into())); + assert!(cache.get_model_budget(Some("p"), "m").await.is_err()); + assert!(cache.get_model_budget(Some("p"), "m").await.is_err()); + assert_eq!(inner.calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn flush_clears_every_entry() { + let (inner, _clock, cache) = harness(Ok(Some(budget()))); + cache.get_model_budget(Some("p"), "m").await.unwrap(); + cache.flush(); + cache.get_model_budget(Some("p"), "m").await.unwrap(); + assert_eq!(inner.calls.load(Ordering::SeqCst), 2); + } +} diff --git a/context-manager/src/adapters/mod.rs b/context-manager/src/adapters/mod.rs index 166ea80b8..8dd822b76 100644 --- a/context-manager/src/adapters/mod.rs +++ b/context-manager/src/adapters/mod.rs @@ -1,5 +1,6 @@ //! Production adapters behind the ports: `llm-router` calls over the //! iii bus and filesystem-backed lease storage. +pub mod cache; pub mod fs_lease; pub mod router; diff --git a/context-manager/src/core/prune.rs b/context-manager/src/core/prune.rs index 927d57eb5..61a625324 100644 --- a/context-manager/src/core/prune.rs +++ b/context-manager/src/core/prune.rs @@ -78,6 +78,28 @@ pub fn prune( messages: &mut [AgentMessage], params: &PruneParams, estimator: &dyn Estimator, +) -> PruneStats { + prune_impl(messages, None, params, estimator) +} + +/// [`prune`] that also keeps a caller-held per-message size memo in +/// sync: `sizes[i]` is re-estimated for every rewritten message, so the +/// memo stays exactly what a from-scratch recount would produce. +pub fn prune_with_sizes( + messages: &mut [AgentMessage], + sizes: &mut [u64], + params: &PruneParams, + estimator: &dyn Estimator, +) -> PruneStats { + debug_assert_eq!(messages.len(), sizes.len()); + prune_impl(messages, Some(sizes), params, estimator) +} + +fn prune_impl( + messages: &mut [AgentMessage], + mut sizes: Option<&mut [u64]>, + params: &PruneParams, + estimator: &dyn Estimator, ) -> PruneStats { let mut scanned: u64 = 0; let mut window_tokens: u64 = 0; @@ -138,6 +160,11 @@ pub fn prune( messages[*idx].set_content(vec![ContentBlock::Text { text: placeholder(*tokens), }]); + if let Some(sizes) = sizes.as_deref_mut() { + // Re-estimate the rewritten message — not freed-token + // arithmetic — so the memo matches a from-scratch recount. + sizes[*idx] = estimator.message(&messages[*idx]); + } } PruneStats { @@ -159,6 +186,27 @@ pub fn emergency_reduce( messages: &mut [AgentMessage], required_tokens: u64, estimator: &dyn Estimator, +) -> PruneStats { + emergency_reduce_impl(messages, None, required_tokens, estimator) +} + +/// [`emergency_reduce`] that keeps a caller-held per-message size memo +/// in sync (see [`prune_with_sizes`]). +pub fn emergency_reduce_with_sizes( + messages: &mut [AgentMessage], + sizes: &mut [u64], + required_tokens: u64, + estimator: &dyn Estimator, +) -> PruneStats { + debug_assert_eq!(messages.len(), sizes.len()); + emergency_reduce_impl(messages, Some(sizes), required_tokens, estimator) +} + +fn emergency_reduce_impl( + messages: &mut [AgentMessage], + mut sizes: Option<&mut [u64]>, + required_tokens: u64, + estimator: &dyn Estimator, ) -> PruneStats { let mut candidates: Vec<(usize, u64)> = messages .iter() @@ -193,6 +241,9 @@ pub fn emergency_reduce( } messages[idx] = replacement; + if let Some(sizes) = sizes.as_deref_mut() { + sizes[idx] = replacement_tokens; + } stats.pruned_tokens = stats.pruned_tokens.saturating_add(freed); stats.pruned_parts += 1; } @@ -522,6 +573,47 @@ mod tests { } } + fn sizes_of(messages: &[AgentMessage]) -> Vec { + messages + .iter() + .map(|m| HeuristicEstimator.message(m)) + .collect() + } + + #[test] + fn prune_with_sizes_keeps_the_memo_equal_to_a_recount() { + let mut messages = history(); + let mut sizes = sizes_of(&messages); + let stats = prune_with_sizes(&mut messages, &mut sizes, ¶ms(), &HeuristicEstimator); + assert_eq!(stats.pruned_parts, 1); + assert_eq!(sizes, sizes_of(&messages)); + } + + #[test] + fn emergency_reduce_with_sizes_keeps_the_memo_equal_to_a_recount() { + let mut messages = vec![ + result("small", 4_000, 1), + result("largest", 40_000, 2), + result("middle", 20_000, 3), + ]; + let mut sizes = sizes_of(&messages); + let stats = + emergency_reduce_with_sizes(&mut messages, &mut sizes, 5_000, &HeuristicEstimator); + assert!(stats.pruned_parts >= 1); + assert_eq!(sizes, sizes_of(&messages)); + } + + #[test] + fn with_sizes_variants_match_the_plain_passes() { + let mut plain = history(); + let mut memoed = history(); + let mut sizes = sizes_of(&memoed); + let plain_stats = prune(&mut plain, ¶ms(), &HeuristicEstimator); + let memo_stats = prune_with_sizes(&mut memoed, &mut sizes, ¶ms(), &HeuristicEstimator); + assert_eq!(plain_stats, memo_stats); + assert_eq!(plain, memoed); + } + #[test] fn emergency_reduces_largest_results_only_as_needed() { let mut messages = vec![ diff --git a/context-manager/src/core/selection.rs b/context-manager/src/core/selection.rs index 80c6d15bc..2a3c6b0f8 100644 --- a/context-manager/src/core/selection.rs +++ b/context-manager/src/core/selection.rs @@ -8,7 +8,6 @@ //! `function_result` message, and never a user message carrying inline //! `function_result` blocks). -use crate::core::estimate::Estimator; use crate::types::{AgentMessage, Role}; /// One conversation turn: starts at a user message, ends right before @@ -60,12 +59,9 @@ fn is_safe_cut(message: &AgentMessage) -> bool { /// Find a partial tail inside `turn` that fits `budget`, scanning from /// the oldest in-turn position forward; only safe cuts qualify. -fn split_turn( - messages: &[AgentMessage], - turn: Turn, - budget: u64, - estimator: &dyn Estimator, -) -> Option { +/// `prefix[i]` is the token sum of `messages[..i]`, so a range costs +/// one subtraction instead of a re-serializing scan per candidate cut. +fn split_turn(messages: &[AgentMessage], prefix: &[u64], turn: Turn, budget: u64) -> Option { if budget == 0 || turn.end.saturating_sub(turn.start) <= 1 { return None; } @@ -73,10 +69,7 @@ fn split_turn( if !is_safe_cut(&messages[start]) { continue; } - let size: u64 = messages[start..turn.end] - .iter() - .map(|m| estimator.message(m)) - .sum(); + let size = prefix[turn.end] - prefix[start]; if size > budget { continue; } @@ -89,12 +82,17 @@ fn split_turn( /// turns that fit `budget` (newest first); when a whole turn does not /// fit, fall back to a safe partial cut inside it. Everything before /// the kept tail is the head to summarise. +/// +/// `sizes[i]` must be the estimated tokens of `messages[i]` — callers +/// hold this memo anyway, and passing it keeps selection free of any +/// per-candidate re-estimation. pub fn select( messages: &[AgentMessage], + sizes: &[u64], budget: u64, tail_turns: usize, - estimator: &dyn Estimator, ) -> Selection { + debug_assert_eq!(messages.len(), sizes.len()); let whole_head = Selection { head_len: messages.len(), tail_start_index: None, @@ -108,13 +106,16 @@ pub fn select( } let recent = &all[all.len().saturating_sub(tail_turns)..]; + let mut prefix: Vec = Vec::with_capacity(sizes.len() + 1); + prefix.push(0); + for size in sizes { + prefix.push(prefix[prefix.len() - 1] + size); + } + let mut total: u64 = 0; let mut keep: Option = None; for turn in recent.iter().rev() { - let size: u64 = messages[turn.start..turn.end] - .iter() - .map(|m| estimator.message(m)) - .sum(); + let size = prefix[turn.end] - prefix[turn.start]; if total + size <= budget { total += size; // A turn whose user message carries inline function_result @@ -127,7 +128,7 @@ pub fn select( continue; } let remaining = budget.saturating_sub(total); - if let Some(split) = split_turn(messages, *turn, remaining, estimator) { + if let Some(split) = split_turn(messages, &prefix, *turn, remaining) { keep = Some(split); } break; @@ -167,9 +168,19 @@ pub fn select( #[cfg(test)] mod tests { use super::*; - use crate::core::estimate::HeuristicEstimator; + use crate::core::estimate::{Estimator, HeuristicEstimator}; use serde_json::json; + /// Test shim preserving the old call shape: estimate every message + /// with the heuristic, exactly as production callers fill the memo. + fn select_est(messages: &[AgentMessage], budget: u64, tail_turns: usize) -> Selection { + let sizes: Vec = messages + .iter() + .map(|m| HeuristicEstimator.message(m)) + .collect(); + select(messages, &sizes, budget, tail_turns) + } + fn user(text: &str, ts: i64) -> AgentMessage { serde_json::from_value(json!({ "role": "user", "content": [{ "type": "text", "text": text }], "timestamp": ts @@ -226,7 +237,7 @@ mod tests { user("recent question", 3), assistant("recent answer", 4), ]; - let sel = select(&messages, 1_000, 1, &HeuristicEstimator); + let sel = select_est(&messages, 1_000, 1); assert_eq!(sel.tail_start_index, Some(2)); assert_eq!(sel.head_len, 2); } @@ -234,7 +245,7 @@ mod tests { #[test] fn zero_tail_turns_summarises_everything() { let messages = vec![user("a", 1), assistant("b", 2)]; - let sel = select(&messages, 1_000, 0, &HeuristicEstimator); + let sel = select_est(&messages, 1_000, 0); assert_eq!(sel.tail_start_index, None); assert_eq!(sel.head_len, 2); } @@ -246,7 +257,7 @@ mod tests { // (head_len 0 → the caller skips compaction) instead of summarising // the whole history into an empty model context. let messages = vec![user("a", 1), assistant("b", 2)]; - let sel = select(&messages, u64::MAX, 2, &HeuristicEstimator); + let sel = select_est(&messages, u64::MAX, 2); assert_eq!(sel.head_len, 0); assert_eq!(sel.tail_start_index, Some(0)); } @@ -261,7 +272,7 @@ mod tests { user("recent question", 3), assistant("recent answer", 4), ]; - let sel = select(&messages, u64::MAX, 2, &HeuristicEstimator); + let sel = select_est(&messages, u64::MAX, 2); assert_eq!(sel.head_len, 0); assert_eq!(sel.tail_start_index, Some(0)); } @@ -277,7 +288,7 @@ mod tests { result("c1", 4_000, 3), assistant("done", 4), ]; - let sel = select(&messages, 100, 1, &HeuristicEstimator); + let sel = select_est(&messages, 100, 1); // Tail = just the final assistant message (index 3): the cut at // index 2 (the function_result) is unsafe even though it fits. assert_eq!(sel.tail_start_index, Some(3)); @@ -317,7 +328,7 @@ mod tests { inline_result_user, // 4: turn 3 (UNSAFE start) assistant("done", 5), // 5 ]; - let sel = select(&messages, 1_000, 2, &HeuristicEstimator); + let sel = select_est(&messages, 1_000, 2); assert_eq!( sel.tail_start_index, Some(2), @@ -336,7 +347,7 @@ mod tests { user("q2", 3), assistant("a2", 4), ]; - let sel = select(&messages, 0, 2, &HeuristicEstimator); + let sel = select_est(&messages, 0, 2); assert_eq!(sel.tail_start_index, Some(2)); assert_eq!(sel.head_len, 2); } @@ -354,7 +365,7 @@ mod tests { result("c1", 4_000, 5), assistant("done", 6), ]; - let sel = select(&messages, 0, 2, &HeuristicEstimator); + let sel = select_est(&messages, 0, 2); assert_eq!(sel.tail_start_index, Some(2)); assert!(sel.head_len > 0 && sel.head_len < messages.len()); } @@ -362,7 +373,7 @@ mod tests { #[test] fn history_without_user_messages_is_all_head() { let messages = vec![assistant("a", 1), assistant("b", 2)]; - let sel = select(&messages, 1_000, 2, &HeuristicEstimator); + let sel = select_est(&messages, 1_000, 2); assert_eq!(sel.tail_start_index, None); } } diff --git a/context-manager/src/functions/assemble.rs b/context-manager/src/functions/assemble.rs index e37390e07..258d58ebf 100644 --- a/context-manager/src/functions/assemble.rs +++ b/context-manager/src/functions/assemble.rs @@ -16,7 +16,7 @@ use serde::{Deserialize, Serialize}; use crate::core::budget::{default_reserved, preserve_recent_budget, usable}; use crate::core::estimate::{estimator_for_model, Estimator}; use crate::core::lease; -use crate::core::prune::{emergency_reduce, prune as run_prune, PruneParams}; +use crate::core::prune::{emergency_reduce_with_sizes, prune_with_sizes, PruneParams}; use crate::core::selection::select; use crate::core::summary::{ build_system_prompt, render_system_prompt, render_user_prompt, strip_media, @@ -129,6 +129,19 @@ pub struct AssembleResponse { pub applied: Applied, } +/// Test-only re-export of [`count_context`] so sibling function tests +/// can pin cross-function counting equivalence (see count_tokens.rs). +#[cfg(test)] +pub(crate) fn count_context_for_tests( + messages: &[AgentMessage], + prompt: &str, + tools: &[AgentFunction], + request_overhead_tokens: u64, + estimator: &dyn Estimator, +) -> u64 { + count_context(messages, prompt, tools, request_overhead_tokens, estimator) +} + fn count_context( messages: &[AgentMessage], prompt: &str, @@ -181,11 +194,26 @@ pub async fn handle(deps: &Deps, req: AssembleRequest) -> Result u64 { - count_context(messages, prompt, tools, request_overhead_tokens, estimator) + // Size memo: every message, tool, and the prompt are estimated once; + // the pipeline's recounts become O(1)-per-message sums, and the + // mutating passes below keep `sizes` in lockstep with `working`. + // The fold order and saturating ops mirror `count_context` exactly + // so totals stay byte-identical with a from-scratch recount. + let mut sizes: Vec = working.iter().map(|m| estimator.message(m)).collect(); + let tool_tokens = tools.iter().fold(0u64, |total, tool| { + total.saturating_add(estimator.function(tool)) + }); + let mut prompt_tokens = estimator.text(&system_prompt); + let total = |sizes: &[u64], prompt_tokens: u64| -> u64 { + sizes + .iter() + .fold(0u64, |total, size| total.saturating_add(*size)) + .saturating_add(prompt_tokens) + .saturating_add(tool_tokens) + .saturating_add(request_overhead_tokens) }; - let initial_token_count = count(&working, &system_prompt); + let initial_token_count = total(&sizes, prompt_tokens); let mut applied = Applied { initial_token_count, pruned: false, @@ -207,10 +235,10 @@ pub async fn handle(deps: &Deps, req: AssembleRequest) -> Result 0; applied.pruned_tokens = stats.pruned_tokens; - token_count = count(&working, &system_prompt); + token_count = total(&sizes, prompt_tokens); } // Step 2: compact the head. @@ -227,22 +255,24 @@ pub async fn handle(deps: &Deps, req: AssembleRequest) -> Result Result usable_budget { - let emergency = emergency_reduce( + let emergency = emergency_reduce_with_sizes( &mut working, + &mut sizes, token_count.saturating_sub(usable_budget), estimator, ); @@ -259,9 +290,21 @@ pub async fn handle(deps: &Deps, req: AssembleRequest) -> Result usable_budget { return Err(ContextError::Overflow { token_count, @@ -299,11 +342,11 @@ async fn try_compact( deps: &Deps, model: &ModelInput, working: &[AgentMessage], + sizes: &[u64], usable_budget: u64, tail_turns: usize, lease_key: &str, previous_summary: Option<&str>, - estimator: &dyn Estimator, ) -> Option { let config = deps.config().await; let leases = deps.leases().await; @@ -312,13 +355,13 @@ async fn try_compact( let outcome = async { let budget = preserve_recent_budget(usable_budget, None); - let selection = select(working, budget, tail_turns, estimator); + let selection = select(working, sizes, budget, tail_turns); let head = &working[..selection.head_len]; if head.is_empty() { return None; } - let tokens_before: u64 = head.iter().map(|m| estimator.message(m)).sum(); + let tokens_before: u64 = sizes[..selection.head_len].iter().sum(); let stripped = strip_media(head, config.max_output_chars); let request = SummarizeRequest { system_prompt: build_system_prompt(previous_summary), diff --git a/context-manager/src/functions/compact.rs b/context-manager/src/functions/compact.rs index cd608b6f1..56f480168 100644 --- a/context-manager/src/functions/compact.rs +++ b/context-manager/src/functions/compact.rs @@ -131,14 +131,16 @@ async fn summarise( previous_summary: Option<&str>, estimator: &dyn Estimator, ) -> CompactResponse { - let selection = select(messages, budget, tail_turns, estimator); + // One estimate per message; select, tokens_before, and tokens_after + // all read this memo instead of re-serializing per pass. + let sizes: Vec = messages.iter().map(|m| estimator.message(m)).collect(); + let selection = select(messages, &sizes, budget, tail_turns); let head = &messages[..selection.head_len]; if head.is_empty() { return CompactResponse::Empty; } - let tail = &messages[selection.head_len..]; - let tokens_before: u64 = head.iter().map(|m| estimator.message(m)).sum(); + let tokens_before: u64 = sizes[..selection.head_len].iter().sum(); let stripped = strip_media(head, deps.config().await.max_output_chars); let request = SummarizeRequest { @@ -166,8 +168,7 @@ async fn summarise( } }; - let tokens_after = - estimator.text(&summary) + tail.iter().map(|m| estimator.message(m)).sum::(); + let tokens_after = estimator.text(&summary) + sizes[selection.head_len..].iter().sum::(); CompactResponse::Ok { summary, diff --git a/context-manager/src/functions/count_tokens.rs b/context-manager/src/functions/count_tokens.rs index 7da3b9fa8..ebdb87d48 100644 --- a/context-manager/src/functions/count_tokens.rs +++ b/context-manager/src/functions/count_tokens.rs @@ -90,3 +90,57 @@ pub async fn handle( }, }) } + +#[cfg(test)] +mod tests { + use crate::core::estimate::{estimate_messages, Estimator, HeuristicEstimator}; + use crate::types::{AgentFunction, AgentMessage}; + use serde_json::json; + + /// Callers (the harness) substitute `context::assemble.token_count` + /// for `count-tokens(...).tokens + overhead` when nothing mutated the + /// request after assembly. That substitution is only valid while the + /// two functions count the same inputs to the same number — pin it. + #[test] + fn count_tokens_and_assemble_count_the_same_context_identically() { + let messages: Vec = vec![ + serde_json::from_value(json!({ + "role": "user", + "content": [{ "type": "text", "text": "question" }], + "timestamp": 1 + })) + .unwrap(), + serde_json::from_value(json!({ + "role": "assistant", + "content": [{ "type": "text", "text": "answer" }], + "stop_reason": "end", "model": "m", "provider": "p", "timestamp": 2 + })) + .unwrap(), + ]; + let tools: Vec = vec![serde_json::from_value(json!({ + "name": "agent_trigger", + "description": "Run an agent function", + "parameters": { "type": "object" } + })) + .unwrap()]; + let prompt = "system prompt"; + let overhead = 37u64; + let estimator = HeuristicEstimator; + + // count-tokens' arithmetic (handle() above), plus the overhead the + // harness adds on top. + let mut count_tokens_total = estimate_messages(&estimator, &messages); + count_tokens_total += estimator.text(prompt); + for tool in &tools { + count_tokens_total += estimator.function(tool); + } + count_tokens_total += overhead; + + // assemble's arithmetic (its overhead-inclusive count). + let assemble_total = crate::functions::assemble::count_context_for_tests( + &messages, prompt, &tools, overhead, &estimator, + ); + + assert_eq!(assemble_total, count_tokens_total); + } +} diff --git a/context-manager/src/main.rs b/context-manager/src/main.rs index dbbc02447..e856d1da5 100644 --- a/context-manager/src/main.rs +++ b/context-manager/src/main.rs @@ -28,6 +28,7 @@ use iii_sdk::runtime::WorkerMetadata; use iii_sdk::{register_worker, InitOptions}; use tokio::sync::RwLock; +use context_manager::adapters::cache::CachingModelResolver; use context_manager::adapters::fs_lease::FsLeaseStore; use context_manager::adapters::router::{RouterModelResolver, RouterSummarizer}; use context_manager::configuration::{self, ConfigCell}; @@ -135,15 +136,24 @@ async fn main() -> Result<()> { .map_err(|e| anyhow::anyhow!("opening the compaction lease directory: {e}"))?, )); + // Budget lookups are TTL-cached and flushed on router::models::changed; + // the decorator drops into Deps.resolver with no handler changes. + let clock: Arc = Arc::new(SystemClock); + let resolver = Arc::new(CachingModelResolver::new( + Arc::new(RouterModelResolver::new(iii.clone())), + clock.clone(), + )); + let deps = Arc::new(Deps { config: cell.clone(), - resolver: Arc::new(RouterModelResolver::new(iii.clone())), + resolver: resolver.clone(), summarizer, leases: leases.clone(), - clock: Arc::new(SystemClock), + clock, }); functions::register_all(&iii, &deps); + context_manager::adapters::cache::register_models_changed_flush(&iii, resolver); // LAST: bind the configuration-change trigger so its handler closes over // the snapshot cell + the lease cell it rebuilds on a lease_dir change. diff --git a/harness/src/clients/router.rs b/harness/src/clients/router.rs index 1b1f788ab..b7728b222 100644 --- a/harness/src/clients/router.rs +++ b/harness/src/clients/router.rs @@ -173,6 +173,9 @@ impl RouterClient { // `_streaming` so UIs can show the command being formed. let mut args_acc: std::collections::HashMap = std::collections::HashMap::new(); + // Cumulative message across slim delta frames (boundary snapshot + + // open-block deltas); fat frames just refresh the snapshot. + let mut tracker = PartialTracker::default(); // Consume frames until reader EOF — but ALSO watch the held-open // trigger: when it resolves, the router is done (ack) or the dispatch @@ -239,20 +242,28 @@ impl RouterClient { other => { if let AssistantMessageEvent::FunctioncallDelta { partial, delta, id } = &other { + // Pre-id producers: guess from the frame's own fat + // partial or, on slim frames, from the last boundary + // snapshot (FunctioncallStart carries the call block). let id = if id.is_empty() { - open_call_id(partial) // pre-id producers: guess + partial + .as_ref() + .or(tracker.base.as_ref()) + .and_then(open_call_id) + .map(str::to_string) } else { - Some(id.as_str()) + Some(id.clone()) }; if let Some(id) = id { - args_acc.entry(id.to_string()).or_default().push_str(delta); + args_acc.entry(id).or_default().push_str(delta); } } - if let Some(partial) = partial_of(&other) { - if last_emit.elapsed() >= coalesce { - match enrich_streaming_args(partial, &args_acc) { + tracker.apply(&other); + if other.is_content() && last_emit.elapsed() >= coalesce { + if let Some(cum) = tracker.current() { + match enrich_streaming_args(&cum, &args_acc) { Some(enriched) => sink.on_update(&enriched).await, - None => sink.on_update(partial).await, + None => sink.on_update(&cum).await, } last_emit = Instant::now(); } @@ -453,6 +464,140 @@ impl StreamSink for CapturingSink { } } +/// Cumulative-message view over the streamed frames: the last +/// block-boundary snapshot plus the locally accumulated deltas of the open +/// text/thinking block. Slim delta frames carry no `partial`, so this +/// tracker is the only live view of the message between block boundaries; +/// legacy fat deltas (partial present) simply replace the snapshot, which +/// preserves the old behavior byte for byte. In-flight function-call +/// arguments intentionally stay in `args_acc` (`enrich_streaming_args`). +#[derive(Default)] +struct PartialTracker { + base: Option, + open: Option, +} + +enum OpenBlock { + Text { seed: String, acc: String }, + Thinking { seed: String, acc: String }, +} + +impl PartialTracker { + fn apply(&mut self, event: &AssistantMessageEvent) { + use crate::types::content::ContentBlock; + match event { + // Block boundaries carry authoritative snapshots (thinking + // signatures and finalized call args exist only here). + AssistantMessageEvent::Start { partial } + | AssistantMessageEvent::TextEnd { partial } + | AssistantMessageEvent::ThinkingEnd { partial } + | AssistantMessageEvent::FunctioncallStart { partial } + | AssistantMessageEvent::FunctioncallEnd { partial } => { + self.base = Some(partial.clone()); + self.open = None; + } + // Producers may or may not include the just-opened (empty) + // block in the Start snapshot — pop a trailing match into the + // seed so the merge never duplicates it. + AssistantMessageEvent::TextStart { partial } => { + let mut base = partial.clone(); + let seed = match base.content.last() { + Some(ContentBlock::Text { text }) => { + let text = text.clone(); + base.content.pop(); + text + } + _ => String::new(), + }; + self.base = Some(base); + self.open = Some(OpenBlock::Text { + seed, + acc: String::new(), + }); + } + AssistantMessageEvent::ThinkingStart { partial } => { + let mut base = partial.clone(); + let seed = match base.content.last() { + Some(ContentBlock::Thinking { text, .. }) => { + let text = text.clone(); + base.content.pop(); + text + } + _ => String::new(), + }; + self.base = Some(base); + self.open = Some(OpenBlock::Thinking { + seed, + acc: String::new(), + }); + } + AssistantMessageEvent::TextDelta { partial, delta } => match partial { + Some(p) => { + self.base = Some(p.clone()); + self.open = None; + } + None => match &mut self.open { + Some(OpenBlock::Text { acc, .. }) => acc.push_str(delta), + // Defensive: producer skipped TextStart. + _ => { + self.open = Some(OpenBlock::Text { + seed: String::new(), + acc: delta.clone(), + }) + } + }, + }, + AssistantMessageEvent::ThinkingDelta { partial, delta } => match partial { + Some(p) => { + self.base = Some(p.clone()); + self.open = None; + } + None => match &mut self.open { + Some(OpenBlock::Thinking { acc, .. }) => acc.push_str(delta), + _ => { + self.open = Some(OpenBlock::Thinking { + seed: String::new(), + acc: delta.clone(), + }) + } + }, + }, + // Raw args accumulate in args_acc; the call block itself is + // already in the FunctioncallStart snapshot. A fat frame + // still refreshes the snapshot (legacy behavior). + AssistantMessageEvent::FunctioncallDelta { + partial: Some(p), .. + } => { + self.base = Some(p.clone()); + self.open = None; + } + AssistantMessageEvent::FunctioncallDelta { partial: None, .. } => {} + _ => {} + } + } + + /// The cumulative message: snapshot + the open block's accumulated text. + fn current(&self) -> Option { + use crate::types::content::ContentBlock; + let mut out = self.base.clone()?; + match &self.open { + Some(OpenBlock::Text { seed, acc }) if !(seed.is_empty() && acc.is_empty()) => { + out.content.push(ContentBlock::Text { + text: format!("{seed}{acc}"), + }); + } + Some(OpenBlock::Thinking { seed, acc }) if !(seed.is_empty() && acc.is_empty()) => { + out.content.push(ContentBlock::Thinking { + text: format!("{seed}{acc}"), + signature: None, + }); + } + _ => {} + } + Some(out) + } +} + /// Fallback attribution for id-less delta frames (pre-id producers): blocks /// stream in order, so guess the last function_call block of the partial. fn open_call_id(partial: &AssistantMessage) -> Option<&str> { @@ -512,22 +657,6 @@ fn utf8_tail(s: &str, max: usize) -> &str { &s[start..] } -fn partial_of(event: &AssistantMessageEvent) -> Option<&AssistantMessage> { - match event { - AssistantMessageEvent::Start { partial } - | AssistantMessageEvent::TextStart { partial } - | AssistantMessageEvent::TextDelta { partial, .. } - | AssistantMessageEvent::TextEnd { partial } - | AssistantMessageEvent::ThinkingStart { partial } - | AssistantMessageEvent::ThinkingDelta { partial, .. } - | AssistantMessageEvent::ThinkingEnd { partial } - | AssistantMessageEvent::FunctioncallStart { partial } - | AssistantMessageEvent::FunctioncallDelta { partial, .. } - | AssistantMessageEvent::FunctioncallEnd { partial } => Some(partial), - _ => None, - } -} - #[cfg(test)] mod streaming_args_tests { use super::*; @@ -606,5 +735,122 @@ mod streaming_args_tests { AssistantMessageEvent::FunctioncallDelta { id, .. } => assert_eq!(id, "c1"), other => panic!("want functioncall_delta, got {other:?}"), } + // Slim frames (no partial at all) must parse too — they are the + // post-contract-change producer shape. + let slim = r#"{"type":"functioncall_delta","delta":"x","id":"c1"}"#; + match serde_json::from_str::(slim).unwrap() { + AssistantMessageEvent::FunctioncallDelta { partial, id, .. } => { + assert!(partial.is_none()); + assert_eq!(id, "c1"); + } + other => panic!("want functioncall_delta, got {other:?}"), + } + let slim_text = r#"{"type":"text_delta","delta":"hi"}"#; + assert!(matches!( + serde_json::from_str::(slim_text).unwrap(), + AssistantMessageEvent::TextDelta { partial: None, .. } + )); + } + + #[test] + fn tracker_accumulates_slim_text_deltas_from_the_boundary_snapshot() { + use crate::types::content::ContentBlock; + let mut tracker = PartialTracker::default(); + let base = empty_assistant("p", "m"); + tracker.apply(&AssistantMessageEvent::Start { + partial: base.clone(), + }); + tracker.apply(&AssistantMessageEvent::TextStart { + partial: base.clone(), + }); + tracker.apply(&AssistantMessageEvent::TextDelta { + partial: None, + delta: "Hel".into(), + }); + tracker.apply(&AssistantMessageEvent::TextDelta { + partial: None, + delta: "lo".into(), + }); + let cum = tracker.current().expect("cumulative message"); + assert_eq!( + cum.content, + vec![ContentBlock::Text { + text: "Hello".into() + }] + ); + // The TextEnd snapshot is authoritative and replaces the open block. + let mut done = base.clone(); + done.content = vec![ContentBlock::Text { + text: "Hello!".into(), + }]; + tracker.apply(&AssistantMessageEvent::TextEnd { + partial: done.clone(), + }); + assert_eq!(tracker.current().unwrap().content, done.content); + } + + #[test] + fn tracker_treats_fat_deltas_as_authoritative_snapshots() { + use crate::types::content::ContentBlock; + let mut tracker = PartialTracker::default(); + let mut fat = empty_assistant("p", "m"); + fat.content = vec![ContentBlock::Text { + text: "cumulative".into(), + }]; + tracker.apply(&AssistantMessageEvent::TextDelta { + partial: Some(fat.clone()), + delta: "e".into(), + }); + // Legacy behavior: the fat partial IS the message; deltas are not + // re-applied on top of it. + assert_eq!(tracker.current().unwrap().content, fat.content); + } + + #[test] + fn tracker_keeps_thinking_and_call_blocks_across_boundaries() { + use crate::types::content::ContentBlock; + use serde_json::json; + let mut tracker = PartialTracker::default(); + let base = empty_assistant("p", "m"); + tracker.apply(&AssistantMessageEvent::ThinkingStart { + partial: base.clone(), + }); + tracker.apply(&AssistantMessageEvent::ThinkingDelta { + partial: None, + delta: "hmm".into(), + }); + assert_eq!( + tracker.current().unwrap().content, + vec![ContentBlock::Thinking { + text: "hmm".into(), + signature: None + }] + ); + // FunctioncallStart snapshot carries the call block; slim call-arg + // deltas leave the tracker alone (args live in args_acc) so the + // current view still shows the call block for open_call_id. + let mut with_call = base.clone(); + with_call.content = vec![ + ContentBlock::Thinking { + text: "hmm".into(), + signature: Some("sig".into()), + }, + ContentBlock::FunctionCall { + id: "c1".into(), + function_id: "agent_trigger".into(), + arguments: json!({}), + }, + ]; + tracker.apply(&AssistantMessageEvent::FunctioncallStart { + partial: with_call.clone(), + }); + tracker.apply(&AssistantMessageEvent::FunctioncallDelta { + partial: None, + delta: r#"{"x":1"#.into(), + id: "c1".into(), + }); + let cum = tracker.current().unwrap(); + assert_eq!(cum.content, with_call.content); + assert_eq!(open_call_id(&cum), Some("c1")); } } diff --git a/harness/src/clients/session.rs b/harness/src/clients/session.rs index fc88f0c5a..0a74689b7 100644 --- a/harness/src/clients/session.rs +++ b/harness/src/clients/session.rs @@ -262,4 +262,57 @@ impl SessionClient { } Ok(out) } + + /// Load only the entries strictly after `after_entry_id` on the active + /// path — the incremental form of [`Self::messages`] for callers holding + /// a watermark. `Ok(None)` means the watermark left the active path + /// (fork / `set_active_leaf`); the caller must fall back to a full load. + pub async fn messages_after( + &self, + session_id: &str, + after_entry_id: &str, + include_custom: bool, + ) -> Result>, HarnessError> { + const PAGE_LIMIT: u64 = 500; + let mut out: Vec = Vec::new(); + let mut cursor: Option = None; + loop { + let mut payload = json!({ + "session_id": session_id, + "include_custom": include_custom, + "limit": PAGE_LIMIT, + }); + match &cursor { + // Later pages resume from the server cursor; only the first + // page anchors on the caller's watermark. + Some(c) => payload["cursor"] = json!(c), + None => payload["after_entry_id"] = json!(after_entry_id), + } + let resp = match self.call("session::messages", payload).await { + Ok(resp) => resp, + Err(HarnessError::Dependency(msg)) if msg.contains("session/invalid_cursor") => { + return Ok(None); + } + Err(e) => return Err(e), + }; + let arr = resp + .get("messages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + for item in arr { + match serde_json::from_value::(item) { + Ok(entry) => out.push(entry), + Err(e) => { + tracing::warn!(session_id, error = %e, "skipping unparseable session entry") + } + } + } + match resp.get("next_cursor").and_then(Value::as_str) { + Some(next) if !next.is_empty() => cursor = Some(next.to_string()), + _ => break, + } + } + Ok(Some(out)) + } } diff --git a/harness/src/turn_loop.rs b/harness/src/turn_loop.rs index d1107f9e9..2f92050e2 100644 --- a/harness/src/turn_loop.rs +++ b/harness/src/turn_loop.rs @@ -51,6 +51,15 @@ const CONTEXT_OVERFLOW_FAILURE: FailureInfo = FailureInfo { /// Keep a small fixed reserve in addition to the deterministic JSON estimate. const PROVIDER_FRAMING_ALLOWANCE_TOKENS: u64 = 64; +/// Assembly headroom reserved when pre-generate hooks are bound: hook +/// appends land AFTER assembly fits the context, so without a reservation +/// any append overflows an exactly-full context and fails the turn. +const PRE_GENERATE_HOOK_ALLOWANCE_TOKENS: u64 = 256; + +/// Extra margin folded into the one-shot re-assembly after a post-assembly +/// overflow, covering hook-output variance on the retry. +const REASSEMBLY_HEADROOM_MARGIN_TOKENS: u64 = 256; + fn estimate_request_overhead_tokens( response_format: Option<&Value>, provider_options: Option<&Value>, @@ -290,134 +299,189 @@ pub async fn run_step( // Assemble the model-ready context (+ compaction persistence). An // impossible fit is a terminal turn outcome, not an unexpected step error. - let assembled = match assemble_context( - deps, - &session, - &record, - &entries, - payload.step, - prev_watermark.as_deref(), - ContextAssemblyInputs { - system_prompt: assembly_system_prompt, - tools: &tools, - request_overhead_tokens, - }, - ) - .await - { - Ok(assembled) => assembled, - Err(HarnessError::ContextOverflow(reason)) => { - return finalize_failed( - deps, - &session, - &mut record, - &reason, - CONTEXT_OVERFLOW_FAILURE, - ) - .await; - } - Err(error) => return Err(error), + // + // Post-assembly, pre_generate hooks and orphan repair may GROW the + // request, while assembly fits to the exact usable ceiling — so on a + // full context even a tiny hook append overflows the final check + // (observed live: a ~20-token guidance injection over a 220k-usable + // request killed a 13-step turn). Two defenses, composable: when + // pre-generate hooks are bound, reserve a small allowance up front; + // and if the final count still overflows, re-assemble ONCE with the + // measured deficit folded into the overhead so assembly makes room — + // failing the turn only when even that is not enough. + let mut extra_overhead_tokens: u64 = if deps.hooks.pre_generate.is_empty() { + 0 + } else { + PRE_GENERATE_HOOK_ALLOWANCE_TOKENS }; - - // pre_generate hooks: extend the system prompt / append bounded messages, - // or veto. Annotations ride the assistant entry's origin (audit trail). - let (gen_system_prompt, appended, gen_annotations) = match deps - .hooks - .run_pre_generate( + let mut reassembled = false; + let (gen_system_prompt, gen_annotations, gen_messages) = loop { + let assembled = match assemble_context( + deps, + &session, &record, + &entries, payload.step, - assembled.system_prompt.clone(), - &assembled.messages, + prev_watermark.as_deref(), + ContextAssemblyInputs { + system_prompt: assembly_system_prompt.clone(), + tools: &tools, + request_overhead_tokens: request_overhead_tokens + .saturating_add(extra_overhead_tokens), + }, ) .await - { - crate::hooks::runner::PreGenerateOutcome::Continue { - system_prompt, - append_messages, - annotations, - } => (system_prompt, append_messages, annotations), - crate::hooks::runner::PreGenerateOutcome::Deny(reason) => { + { + Ok(assembled) => assembled, + Err(HarnessError::ContextOverflow(reason)) => { + return finalize_failed( + deps, + &session, + &mut record, + &reason, + CONTEXT_OVERFLOW_FAILURE, + ) + .await; + } + Err(error) => return Err(error), + }; + + // pre_generate hooks: extend the system prompt / append bounded messages, + // or veto. Annotations ride the assistant entry's origin (audit trail). + let (gen_system_prompt, appended, gen_annotations) = match deps + .hooks + .run_pre_generate( + &record, + payload.step, + assembled.system_prompt.clone(), + &assembled.messages, + ) + .await + { + crate::hooks::runner::PreGenerateOutcome::Continue { + system_prompt, + append_messages, + annotations, + } => (system_prompt, append_messages, annotations), + crate::hooks::runner::PreGenerateOutcome::Deny(reason) => { + return finalize_failed( + deps, + &session, + &mut record, + &format!("pre_generate hook denied: {reason}"), + FailureInfo { + code: "harness.pre_generate_denied", + phase: "pre_generate", + retryable: false, + }, + ) + .await; + } + }; + let hook_appended = !appended.is_empty(); + let mut gen_messages = assembled.messages.clone(); + gen_messages.extend(appended); + + // Post-assembly invariant guard: providers reject a context where an + // assistant function_call has no function_result. Compaction can cut a + // pair (result summarized away, call kept) even when the TRANSCRIPT is + // fully paired — patch the assembled copy only. + let patched = patch_orphaned_calls(&mut gen_messages); + if patched > 0 { + tracing::warn!( + session_id = %record.session_id, + turn_id = %record.turn_id, + patched, + "assembled context contained orphaned function_calls; injected elided results (compaction cut a call/result pair)" + ); + } + + // Never hand the provider an empty messages array (Anthropic 400: + // "messages: at least one message is required"). Assembly's own guards + // make this unreachable in practice; if it still happens (e.g. a + // transcript with no user message at all), fail the turn with a clear + // harness error instead of emitting a cryptic provider error. + if gen_messages.is_empty() { return finalize_failed( deps, &session, &mut record, - &format!("pre_generate hook denied: {reason}"), + "assembled context is empty; refusing to call the provider with no messages", FailureInfo { - code: "harness.pre_generate_denied", - phase: "pre_generate", + code: "harness.empty_context", + phase: "context_assembly", retryable: false, }, ) .await; } - }; - let mut gen_messages = assembled.messages.clone(); - gen_messages.extend(appended); - - // Post-assembly invariant guard: providers reject a context where an - // assistant function_call has no function_result. Compaction can cut a - // pair (result summarized away, call kept) even when the TRANSCRIPT is - // fully paired — patch the assembled copy only. - let patched = patch_orphaned_calls(&mut gen_messages); - if patched > 0 { + + // Hooks and orphan repair can change the assembled request. When nothing + // did — no appended messages, no orphan patches, prompt unchanged — the + // request IS the assemble output, whose `token_count` already covers + // messages + prompt + tools + the request overhead, so the second + // full-payload count-tokens round trip is pure repetition. Re-count only + // when the request actually diverged. (The re-count compares against the + // BASE overhead: the extra reservation exists only to make assembly + // leave room; it is not part of the real request.) + let request_unchanged = final_request_unchanged( + hook_appended, + patched, + &gen_system_prompt, + &assembled.system_prompt, + ); + let final_request_tokens = if request_unchanged { + assembled.token_count + } else { + let final_count = deps + .context() + .await + .count_tokens(crate::clients::context::CountTokensParams { + messages: gen_messages.clone(), + model_id: record.options.model.clone(), + provider: record.options.provider.clone(), + system_prompt: gen_system_prompt.clone(), + tools: tools.clone(), + }) + .await + .map_err(HarnessError::Dependency)?; + final_count.tokens.saturating_add(request_overhead_tokens) + }; + if final_request_tokens <= assembled.usable { + break (gen_system_prompt, gen_annotations, gen_messages); + } + if reassembled { + let reason = format!( + "final model request exceeds the assembled usable budget: {final_request_tokens} tokens > {} usable (persists after re-assembly with reserved headroom)", + assembled.usable + ); + return finalize_failed( + deps, + &session, + &mut record, + &reason, + CONTEXT_OVERFLOW_FAILURE, + ) + .await; + } + // One-shot recovery: fold the measured overshoot (plus margin for + // hook variance on the retry — hooks re-run against the smaller + // context) into the reservation and re-assemble. The compaction + // bookkeeping entry id is per (turn, step), so a re-assembled + // compaction dedupes instead of double-writing. + reassembled = true; + let deficit = final_request_tokens - assembled.usable; + extra_overhead_tokens = extra_overhead_tokens + .saturating_add(deficit) + .saturating_add(REASSEMBLY_HEADROOM_MARGIN_TOKENS); tracing::warn!( session_id = %record.session_id, turn_id = %record.turn_id, - patched, - "assembled context contained orphaned function_calls; injected elided results (compaction cut a call/result pair)" + deficit, + extra_overhead_tokens, + "post-assembly additions exceeded the usable budget; re-assembling with reserved headroom" ); - } - - // Never hand the provider an empty messages array (Anthropic 400: - // "messages: at least one message is required"). Assembly's own guards - // make this unreachable in practice; if it still happens (e.g. a - // transcript with no user message at all), fail the turn with a clear - // harness error instead of emitting a cryptic provider error. - if gen_messages.is_empty() { - return finalize_failed( - deps, - &session, - &mut record, - "assembled context is empty; refusing to call the provider with no messages", - FailureInfo { - code: "harness.empty_context", - phase: "context_assembly", - retryable: false, - }, - ) - .await; - } - - // Hooks and orphan repair can change the assembled request. Re-count the - // exact final prompt/messages/tools and include the same non-context - // overhead reserve before creating an assistant entry or calling router. - let final_count = deps - .context() - .await - .count_tokens(crate::clients::context::CountTokensParams { - messages: gen_messages.clone(), - model_id: record.options.model.clone(), - provider: record.options.provider.clone(), - system_prompt: gen_system_prompt.clone(), - tools: tools.clone(), - }) - .await - .map_err(HarnessError::Dependency)?; - let final_request_tokens = final_count.tokens.saturating_add(request_overhead_tokens); - if final_request_tokens > assembled.usable { - let reason = format!( - "final model request exceeds the assembled usable budget: {final_request_tokens} tokens > {} usable", - assembled.usable - ); - return finalize_failed( - deps, - &session, - &mut record, - &reason, - CONTEXT_OVERFLOW_FAILURE, - ) - .await; - } + }; let assistant_origin = origin_with(&record.turn_id, &gen_annotations); @@ -1509,7 +1573,17 @@ async fn has_user_after_watermark( }; // include_custom must match the watermark's source list (the step entry // loads with `true`): a watermark landing on a custom entry would - // otherwise never be found and the steering check silently dies. + // otherwise be filtered off the path and the delta fetch would error. + // The incremental fetch returns only post-watermark entries; when the + // watermark left the active path (fork), fall back to the full scan. + if let Some(suffix) = session + .messages_after(&record.session_id, watermark, true) + .await? + { + return Ok(suffix + .iter() + .any(|entry| matches!(&entry.message, Some(AgentMessage::User(_))))); + } let entries = session.messages(&record.session_id, true).await?; let mut after = false; for entry in entries { @@ -1669,10 +1743,13 @@ async fn assemble_context( .map(|m| serde_json::to_value(m).unwrap_or(Value::Null)) .collect(); rotate_mid_generation_users(&mut messages, new_suffix_len); + // Rotation only reorders; the estimator is an order-independent + // per-message sum, so assemble's count still describes this list. Ok(Assembled { system_prompt: Some(out.system_prompt), messages, usable: out.usable, + token_count: out.token_count, }) } @@ -1745,6 +1822,10 @@ struct Assembled { system_prompt: Option, messages: Vec, usable: u64, + /// `context::assemble`'s estimate of this exact context (messages + + /// prompt + tools + request overhead). Reused as the final request + /// count when nothing mutates the request after assembly. + token_count: u64, } struct ContextAssemblyInputs<'a> { @@ -1804,6 +1885,19 @@ fn rotate_mid_generation_users(messages: &mut Vec, new_suffix_len: usize) /// message directly after each orphaned call's assistant message. Returns how /// many results were injected. The durable transcript is never touched; the /// orphan usually means compaction cut a call/result pair. +/// Whether the final model request is exactly the assemble output — no +/// hook-appended messages, no orphan patches, prompt unchanged — so +/// `context::assemble.token_count` already describes it and the second +/// count-tokens round trip would be pure repetition. +fn final_request_unchanged( + hook_appended: bool, + patched: usize, + gen_system_prompt: &Option, + assembled_system_prompt: &Option, +) -> bool { + !hook_appended && patched == 0 && gen_system_prompt == assembled_system_prompt +} + fn patch_orphaned_calls(messages: &mut Vec) -> usize { let mut resolved: std::collections::HashSet = std::collections::HashSet::new(); for m in messages.iter() { @@ -1996,6 +2090,24 @@ mod tests { ); } + #[test] + fn final_count_is_skipped_only_when_nothing_mutated_the_request() { + let prompt = Some("base".to_string()); + // Untouched request: assemble's token_count is authoritative. + assert!(super::final_request_unchanged(false, 0, &prompt, &prompt)); + // Any mutation forces the re-count: hook append, orphan patch, + // or a hook-rewritten system prompt. + assert!(!super::final_request_unchanged(true, 0, &prompt, &prompt)); + assert!(!super::final_request_unchanged(false, 1, &prompt, &prompt)); + assert!(!super::final_request_unchanged( + false, + 0, + &Some("hooked".to_string()), + &prompt + )); + assert!(!super::final_request_unchanged(false, 0, &None, &prompt)); + } + #[test] fn context_overflow_classification_requires_the_stable_context_code() { assert!(super::is_context_overflow_error( diff --git a/harness/src/types/event.rs b/harness/src/types/event.rs index 09c9180b7..f379ff37a 100644 --- a/harness/src/types/event.rs +++ b/harness/src/types/event.rs @@ -62,7 +62,10 @@ pub enum AssistantMessageEvent { partial: AssistantMessage, }, TextDelta { - partial: AssistantMessage, + /// Legacy fat-frame snapshot; slim producers omit it and readers + /// accumulate `delta`s from the last block-boundary snapshot. + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, delta: String, }, TextEnd { @@ -72,7 +75,8 @@ pub enum AssistantMessageEvent { partial: AssistantMessage, }, ThinkingDelta { - partial: AssistantMessage, + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, delta: String, }, ThinkingEnd { @@ -82,7 +86,8 @@ pub enum AssistantMessageEvent { partial: AssistantMessage, }, FunctioncallDelta { - partial: AssistantMessage, + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, delta: String, /// Call id receiving this delta; empty from pre-id producers. #[serde(default, skip_serializing_if = "String::is_empty")] @@ -117,4 +122,21 @@ impl AssistantMessageEvent { AssistantMessageEvent::Done { .. } | AssistantMessageEvent::Error { .. } ) } + + /// The ten content-bearing variants: block boundaries and deltas. + pub fn is_content(&self) -> bool { + matches!( + self, + AssistantMessageEvent::Start { .. } + | AssistantMessageEvent::TextStart { .. } + | AssistantMessageEvent::TextDelta { .. } + | AssistantMessageEvent::TextEnd { .. } + | AssistantMessageEvent::ThinkingStart { .. } + | AssistantMessageEvent::ThinkingDelta { .. } + | AssistantMessageEvent::ThinkingEnd { .. } + | AssistantMessageEvent::FunctioncallStart { .. } + | AssistantMessageEvent::FunctioncallDelta { .. } + | AssistantMessageEvent::FunctioncallEnd { .. } + ) + } } diff --git a/llm-router/src/chat/accumulate.rs b/llm-router/src/chat/accumulate.rs new file mode 100644 index 000000000..d6a776b78 --- /dev/null +++ b/llm-router/src/chat/accumulate.rs @@ -0,0 +1,340 @@ +//! Cumulative-message reconstruction over the streaming vocabulary. +//! +//! Delta frames carry only their `delta` (a per-chunk cumulative snapshot +//! made streams O(N²) — see `types::events`); block-boundary frames carry +//! authoritative snapshots. This accumulator folds the two back into "the +//! message so far": the last boundary snapshot plus the accumulated deltas +//! of the open block. The relay feeds abort / no-terminal synthesis from +//! it, so a provider dying mid-block still yields every delta received. +//! +//! Legacy fat deltas (`partial: Some`) replace the snapshot wholesale — +//! byte-identical to the old behavior, so old producers interop. + +use crate::types::content::ContentBlock; +use crate::types::events::AssistantMessageEvent; +use crate::types::messages::{degraded_arguments, AssistantMessage}; + +#[derive(Default)] +pub struct PartialAccumulator { + /// Last block-boundary snapshot (authoritative: signatures, final args). + base: Option, + /// Deltas since that snapshot. + open: Option, +} + +enum OpenBlock { + Text { seed: String, acc: String }, + Thinking { seed: String, acc: String }, + Call { args: String }, +} + +impl PartialAccumulator { + pub fn apply(&mut self, event: &AssistantMessageEvent) { + match event { + AssistantMessageEvent::Start { partial } + | AssistantMessageEvent::TextEnd { partial } + | AssistantMessageEvent::ThinkingEnd { partial } + | AssistantMessageEvent::FunctioncallEnd { partial } => { + self.base = Some(partial.clone()); + self.open = None; + } + // Producers differ on whether the Start snapshot already carries + // the just-opened (empty) block — pop a trailing match into the + // seed so the merge never duplicates it. + AssistantMessageEvent::TextStart { partial } => { + let mut base = partial.clone(); + let seed = match base.content.last() { + Some(ContentBlock::Text { text }) => { + let text = text.clone(); + base.content.pop(); + text + } + _ => String::new(), + }; + self.base = Some(base); + self.open = Some(OpenBlock::Text { + seed, + acc: String::new(), + }); + } + AssistantMessageEvent::ThinkingStart { partial } => { + let mut base = partial.clone(); + let seed = match base.content.last() { + Some(ContentBlock::Thinking { text, .. }) => { + let text = text.clone(); + base.content.pop(); + text + } + _ => String::new(), + }; + self.base = Some(base); + self.open = Some(OpenBlock::Thinking { + seed, + acc: String::new(), + }); + } + // The Start snapshot keeps the call block (placeholder args); + // deltas accumulate the raw argument text for the merge. + AssistantMessageEvent::FunctioncallStart { partial } => { + self.base = Some(partial.clone()); + self.open = Some(OpenBlock::Call { + args: String::new(), + }); + } + AssistantMessageEvent::TextDelta { partial, delta } => match partial { + Some(p) => { + self.base = Some(p.clone()); + self.open = None; + } + None => match &mut self.open { + Some(OpenBlock::Text { acc, .. }) => acc.push_str(delta), + // Defensive: producer skipped the Start frame. + _ => { + self.open = Some(OpenBlock::Text { + seed: String::new(), + acc: delta.clone(), + }) + } + }, + }, + AssistantMessageEvent::ThinkingDelta { partial, delta } => match partial { + Some(p) => { + self.base = Some(p.clone()); + self.open = None; + } + None => match &mut self.open { + Some(OpenBlock::Thinking { acc, .. }) => acc.push_str(delta), + _ => { + self.open = Some(OpenBlock::Thinking { + seed: String::new(), + acc: delta.clone(), + }) + } + }, + }, + AssistantMessageEvent::FunctioncallDelta { partial, delta, .. } => match partial { + Some(p) => { + self.base = Some(p.clone()); + self.open = None; + } + None => match &mut self.open { + Some(OpenBlock::Call { args }) => args.push_str(delta), + _ => { + self.open = Some(OpenBlock::Call { + args: delta.clone(), + }) + } + }, + }, + AssistantMessageEvent::Usage { .. } + | AssistantMessageEvent::Ping + | AssistantMessageEvent::Stop { .. } + | AssistantMessageEvent::Done { .. } + | AssistantMessageEvent::Error { .. } => {} + } + } + + /// The message so far: the boundary snapshot plus the open block. + /// `None` before any content frame. + pub fn current(&self) -> Option { + let mut out = self.base.clone()?; + match &self.open { + Some(OpenBlock::Text { seed, acc }) if !(seed.is_empty() && acc.is_empty()) => { + out.content.push(ContentBlock::Text { + text: format!("{seed}{acc}"), + }); + } + Some(OpenBlock::Thinking { seed, acc }) if !(seed.is_empty() && acc.is_empty()) => { + out.content.push(ContentBlock::Thinking { + text: format!("{seed}{acc}"), + signature: None, + }); + } + // Mid-call death: replace the open call block's placeholder args + // with the same replay-safe degraded object providers produce. + Some(OpenBlock::Call { args }) if !args.is_empty() => { + if let Some(ContentBlock::FunctionCall { arguments, .. }) = out + .content + .iter_mut() + .rev() + .find(|b| matches!(b, ContentBlock::FunctionCall { .. })) + { + *arguments = serde_json::from_str(args) + .ok() + .filter(serde_json::Value::is_object) + .unwrap_or_else(|| degraded_arguments(args)); + } + } + _ => {} + } + Some(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::chat::synthesize::empty_partial; + use crate::types::events::AssistantMessageEvent as Ev; + use serde_json::json; + + fn base() -> AssistantMessage { + empty_partial("m", "p", 1) + } + + #[test] + fn slim_text_deltas_accumulate_from_the_boundary_snapshot() { + let mut acc = PartialAccumulator::default(); + assert!(acc.current().is_none()); + acc.apply(&Ev::Start { partial: base() }); + acc.apply(&Ev::TextStart { partial: base() }); + acc.apply(&Ev::TextDelta { + partial: None, + delta: "Hel".into(), + }); + acc.apply(&Ev::TextDelta { + partial: None, + delta: "lo".into(), + }); + assert_eq!( + acc.current().unwrap().content, + vec![ContentBlock::Text { + text: "Hello".into() + }] + ); + } + + #[test] + fn start_snapshots_with_and_without_the_open_block_merge_identically() { + // Producer A: TextStart snapshot excludes the new (empty) block. + let mut a = PartialAccumulator::default(); + a.apply(&Ev::TextStart { partial: base() }); + a.apply(&Ev::TextDelta { + partial: None, + delta: "x".into(), + }); + // Producer B: TextStart snapshot includes the empty open block. + let mut with_empty = base(); + with_empty.content.push(ContentBlock::Text { + text: String::new(), + }); + let mut b = PartialAccumulator::default(); + b.apply(&Ev::TextStart { + partial: with_empty, + }); + b.apply(&Ev::TextDelta { + partial: None, + delta: "x".into(), + }); + assert_eq!(a.current().unwrap().content, b.current().unwrap().content); + } + + #[test] + fn thinking_end_snapshot_carries_the_signature() { + let mut acc = PartialAccumulator::default(); + acc.apply(&Ev::ThinkingStart { partial: base() }); + acc.apply(&Ev::ThinkingDelta { + partial: None, + delta: "reasoning".into(), + }); + // Mid-block view has no signature (it only exists at the end). + assert_eq!( + acc.current().unwrap().content, + vec![ContentBlock::Thinking { + text: "reasoning".into(), + signature: None + }] + ); + let mut done = base(); + done.content = vec![ContentBlock::Thinking { + text: "reasoning".into(), + signature: Some("sig".into()), + }]; + acc.apply(&Ev::ThinkingEnd { + partial: done.clone(), + }); + assert_eq!(acc.current().unwrap().content, done.content); + } + + #[test] + fn mid_call_death_yields_degraded_replay_safe_arguments() { + let mut with_call = base(); + with_call.content = vec![ContentBlock::FunctionCall { + id: "c1".into(), + function_id: "agent_trigger".into(), + arguments: json!({}), + }]; + let mut acc = PartialAccumulator::default(); + acc.apply(&Ev::FunctioncallStart { partial: with_call }); + acc.apply(&Ev::FunctioncallDelta { + partial: None, + delta: r#"{"function":"state::set","payload":{"x":"#.into(), + id: "c1".into(), + }); + let cum = acc.current().unwrap(); + let ContentBlock::FunctionCall { arguments, .. } = &cum.content[0] else { + panic!("want function_call"); + }; + // Same degraded shape providers produce for incomplete args. + assert_eq!(arguments["function"], "state::set"); + assert_eq!(arguments["_partial"], true); + } + + #[test] + fn complete_call_args_parse_into_a_real_object() { + let mut with_call = base(); + with_call.content = vec![ContentBlock::FunctionCall { + id: "c1".into(), + function_id: "agent_trigger".into(), + arguments: json!({}), + }]; + let mut acc = PartialAccumulator::default(); + acc.apply(&Ev::FunctioncallStart { partial: with_call }); + acc.apply(&Ev::FunctioncallDelta { + partial: None, + delta: r#"{"function":"state::get"}"#.into(), + id: "c1".into(), + }); + let cum = acc.current().unwrap(); + let ContentBlock::FunctionCall { arguments, .. } = &cum.content[0] else { + panic!("want function_call"); + }; + assert_eq!(arguments, &json!({"function":"state::get"})); + } + + #[test] + fn legacy_fat_delta_replaces_the_snapshot_wholesale() { + let mut fat = base(); + fat.content = vec![ContentBlock::Text { + text: "cumulative".into(), + }]; + let mut acc = PartialAccumulator::default(); + acc.apply(&Ev::TextDelta { + partial: Some(fat.clone()), + delta: "e".into(), + }); + // The fat partial IS the message: the delta is not re-applied. + assert_eq!(acc.current().unwrap().content, fat.content); + } + + #[test] + fn delta_without_a_start_frame_does_not_panic() { + let mut acc = PartialAccumulator::default(); + acc.apply(&Ev::TextDelta { + partial: None, + delta: "orphan".into(), + }); + // No boundary snapshot yet → nothing to attach the block to. + assert!(acc.current().is_none()); + // A late Start attaches subsequent deltas normally. + acc.apply(&Ev::Start { partial: base() }); + acc.apply(&Ev::TextDelta { + partial: None, + delta: "ok".into(), + }); + assert_eq!( + acc.current().unwrap().content, + vec![ContentBlock::Text { text: "ok".into() }] + ); + } +} diff --git a/llm-router/src/chat/mod.rs b/llm-router/src/chat/mod.rs index 4d4977327..f0c0fd4d1 100644 --- a/llm-router/src/chat/mod.rs +++ b/llm-router/src/chat/mod.rs @@ -1,4 +1,5 @@ pub mod abort; +pub mod accumulate; #[allow(clippy::module_inception)] // chat/chat.rs is the plan-locked file layout pub mod chat; pub mod complete; diff --git a/llm-router/src/chat/relay.rs b/llm-router/src/chat/relay.rs index ca6a4a9e9..adaeaa7a6 100644 --- a/llm-router/src/chat/relay.rs +++ b/llm-router/src/chat/relay.rs @@ -92,26 +92,10 @@ pub struct RelayOpts { pub aborted: Arc, } -fn partial_of(ev: &AssistantMessageEvent) -> Option<&AssistantMessage> { - use AssistantMessageEvent as E; - match ev { - E::Start { partial } - | E::TextStart { partial } - | E::TextDelta { partial, .. } - | E::TextEnd { partial } - | E::ThinkingStart { partial } - | E::ThinkingDelta { partial, .. } - | E::ThinkingEnd { partial } - | E::FunctioncallStart { partial } - | E::FunctioncallDelta { partial, .. } - | E::FunctioncallEnd { partial } => Some(partial), - _ => None, - } -} - /// One provider attempt (design § chat flow step 5). Reads provider frames, /// enforces the idle budget, fills cost_usd, forwards in order, tracks the -/// partial, classifies how the stream ended. +/// cumulative message (boundary snapshots + accumulated deltas), classifies +/// how the stream ended. /// /// Idle semantics: before any content, every frame (incl. ping) resets the /// budget — a slow first token behind provider keepalives is legitimate @@ -125,7 +109,9 @@ pub async fn relay_frames( opts: &RelayOpts, ) -> RelayResult { let mut forwarded = false; // a non-ping frame reached the caller (gates retry) - let mut partial: Option = None; + // Cumulative message for abort/no-terminal synthesis: slim deltas are + // cheap string appends; legacy fat partials replace the snapshot. + let mut acc = crate::chat::accumulate::PartialAccumulator::default(); let mut usage: Option = None; let mut content_started = false; let mut last_progress = std::time::Instant::now(); @@ -145,7 +131,7 @@ pub async fn relay_frames( if opts.aborted.load(Ordering::SeqCst) { reader.close(); return RelayResult::Aborted { - partial, + partial: acc.current(), usage, forwarded, }; @@ -154,7 +140,7 @@ pub async fn relay_frames( ReadEvent::Eof => { return RelayResult::NoTerminal { reason: NoTerminalReason::Closed, - partial, + partial: acc.current(), usage, forwarded, } @@ -163,7 +149,7 @@ pub async fn relay_frames( reader.close(); return RelayResult::NoTerminal { reason: NoTerminalReason::Idle, - partial, + partial: acc.current(), usage, forwarded, }; @@ -176,57 +162,58 @@ pub async fn relay_frames( last_progress = std::time::Instant::now(); // Start carries an empty partial pre-content; it must not // arm the strict budget or slow-first-token providers trip. - if partial_of(&ev).is_some() - && !matches!(ev, AssistantMessageEvent::Start { .. }) - { + if ev.is_content() && !matches!(ev, AssistantMessageEvent::Start { .. }) { content_started = true; } } - if let Some(p) = partial_of(&ev) { - partial = Some(p.clone()); - } + acc.apply(&ev); + let is_ping = matches!(ev, AssistantMessageEvent::Ping); - // Enrich usage with cost before forwarding (router fills cost_usd). - let out = match ev { + // Only Usage/Done/Error are enriched (cost_usd) and need + // re-serialization; every other frame is forwarded as the + // original string — no per-frame re-serialize on the hot path. + let enriched = match ev { AssistantMessageEvent::Usage { usage: u } => { let filled = fill_cost_usd(&u, opts.pricing.as_ref()); usage = Some(filled.clone()); - AssistantMessageEvent::Usage { usage: filled } + Some(AssistantMessageEvent::Usage { usage: filled }) } AssistantMessageEvent::Done { mut message } => { message.usage = message .usage .map(|u| fill_cost_usd(&u, opts.pricing.as_ref())); - AssistantMessageEvent::Done { message } + Some(AssistantMessageEvent::Done { message }) } AssistantMessageEvent::Error { mut error } => { error.usage = error .usage .map(|u| fill_cost_usd(&u, opts.pricing.as_ref())) .or_else(|| usage.clone()); - AssistantMessageEvent::Error { error } + Some(AssistantMessageEvent::Error { error }) } - other => other, + _ => None, }; + let is_done = matches!(enriched, Some(AssistantMessageEvent::Done { .. })); + let is_error = matches!(enriched, Some(AssistantMessageEvent::Error { .. })); // Terminal-holdback: a pre-content error terminal stays with us // so chat.rs can retry the attempt invisibly. - let is_done = matches!(out, AssistantMessageEvent::Done { .. }); - let is_error = matches!(out, AssistantMessageEvent::Error { .. }); if is_error && !forwarded { reader.close(); return RelayResult::ErrorFrame { - terminal: out, + terminal: enriched.expect("error frame just matched"), forwarded: false, terminal_forwarded: false, }; } - let is_ping = matches!(out, AssistantMessageEvent::Ping); - if sink - .send(&serde_json::to_string(&out).expect("serializable frame")) - .is_err() - { + let frame = match &enriched { + Some(out) => std::borrow::Cow::Owned( + serde_json::to_string(out).expect("serializable frame"), + ), + None => std::borrow::Cow::Borrowed(msg.as_str()), + }; + if sink.send(&frame).is_err() { reader.close(); // closure propagates: the provider's writes now fail return RelayResult::CallerGone { forwarded }; } @@ -235,13 +222,13 @@ pub async fn relay_frames( } if is_done { return RelayResult::Done { - terminal: out, + terminal: enriched.expect("done frame just matched"), forwarded, }; } if is_error { return RelayResult::ErrorFrame { - terminal: out, + terminal: enriched.expect("error frame just matched"), forwarded, terminal_forwarded: true, }; @@ -492,7 +479,7 @@ mod loop_tests { send( &provider, &AssistantMessageEvent::TextDelta { - partial: partial("hi"), + partial: Some(partial("hi")), delta: "hi".into(), }, ); @@ -563,4 +550,80 @@ mod loop_tests { }; assert_eq!(p.unwrap().content.len(), 1); } + + #[tokio::test] + async fn slim_delta_stream_synthesis_carries_the_accumulated_text() { + // A provider dies mid-block after slim (partial-less) deltas: the + // no-terminal result must still carry every delta received. + let provider = FakeChannel::new(); + let caller = FakeChannel::new(); + send( + &provider, + &AssistantMessageEvent::Start { + partial: partial(""), + }, + ); + send( + &provider, + &AssistantMessageEvent::TextStart { + partial: partial(""), + }, + ); + send( + &provider, + &AssistantMessageEvent::TextDelta { + partial: None, + delta: "Hel".into(), + }, + ); + send( + &provider, + &AssistantMessageEvent::TextDelta { + partial: None, + delta: "lo".into(), + }, + ); + provider.writer.close(); // crash: no terminal frame + let (_, o) = opts(1000); + let mut reader: Box = Box::new(provider.reader); + let RelayResult::NoTerminal { partial: p, .. } = + relay_frames(&mut reader, &caller.writer.clone(), &o).await + else { + panic!("want no-terminal"); + }; + let msg = p.expect("accumulated partial"); + assert_eq!( + msg.content, + vec![crate::types::content::ContentBlock::Text { + text: "Hello".into() + }] + ); + } + + #[tokio::test] + async fn non_enriched_frames_are_forwarded_byte_identical() { + // The relay must not re-serialize pass-through frames: the caller + // receives the provider's exact bytes (key order, whitespace, all). + let provider = FakeChannel::new(); + let mut caller = FakeChannel::new(); + let odd_spacing = r#"{ "type": "text_delta", "delta": "hi" }"#; + provider.writer.send(odd_spacing).unwrap(); + let done = serde_json::to_string(&AssistantMessageEvent::Done { + message: partial("hi"), + }) + .unwrap(); + provider.writer.send(&done).unwrap(); + provider.writer.close(); + + let (_, o) = opts(1000); + let mut reader: Box = Box::new(provider.reader); + let result = relay_frames(&mut reader, &caller.writer.clone(), &o).await; + assert!(matches!(result, RelayResult::Done { .. })); + + let mut raw = vec![]; + while let ReadEvent::Msg(m) = caller.reader.next(Duration::from_millis(50)).await { + raw.push(m); + } + assert_eq!(raw[0], odd_spacing, "pass-through frame was rewritten"); + } } diff --git a/llm-router/src/lib.rs b/llm-router/src/lib.rs index 348910da1..5ca17b899 100644 --- a/llm-router/src/lib.rs +++ b/llm-router/src/lib.rs @@ -7,6 +7,7 @@ pub mod chat; pub mod config; pub mod embed; pub mod manifest; +pub mod provider_scaffold; pub mod register; pub mod registry; pub mod routing; diff --git a/llm-router/src/provider_scaffold/cache.rs b/llm-router/src/provider_scaffold/cache.rs new file mode 100644 index 000000000..f8b5a3155 --- /dev/null +++ b/llm-router/src/provider_scaffold/cache.rs @@ -0,0 +1,138 @@ +//! Per-provider caches for the two engine round trips every stream call +//! paid before its first upstream byte: the registration token +//! (`state::get`) and `router::provider::resolve`. +//! +//! The token changes only on re-registration, which happens in-process — +//! `store` updates the cell alongside iii-state, so reads never miss after +//! boot. The resolve response (credential, api_url, max_tokens) changes when +//! an operator edits the ROUTER's configuration entry, which providers +//! cannot observe; a short TTL bounds that staleness, and callers invalidate +//! eagerly on `router::ready` (router restart) and on upstream auth errors +//! (credential rotated out from under us). Invalidation only drops the +//! cache — retry/failover stays the router's job. + +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; + +use iii_sdk::errors::Error; +use iii_sdk::IIIClient; + +use crate::types::router::ProviderResolveResponse; + +/// Staleness bound after an unobservable operator edit to the router's +/// provider slice (credential/api_url/max_tokens). +pub const RESOLVE_TTL: Duration = Duration::from_secs(30); + +/// The per-provider scaffold state a stream/register/ready handler clones. +#[derive(Clone, Default)] +pub struct ScaffoldCache { + token: Arc>>, + resolve: Arc>>, +} + +impl ScaffoldCache { + pub fn new() -> Self { + Self::default() + } + + /// Cached token, falling back to iii-state and populating the cell. + pub async fn load_token(&self, iii: &IIIClient, scope: &str) -> Option { + if let Some(token) = self + .token + .read() + .expect("token cache lock poisoned") + .clone() + { + return Some(token); + } + let token = super::state::load_token(iii, scope).await?; + *self.token.write().expect("token cache lock poisoned") = Some(token.clone()); + Some(token) + } + + /// Persist a (re)issued token and refresh the cell. + pub async fn store_token( + &self, + iii: &IIIClient, + scope: &str, + token: &str, + ) -> Result<(), Error> { + super::state::store_token(iii, scope, token).await?; + *self.token.write().expect("token cache lock poisoned") = Some(token.to_string()); + Ok(()) + } + + /// `router::provider::resolve`, served from cache within [`RESOLVE_TTL`]. + pub async fn resolve( + &self, + iii: &IIIClient, + provider_id: &str, + token: Option<&str>, + ) -> Result { + if let Some(resolved) = self.fresh_resolve() { + return Ok(resolved); + } + // The lock is never held across the await; concurrent misses may + // duplicate one resolve, which is harmless. + let resolved = super::router_client::resolve(iii, provider_id, token).await?; + *self.resolve.write().expect("resolve cache lock poisoned") = + Some((resolved.clone(), Instant::now())); + Ok(resolved) + } + + fn fresh_resolve(&self) -> Option { + let guard = self.resolve.read().expect("resolve cache lock poisoned"); + let (resolved, stored_at) = guard.as_ref()?; + (stored_at.elapsed() <= RESOLVE_TTL).then(|| resolved.clone()) + } + + /// Drop the cached resolve response AND token. Call on `router::ready` + /// (a restarted router may carry new config; the token survives but the + /// re-declare path refreshes it anyway) and on upstream auth errors + /// (the operator rotated the credential; the next attempt re-resolves). + pub fn invalidate(&self) { + *self.resolve.write().expect("resolve cache lock poisoned") = None; + *self.token.write().expect("token cache lock poisoned") = None; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn resolved(url: &str) -> ProviderResolveResponse { + serde_json::from_value(serde_json::json!({ + "configured": true, + "source": "config", + "credential": null, + "api_url": url, + })) + .expect("minimal resolve response shape") + } + + #[test] + fn resolve_cache_serves_within_ttl_and_invalidates() { + let cache = ScaffoldCache::new(); + assert!(cache.fresh_resolve().is_none()); + *cache.resolve.write().unwrap() = Some((resolved("https://a"), Instant::now())); + assert!(cache.fresh_resolve().is_some()); + cache.invalidate(); + assert!(cache.fresh_resolve().is_none()); + } + + #[test] + fn resolve_cache_expires_after_ttl() { + let cache = ScaffoldCache::new(); + let stale = Instant::now() - (RESOLVE_TTL + Duration::from_secs(1)); + *cache.resolve.write().unwrap() = Some((resolved("https://a"), stale)); + assert!(cache.fresh_resolve().is_none()); + } + + #[test] + fn invalidate_also_clears_the_token_cell() { + let cache = ScaffoldCache::new(); + *cache.token.write().unwrap() = Some("tok".into()); + cache.invalidate(); + assert!(cache.token.read().unwrap().is_none()); + } +} diff --git a/llm-router/src/provider_scaffold/mod.rs b/llm-router/src/provider_scaffold/mod.rs new file mode 100644 index 000000000..22c975ac7 --- /dev/null +++ b/llm-router/src/provider_scaffold/mod.rs @@ -0,0 +1,16 @@ +//! Shared scaffolding for provider workers. Every provider worker carries +//! the same plumbing around its provider-specific wire code: the router +//! protocol client, registration-token persistence, the event pump into the +//! router-owned channel, SSE transport buffering, and the tool-name codec. +//! These were verbatim copies across the provider crates (the pump literally +//! carried a "shared extraction into llm-router is a listed follow-up" +//! comment); this module is that extraction. Providers keep only what is +//! genuinely theirs: request building, SSE event decoding, and error +//! classification. + +pub mod cache; +pub mod names; +pub mod pump; +pub mod router_client; +pub mod sse_transport; +pub mod state; diff --git a/llm-router/src/provider_scaffold/names.rs b/llm-router/src/provider_scaffold/names.rs new file mode 100644 index 000000000..7064750f3 --- /dev/null +++ b/llm-router/src/provider_scaffold/names.rs @@ -0,0 +1,36 @@ +//! iii function ids ↔ provider tool names. Upstream APIs enforce +//! `^[a-zA-Z0-9_-]{1,128}$`-style names; bus ids use `::` separators. + +pub fn encode_tool_name(name: &str) -> String { + name.replace("::", "__") +} + +/// Inverse of `encode_tool_name`. Lossy precondition: an id containing a +/// literal `__` decodes to `::` — such ids are not in use on the bus today. +pub fn decode_tool_name(name: &str) -> String { + name.replace("__", "::") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn round_trips_bus_ids() { + assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); + assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); + assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); + } + + #[test] + fn plain_names_pass_through() { + assert_eq!(encode_tool_name("submit_result"), "submit_result"); + assert_eq!(decode_tool_name("submit_result"), "submit_result"); + } + + #[test] + fn literal_double_underscore_is_the_accepted_lossy_case() { + // known limitation: a literal `__` in an id decodes to `::` + assert_eq!(decode_tool_name("a__b"), "a::b"); + } +} diff --git a/llm-router/src/provider_scaffold/pump.rs b/llm-router/src/provider_scaffold/pump.rs new file mode 100644 index 000000000..93e53cb73 --- /dev/null +++ b/llm-router/src/provider_scaffold/pump.rs @@ -0,0 +1,177 @@ +//! Forward upstream events into the router-owned channel (spec § Provider +//! stream contract): AssistantMessageEvent frames as JSON text messages, +//! terminal done/error last, pings through silence. Previously a verbatim +//! copy in every provider crate. +use crate::chat::relay::FrameSink; +use crate::types::events::{AssistantMessageEvent, ErrorKind}; +use std::time::Duration; +use tokio::sync::mpsc; + +/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). +pub const PING_INTERVAL: Duration = Duration::from_secs(30); + +/// `Err(())` means only "the sink is gone — stop writing"; there is no +/// error detail to carry, so the unit error is the whole contract. +#[allow(clippy::result_unit_err)] +pub fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { + let frame = serde_json::to_string(ev).expect("serializable event"); + sink.send(&frame).map_err(|_| ()) +} + +/// Forward upstream events to the sink; ping through silence; stop on the +/// terminal event or on a failed write (caller gone → dropping `rx` aborts +/// the upstream task and its in-flight HTTP request). +/// +/// Returns the `error_kind` of a forwarded terminal Error frame (`None` for +/// done / no terminal / caller gone) so providers can react to specific +/// failures — e.g. dropping a cached credential on `AuthExpired`. +pub async fn pump( + mut rx: mpsc::Receiver, + sink: &dyn FrameSink, + ping_interval: Duration, +) -> Option { + loop { + match tokio::time::timeout(ping_interval, rx.recv()).await { + Ok(Some(ev)) => { + let terminal = ev.is_terminal(); + let error_kind = match &ev { + AssistantMessageEvent::Error { error } => error.error_kind, + _ => None, + }; + if send_event(sink, &ev).is_err() { + return None; + } + if terminal { + return error_kind; + } + } + // Upstream task ended without a terminal (panic/abort): the + // router synthesizes the terminal frame — never two terminals. + Ok(None) => return None, + // Silent stretch: heartbeat (also probes for a gone caller). + Err(_elapsed) => { + if send_event(sink, &AssistantMessageEvent::Ping).is_err() { + return None; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::chat::relay::{ReadEvent, RelayRead}; + use crate::chat::synthesize::empty_partial; + use crate::testkit::fake_channels::FakeChannel; + use crate::types::messages::AssistantMessage; + use serde_json::Value; + + fn now_ms() -> i64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0) + } + + fn empty_assistant(model: &str) -> AssistantMessage { + empty_partial(model, "test-provider", now_ms()) + } + + fn done_event() -> AssistantMessageEvent { + AssistantMessageEvent::Done { + message: empty_assistant("model-test"), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn forwards_events_and_stops_at_terminal() { + let ch = FakeChannel::new(); + let (tx, rx) = mpsc::channel(8); + tx.send(AssistantMessageEvent::Start { + partial: empty_assistant("m"), + }) + .await + .unwrap(); + tx.send(done_event()).await.unwrap(); + // a frame after the terminal must never be forwarded + tx.send(AssistantMessageEvent::Ping).await.unwrap(); + drop(tx); + + pump(rx, &ch.writer, Duration::from_secs(30)).await; + ch.writer.close(); + + let mut frames = Vec::new(); + let mut reader = ch.reader; + while let ReadEvent::Msg(m) = reader.next(Duration::from_millis(100)).await { + frames.push(m); + } + assert_eq!(frames.len(), 2); + let last: Value = serde_json::from_str(&frames[1]).unwrap(); + assert_eq!(last["type"], "done"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn pings_through_silence() { + let ch = FakeChannel::new(); + let (tx, rx) = mpsc::channel::(8); + // hold tx open, send nothing for > 2 ping intervals, then terminate + let pump_task = { + let writer = ch.writer.clone(); + tokio::spawn(async move { pump(rx, &writer, Duration::from_millis(50)).await }) + }; + tokio::time::sleep(Duration::from_millis(140)).await; + tx.send(done_event()).await.unwrap(); + drop(tx); + pump_task.await.unwrap(); + ch.writer.close(); + + let mut frames = Vec::new(); + let mut reader = ch.reader; + while let ReadEvent::Msg(m) = reader.next(Duration::from_millis(100)).await { + frames.push(m); + } + let pings = frames + .iter() + .filter(|f| serde_json::from_str::(f).unwrap()["type"] == "ping") + .count(); + assert!( + pings >= 2, + "want >=2 pings through 140ms of silence, got {pings}" + ); + assert_eq!( + serde_json::from_str::(frames.last().unwrap()).unwrap()["type"], + "done" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn terminal_error_kind_is_returned_to_the_caller() { + let ch = FakeChannel::new(); + let (tx, rx) = mpsc::channel(8); + let mut error = empty_assistant("m"); + error.error_kind = Some(ErrorKind::AuthExpired); + tx.send(AssistantMessageEvent::Error { error }) + .await + .unwrap(); + drop(tx); + let kind = pump(rx, &ch.writer, Duration::from_secs(30)).await; + assert_eq!(kind, Some(ErrorKind::AuthExpired)); + } + + #[tokio::test(flavor = "multi_thread")] + async fn reader_close_stops_the_pump_and_drops_the_receiver() { + let ch = FakeChannel::new(); + ch.reader.close(); // caller gone before anything is written + let (tx, rx) = mpsc::channel(8); + tx.send(AssistantMessageEvent::Start { + partial: empty_assistant("m"), + }) + .await + .unwrap(); + pump(rx, &ch.writer, Duration::from_secs(30)).await; // returns immediately + // the receiver was consumed and dropped by pump → upstream send fails + assert!(tx.send(done_event()).await.is_err()); + } +} diff --git a/llm-router/src/provider_scaffold/router_client.rs b/llm-router/src/provider_scaffold/router_client.rs new file mode 100644 index 000000000..3e07ee4db --- /dev/null +++ b/llm-router/src/provider_scaffold/router_client.rs @@ -0,0 +1,72 @@ +//! Thin wrappers over the router's provider-protocol functions. All calls +//! carry the registration token (identity binding, spec adaptation #1). +//! `provider_id` is the caller's declared provider id (e.g. "anthropic"). +use crate::types::model::Model; +use crate::types::router::ProviderResolveResponse; +use iii_sdk::errors::Error; +use iii_sdk::protocol::TriggerRequest; +use iii_sdk::IIIClient; +use serde_json::{json, Value}; + +pub async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { + iii.trigger(TriggerRequest { + function_id: function_id.into(), + payload, + action: None, + timeout_ms: Some(15_000), + }) + .await +} + +/// `router::provider::resolve` — credential + effective settings. +pub async fn resolve( + iii: &IIIClient, + provider_id: &str, + token: Option<&str>, +) -> Result { + let mut payload = json!({ "id": provider_id }); + if let Some(t) = token { + payload["token"] = json!(t); + } + let raw = call(iii, "router::provider::resolve", payload).await?; + serde_json::from_value(raw).map_err(|e| Error::Remote { + code: "provider/bad_resolve_response".into(), + message: e.to_string(), + stacktrace: None, + }) +} + +/// `router::models::reconcile` — replace this provider's catalog slice. +pub async fn reconcile( + iii: &IIIClient, + provider_id: &str, + models: Vec, + token: Option<&str>, +) -> Result<(), Error> { + let mut payload = json!({ + "provider": provider_id, + "models": serde_json::to_value(models).expect("serializable models"), + }); + if let Some(t) = token { + payload["token"] = json!(t); + } + call(iii, "router::models::reconcile", payload).await?; + Ok(()) +} + +/// `router::models::get` — authoritative catalog record (None when absent). +pub async fn models_get(iii: &IIIClient, provider_id: &str, model_id: &str) -> Option { + let raw = call( + iii, + "router::models::get", + json!({ "provider": provider_id, "id": model_id }), + ) + .await + .ok()?; + serde_json::from_value(raw.get("model")?.clone()).ok() +} + +/// `router::provider::register` — returns the registration token to persist. +pub async fn register(iii: &IIIClient, declaration: Value) -> Result { + call(iii, "router::provider::register", declaration).await +} diff --git a/llm-router/src/provider_scaffold/sse_transport.rs b/llm-router/src/provider_scaffold/sse_transport.rs new file mode 100644 index 000000000..def6e4a06 --- /dev/null +++ b/llm-router/src/provider_scaffold/sse_transport.rs @@ -0,0 +1,164 @@ +//! Transport-layer SSE helpers shared by provider upstream readers: +//! error-chain flattening, cross-chunk UTF-8 buffering, CRLF +//! normalization, and `\n\n`-delimited block draining. Decoding the +//! blocks into events stays provider-specific (the closure). +use crate::types::events::AssistantMessageEvent; +use tokio::sync::mpsc; + +/// Flatten an error and its `source()` chain into one string. reqwest's +/// top-level Display for a builder error is just "builder error"; the real +/// cause (invalid header value, bad URL) lives in the source chain, so without +/// this the message is undiagnosable. +pub fn error_chain(e: &dyn std::error::Error) -> String { + let mut msg = e.to_string(); + let mut src = e.source(); + while let Some(s) = src { + let next = s.to_string(); + // reqwest sometimes nests the same text; skip exact repeats. + if !msg.ends_with(&next) { + msg.push_str(": "); + msg.push_str(&next); + } + src = s.source(); + } + msg +} + +/// Append chunk bytes to `text`, retaining any trailing incomplete UTF-8 +/// sequence in `byte_buf`. Network chunks split multibyte codepoints; a +/// per-chunk lossy conversion corrupts them to U+FFFD. +pub fn append_utf8_chunk(byte_buf: &mut Vec, text: &mut String, chunk: &[u8]) { + byte_buf.extend_from_slice(chunk); + let mut consumed = 0usize; + loop { + match std::str::from_utf8(&byte_buf[consumed..]) { + Ok(s) => { + text.push_str(s); + byte_buf.clear(); + return; + } + Err(e) => { + let valid = e.valid_up_to(); + if valid > 0 { + // SAFETY: valid_up_to guarantees valid UTF-8 in this prefix. + text.push_str(unsafe { + std::str::from_utf8_unchecked(&byte_buf[consumed..consumed + valid]) + }); + consumed += valid; + } + match e.error_len() { + Some(invalid) => { + byte_buf.drain(..consumed + invalid); + text.push('\u{FFFD}'); + consumed = 0; + } + None => { + if consumed > 0 { + byte_buf.drain(..consumed); + } + return; + } + } + } + } + } +} + +/// SSE allows CRLF line endings; normalize so `\n\n` block framing holds. +pub fn normalize_crlf(text: &mut String) { + if text.contains("\r\n") { + *text = text.replace("\r\n", "\n"); + } +} + +/// Drain complete `\n\n`-delimited SSE blocks from `text`, decoding each with +/// `handle_block` and forwarding the events. Returns true when a terminal +/// event was forwarded or the receiver is gone. +pub async fn drain_sse_blocks( + text: &mut String, + tx: &mpsc::Sender, + handle_block: &mut F, +) -> bool +where + F: FnMut(&str) -> Vec, +{ + normalize_crlf(text); + while let Some(idx) = text.find("\n\n") { + let block: String = text.drain(..idx + 2).collect(); + for ev in handle_block(&block) { + let terminal = ev.is_terminal(); + if tx.send(ev).await.is_err() { + return true; + } + if terminal { + return true; + } + } + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn utf8_split_across_chunks_is_preserved() { + let emoji = "Hello 🌍"; + let bytes = emoji.as_bytes(); + let split = "Hello ".len() + 1; // split inside the 4-byte emoji + let mut byte_buf = Vec::new(); + let mut text = String::new(); + append_utf8_chunk(&mut byte_buf, &mut text, &bytes[..split]); + assert_eq!(text, "Hello "); + append_utf8_chunk(&mut byte_buf, &mut text, &bytes[split..]); + assert_eq!(text, emoji); + } + + #[test] + fn utf8_split_across_three_chunks_is_preserved() { + // 4-byte emoji delivered one byte at a time — the worst case a + // network stream can produce. + let s = "héllo 🌍!"; + let bytes = s.as_bytes(); + let mut byte_buf = Vec::new(); + let mut text = String::new(); + for b in bytes { + append_utf8_chunk(&mut byte_buf, &mut text, std::slice::from_ref(b)); + } + assert_eq!(text, s); + assert!(byte_buf.is_empty()); + } + + #[test] + fn invalid_bytes_become_replacement_chars_without_stalling() { + let mut byte_buf = Vec::new(); + let mut text = String::new(); + append_utf8_chunk(&mut byte_buf, &mut text, b"ok\xFF\xFEok"); + assert_eq!(text, "ok\u{FFFD}\u{FFFD}ok"); + assert!(byte_buf.is_empty()); + } + + #[test] + fn crlf_blocks_are_reframed() { + let mut text = "data: a\r\n\r\ndata: b".to_string(); + normalize_crlf(&mut text); + assert_eq!(text, "data: a\n\ndata: b"); + } + + #[tokio::test] + async fn drains_blocks_and_stops_on_terminal() { + let (tx, mut rx) = mpsc::channel(8); + let mut text = "one\n\ntwo\n\nrest".to_string(); + let mut seen = Vec::new(); + let done = drain_sse_blocks(&mut text, &tx, &mut |block: &str| { + seen.push(block.to_string()); + vec![AssistantMessageEvent::Ping] + }) + .await; + assert!(!done); + assert_eq!(seen, vec!["one\n\n", "two\n\n"]); + assert_eq!(text, "rest"); + assert!(matches!(rx.try_recv(), Ok(AssistantMessageEvent::Ping))); + } +} diff --git a/llm-router/src/provider_scaffold/state.rs b/llm-router/src/provider_scaffold/state.rs new file mode 100644 index 000000000..fbbd746fd --- /dev/null +++ b/llm-router/src/provider_scaffold/state.rs @@ -0,0 +1,33 @@ +//! Registration-token persistence in iii-state (engine `state::*` functions, +//! binary-worker.md § 7). The raw token lives under the provider's own scope +//! (its worker id); the router persists only its sha256 hash. +use iii_sdk::errors::Error; +use iii_sdk::protocol::TriggerRequest; +use iii_sdk::IIIClient; +use serde_json::{json, Value}; + +const TOKEN_KEY: &str = "registration_token"; + +pub async fn load_token(iii: &IIIClient, scope: &str) -> Option { + let value = iii + .trigger(TriggerRequest { + function_id: "state::get".into(), + payload: json!({ "scope": scope, "key": TOKEN_KEY }), + action: None, + timeout_ms: None, + }) + .await + .ok()?; + value.as_str().map(String::from) +} + +pub async fn store_token(iii: &IIIClient, scope: &str, token: &str) -> Result<(), Error> { + iii.trigger(TriggerRequest { + function_id: "state::set".into(), + payload: json!({ "scope": scope, "key": TOKEN_KEY, "value": Value::from(token) }), + action: None, + timeout_ms: None, + }) + .await?; + Ok(()) +} diff --git a/llm-router/src/types/events.rs b/llm-router/src/types/events.rs index d92499037..1e661233c 100644 --- a/llm-router/src/types/events.rs +++ b/llm-router/src/types/events.rs @@ -48,6 +48,16 @@ pub struct Usage { /// The frozen 15-variant streaming vocabulary (README § Streaming events). /// New frame types are a contract revision, not a provider choice. +/// +/// Block-boundary frames (Start / *Start / *End) carry a required +/// cumulative `partial` snapshot — thinking signatures and finalized +/// function-call arguments exist only there. The three per-chunk delta +/// variants carry only `delta`: a cumulative snapshot per chunk made every +/// stream O(N²) in output length (providers rebuilt + re-serialized the +/// whole message per delta; every hop re-parsed it). Readers reconstruct +/// the cumulative message as last-boundary-snapshot + accumulated deltas +/// (see `chat::accumulate`); a legacy fat delta (`partial: Some`) is +/// honored as an authoritative snapshot for old producers. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] #[serde(tag = "type", rename_all = "snake_case")] pub enum AssistantMessageEvent { @@ -58,7 +68,8 @@ pub enum AssistantMessageEvent { partial: AssistantMessage, }, TextDelta { - partial: AssistantMessage, + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, delta: String, }, TextEnd { @@ -68,7 +79,8 @@ pub enum AssistantMessageEvent { partial: AssistantMessage, }, ThinkingDelta { - partial: AssistantMessage, + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, delta: String, }, ThinkingEnd { @@ -78,7 +90,8 @@ pub enum AssistantMessageEvent { partial: AssistantMessage, }, FunctioncallDelta { - partial: AssistantMessage, + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, delta: String, /// Call id receiving this delta; empty from pre-id producers. #[serde(default, skip_serializing_if = "String::is_empty")] @@ -113,6 +126,23 @@ impl AssistantMessageEvent { AssistantMessageEvent::Done { .. } | AssistantMessageEvent::Error { .. } ) } + + /// The ten content-bearing variants: block boundaries and deltas. + pub fn is_content(&self) -> bool { + matches!( + self, + AssistantMessageEvent::Start { .. } + | AssistantMessageEvent::TextStart { .. } + | AssistantMessageEvent::TextDelta { .. } + | AssistantMessageEvent::TextEnd { .. } + | AssistantMessageEvent::ThinkingStart { .. } + | AssistantMessageEvent::ThinkingDelta { .. } + | AssistantMessageEvent::ThinkingEnd { .. } + | AssistantMessageEvent::FunctioncallStart { .. } + | AssistantMessageEvent::FunctioncallDelta { .. } + | AssistantMessageEvent::FunctioncallEnd { .. } + ) + } } #[cfg(test)] @@ -142,7 +172,7 @@ mod tests { (AssistantMessageEvent::Start { partial: partial() }, "start"), ( AssistantMessageEvent::TextDelta { - partial: partial(), + partial: Some(partial()), delta: "x".into(), }, "text_delta", @@ -185,6 +215,53 @@ mod tests { assert_eq!(serde_json::to_value(&ev).unwrap(), json); } + #[test] + fn slim_and_fat_delta_frames_both_round_trip() { + // Slim (post-contract-change producer shape): no partial key at all. + let slim = serde_json::json!({ "type": "text_delta", "delta": "x" }); + let ev: AssistantMessageEvent = serde_json::from_value(slim.clone()).unwrap(); + assert!(matches!( + &ev, + AssistantMessageEvent::TextDelta { partial: None, .. } + )); + // partial: None is omitted on the wire, not serialized as null. + assert_eq!(serde_json::to_value(&ev).unwrap(), slim); + + // Fat (legacy producer shape): partial still parses and re-emits. + let fat = serde_json::json!({ + "type": "functioncall_delta", + "partial": serde_json::to_value(partial()).unwrap(), + "delta": "{\"x\":", + "id": "c1" + }); + let ev: AssistantMessageEvent = serde_json::from_value(fat.clone()).unwrap(); + assert!(matches!( + &ev, + AssistantMessageEvent::FunctioncallDelta { + partial: Some(_), + .. + } + )); + assert_eq!(serde_json::to_value(&ev).unwrap(), fat); + } + + #[test] + fn is_content_covers_exactly_the_ten_content_variants() { + assert!(AssistantMessageEvent::Start { partial: partial() }.is_content()); + assert!(AssistantMessageEvent::TextDelta { + partial: None, + delta: "x".into() + } + .is_content()); + assert!(AssistantMessageEvent::FunctioncallEnd { partial: partial() }.is_content()); + assert!(!AssistantMessageEvent::Ping.is_content()); + assert!(!AssistantMessageEvent::Usage { + usage: Usage::default() + } + .is_content()); + assert!(!AssistantMessageEvent::Done { message: partial() }.is_content()); + } + #[test] fn is_terminal_only_for_done_and_error() { assert!(AssistantMessageEvent::Done { message: partial() }.is_terminal()); diff --git a/llm-router/tests/integration.rs b/llm-router/tests/integration.rs index af2f8af03..ed8a010b1 100644 --- a/llm-router/tests/integration.rs +++ b/llm-router/tests/integration.rs @@ -284,14 +284,28 @@ async fn start_live_provider(url: &str, opts: ProviderOptions) -> LiveProvider { } } } + // Slim streaming shape (contract: deltas carry no partial; + // boundary frames carry the cumulative snapshot). let message = json!({ "role": "assistant", "content": [{ "type": "text", "text": "live" }], "stop_reason": "end", "model": model, "provider": "real", "timestamp": 2 }); - writer - .send_message(&json!({ "type": "done", "message": message }).to_string()) - .await - .map_err(|e| Error::Handler(e.to_string()))?; + let start_snapshot = json!({ + "role": "assistant", "content": [], "stop_reason": "end", + "model": model, "provider": "real", "timestamp": 1 + }); + for frame in [ + json!({ "type": "text_start", "partial": start_snapshot }), + json!({ "type": "text_delta", "delta": "li" }), + json!({ "type": "text_delta", "delta": "ve" }), + json!({ "type": "text_end", "partial": message }), + json!({ "type": "done", "message": message }), + ] { + writer + .send_message(&frame.to_string()) + .await + .map_err(|e| Error::Handler(e.to_string()))?; + } let _ = writer.close().await; Ok(json!({ "ok": true })) } @@ -416,6 +430,20 @@ async fn end_to_end_relay_over_a_live_engine() { assert!(frames.len() >= 2, "want >=2 frames, got {}", frames.len()); let last: Value = serde_json::from_str(frames.last().unwrap()).unwrap(); assert_eq!(last["type"], "done"); + // Slim deltas must reach the consumer untouched: no partial key + // materialized anywhere between provider and consumer channel. + let deltas: Vec = frames + .iter() + .map(|f| serde_json::from_str::(f).unwrap()) + .filter(|v| v["type"] == "text_delta") + .collect(); + assert_eq!(deltas.len(), 2, "want the 2 scripted slim deltas"); + for d in &deltas { + assert!( + d.get("partial").is_none(), + "slim delta grew a partial in transit: {d}" + ); + } } let completed = call( diff --git a/provider-anthropic/Cargo.lock b/provider-anthropic/Cargo.lock index f63af11a4..da4b1b5ba 100644 --- a/provider-anthropic/Cargo.lock +++ b/provider-anthropic/Cargo.lock @@ -776,7 +776,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.1.1" +version = "1.1.3" dependencies = [ "async-trait", "clap", diff --git a/provider-anthropic/src/register.rs b/provider-anthropic/src/register.rs index ca516900a..f1d6a6f81 100644 --- a/provider-anthropic/src/register.rs +++ b/provider-anthropic/src/register.rs @@ -9,6 +9,7 @@ use crate::{router_client, state, PROVIDER_ID}; use iii_sdk::errors::Error; use iii_sdk::protocol::RegisterTriggerInput; use iii_sdk::{IIIClient, RegisterFunction}; +use llm_router::provider_scaffold::cache::ScaffoldCache; use llm_router::types::router::{ ProviderDeclaration, ProviderDefaults, ProviderReadyAck, RouterReadyEvent, }; @@ -118,6 +119,11 @@ fn read_timeout() -> Duration { } pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { + // Shared per-process cache for the registration token and the resolve + // response (see llm_router::provider_scaffold::cache). Invalidated on + // router::ready — a restarted router may carry new config and reissues + // declare/refresh anyway — and on upstream auth errors (stream_fn). + let cache = ScaffoldCache::new(); // Streaming uses no total timeout (the router owns stream budgets), but // reads are silence-bounded: a stalled upstream otherwise pings the router // past its idle guard until the engine kills the call at stream_timeout @@ -132,7 +138,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { iii.register_function( surface::STREAM_ID, RegisterFunction::new_async_with_bad_request( - make_stream(iii.clone(), http.clone()), + make_stream(iii.clone(), http.clone(), cache.clone()), invalid_request_from_serde, ) .description(surface::STREAM_DESC) @@ -149,10 +155,12 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { { let iii_ready = iii.clone(); let http_ready = http.clone(); + let cache_ready = cache.clone(); iii.register_function( surface::ON_ROUTER_READY_ID, RegisterFunction::new_async(move |_event: RouterReadyEvent| { let (iii, http) = (iii_ready.clone(), http_ready.clone()); + cache_ready.invalidate(); async move { tokio::spawn(declare_and_refresh(iii, http)); Ok::<_, Error>(ProviderReadyAck { ok: true }) diff --git a/provider-anthropic/src/router_client.rs b/provider-anthropic/src/router_client.rs index 373c350e5..175f9e844 100644 --- a/provider-anthropic/src/router_client.rs +++ b/provider-anthropic/src/router_client.rs @@ -1,38 +1,20 @@ -//! Thin wrappers over the router's provider-protocol functions. All calls -//! carry the registration token (identity binding, spec adaptation #1). +//! Provider-scoped shims over the shared router-protocol client +//! (`llm_router::provider_scaffold::router_client`): every call binds this +//! crate's `PROVIDER_ID` and carries the registration token. use crate::PROVIDER_ID; use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::router_client as scaffold; use llm_router::types::model::Model; use llm_router::types::router::ProviderResolveResponse; -use serde_json::{json, Value}; - -async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { - iii.trigger(TriggerRequest { - function_id: function_id.into(), - payload, - action: None, - timeout_ms: Some(15_000), - }) - .await -} +use serde_json::Value; /// `router::provider::resolve` — credential + effective settings. pub async fn resolve( iii: &IIIClient, token: Option<&str>, ) -> Result { - let mut payload = json!({ "id": PROVIDER_ID }); - if let Some(t) = token { - payload["token"] = json!(t); - } - let raw = call(iii, "router::provider::resolve", payload).await?; - serde_json::from_value(raw).map_err(|e| Error::Remote { - code: "provider/bad_resolve_response".into(), - message: e.to_string(), - stacktrace: None, - }) + scaffold::resolve(iii, PROVIDER_ID, token).await } /// `router::models::reconcile` — replace this provider's catalog slice. @@ -41,30 +23,15 @@ pub async fn reconcile( models: Vec, token: Option<&str>, ) -> Result<(), Error> { - let mut payload = json!({ - "provider": PROVIDER_ID, - "models": serde_json::to_value(models).expect("serializable models"), - }); - if let Some(t) = token { - payload["token"] = json!(t); - } - call(iii, "router::models::reconcile", payload).await?; - Ok(()) + scaffold::reconcile(iii, PROVIDER_ID, models, token).await } /// `router::models::get` — authoritative catalog record (None when absent). pub async fn models_get(iii: &IIIClient, model_id: &str) -> Option { - let raw = call( - iii, - "router::models::get", - json!({ "provider": PROVIDER_ID, "id": model_id }), - ) - .await - .ok()?; - serde_json::from_value(raw.get("model")?.clone()).ok() + scaffold::models_get(iii, PROVIDER_ID, model_id).await } /// `router::provider::register` — returns the registration token to persist. pub async fn register(iii: &IIIClient, declaration: Value) -> Result { - call(iii, "router::provider::register", declaration).await + scaffold::register(iii, declaration).await } diff --git a/provider-anthropic/src/sse.rs b/provider-anthropic/src/sse.rs index 7e3a8424b..284e1741a 100644 --- a/provider-anthropic/src/sse.rs +++ b/provider-anthropic/src/sse.rs @@ -344,7 +344,7 @@ pub fn handle_sse_event( .unwrap_or(""); state.text_blocks[idx].push_str(text); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta: text.to_string(), }); } @@ -355,7 +355,7 @@ pub fn handle_sse_event( .unwrap_or(""); state.function_calls[idx].args_json.push_str(json); events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta: json.to_string(), id: state.function_calls[idx].id.clone(), }); @@ -367,7 +367,7 @@ pub fn handle_sse_event( .unwrap_or(""); state.thinking_blocks[idx].text.push_str(text); events.push(AssistantMessageEvent::ThinkingDelta { - partial: build_partial(state, model), + partial: None, delta: text.to_string(), }); } @@ -465,6 +465,39 @@ mod tests { (state, events) } + /// Contract pin (llm-router types::events): delta frames are slim — + /// no cumulative partial per chunk — while block-boundary frames carry + /// the authoritative snapshot (cumulative text here; signatures and + /// final call args on their End frames). Readers reconstruct via + /// llm_router::chat::accumulate. + #[test] + fn deltas_are_slim_and_boundary_snapshots_are_cumulative() { + let (_, events) = run(&[ + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"He\"}}", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"llo\"}}", + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}", + ]); + for ev in &events { + if let AssistantMessageEvent::TextDelta { partial, .. } = ev { + assert!(partial.is_none(), "delta frames must not carry partial"); + } + } + let Some(AssistantMessageEvent::TextEnd { partial }) = events + .iter() + .find(|e| matches!(e, AssistantMessageEvent::TextEnd { .. })) + else { + panic!("want a text_end frame"); + }; + assert!( + matches!( + &partial.content[0], + llm_router::types::content::ContentBlock::Text { text } if text == "Hello" + ), + "the End snapshot must carry the cumulative block text" + ); + } + #[test] fn text_stream_produces_start_delta_end_and_final_content() { let (state, events) = run(&[ diff --git a/provider-anthropic/src/state.rs b/provider-anthropic/src/state.rs index 19e459de6..a4c417575 100644 --- a/provider-anthropic/src/state.rs +++ b/provider-anthropic/src/state.rs @@ -1,34 +1,15 @@ -//! Registration-token persistence in iii-state (engine `state::*` functions, -//! binary-worker.md § 7). The raw token lives here, under the provider's own -//! scope; the router persists only its sha256 hash. +//! Registration-token persistence, scoped to this provider's worker id +//! (shared logic in `llm_router::provider_scaffold::state`). use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; -use serde_json::{json, Value}; +use llm_router::provider_scaffold::state as scaffold; pub const STATE_SCOPE: &str = "provider-anthropic"; -const TOKEN_KEY: &str = "registration_token"; pub async fn load_token(iii: &IIIClient) -> Option { - let value = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY }), - action: None, - timeout_ms: None, - }) - .await - .ok()?; - value.as_str().map(String::from) + scaffold::load_token(iii, STATE_SCOPE).await } pub async fn store_token(iii: &IIIClient, token: &str) -> Result<(), Error> { - iii.trigger(TriggerRequest { - function_id: "state::set".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY, "value": Value::from(token) }), - action: None, - timeout_ms: None, - }) - .await?; - Ok(()) + scaffold::store_token(iii, STATE_SCOPE, token).await } diff --git a/provider-anthropic/src/stream_fn.rs b/provider-anthropic/src/stream_fn.rs index eea102e43..fee3b1544 100644 --- a/provider-anthropic/src/stream_fn.rs +++ b/provider-anthropic/src/stream_fn.rs @@ -14,26 +14,24 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; use llm_router::channels::open_sink; use llm_router::chat::relay::FrameSink; -use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::provider_scaffold::pump::{pump, send_event, PING_INTERVAL}; +use llm_router::types::events::ErrorKind; use llm_router::types::router::{ProviderStreamInput, ProviderStreamOutput}; -use std::time::Duration; -use tokio::sync::mpsc; - -/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). -pub const PING_INTERVAL: Duration = Duration::from_secs(30); pub fn make_stream( iii: IIIClient, http: reqwest::Client, + cache: ScaffoldCache, ) -> impl Fn(ProviderStreamInput) -> BoxFuture<'static, Result> + Send + Sync + 'static { move |input: ProviderStreamInput| { - let (iii, http) = (iii.clone(), http.clone()); + let (iii, http, cache) = (iii.clone(), http.clone(), cache.clone()); Box::pin(async move { let sink = open_sink(&iii, &input.writer_ref).await?; - run_stream_call(&iii, http, input, sink.as_ref()).await; + run_stream_call(&iii, http, &cache, input, sink.as_ref()).await; sink.close(); // ProviderStreamOutput (spec § stream contract) Ok(ProviderStreamOutput { ok: true }) @@ -41,14 +39,10 @@ pub fn make_stream( } } -fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { - let frame = serde_json::to_string(ev).expect("serializable event"); - sink.send(&frame).map_err(|_| ()) -} - async fn run_stream_call( iii: &IIIClient, http: reqwest::Client, + cache: &ScaffoldCache, input: ProviderStreamInput, sink: &dyn FrameSink, ) { @@ -63,16 +57,27 @@ async fn run_stream_call( ); } - let token = state::load_token(iii).await; - let resolved = match router_client::resolve(iii, token.as_deref()).await { + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + { Ok(r) => r, Err(e) => { + let kind = classify_bus_error(&e); + if kind == ErrorKind::AuthExpired { + cache.invalidate(); + } let _ = send_event( sink, &synthetic_error_event( &format!("router::provider::resolve failed: {e}"), &model, - classify_bus_error(&e), + kind, ), ); return; @@ -141,136 +146,9 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; -} - -/// Forward upstream events to the sink; ping through silence; stop on the -/// terminal event or on a failed write (caller gone → dropping `rx` aborts -/// the upstream task and its in-flight HTTP request). -pub async fn pump( - mut rx: mpsc::Receiver, - sink: &dyn FrameSink, - ping_interval: Duration, -) { - loop { - match tokio::time::timeout(ping_interval, rx.recv()).await { - Ok(Some(ev)) => { - let terminal = ev.is_terminal(); - if send_event(sink, &ev).is_err() { - return; - } - if terminal { - return; - } - } - // Upstream task ended without a terminal (panic/abort): the - // router synthesizes the terminal frame — never two terminals. - Ok(None) => return, - // Silent stretch: heartbeat (also probes for a gone caller). - Err(_elapsed) => { - if send_event(sink, &AssistantMessageEvent::Ping).is_err() { - return; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use llm_router::chat::relay::RelayRead; - use llm_router::testkit::fake_channels::FakeChannel; - use llm_router::types::messages::AssistantMessage; - use serde_json::Value; - - fn empty_assistant(model: &str) -> AssistantMessage { - llm_router::chat::synthesize::empty_partial(model, crate::PROVIDER_ID, crate::now_ms()) - } - - fn done_event() -> AssistantMessageEvent { - AssistantMessageEvent::Done { - message: empty_assistant("claude-test"), - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn forwards_events_and_stops_at_terminal() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - tx.send(done_event()).await.unwrap(); - // a frame after the terminal must never be forwarded - tx.send(AssistantMessageEvent::Ping).await.unwrap(); - drop(tx); - - pump(rx, &ch.writer, Duration::from_secs(30)).await; - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - assert_eq!(frames.len(), 2); - let last: Value = serde_json::from_str(&frames[1]).unwrap(); - assert_eq!(last["type"], "done"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn pings_through_silence() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel::(8); - // hold tx open, send nothing for > 2 ping intervals, then terminate - let pump_task = { - let writer = ch.writer.clone(); - tokio::spawn(async move { pump(rx, &writer, Duration::from_millis(50)).await }) - }; - tokio::time::sleep(Duration::from_millis(140)).await; - tx.send(done_event()).await.unwrap(); - drop(tx); - pump_task.await.unwrap(); - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - let pings = frames - .iter() - .filter(|f| serde_json::from_str::(f).unwrap()["type"] == "ping") - .count(); - assert!( - pings >= 2, - "want >=2 pings through 140ms of silence, got {pings}" - ); - assert_eq!( - serde_json::from_str::(frames.last().unwrap()).unwrap()["type"], - "done" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn reader_close_stops_the_pump_and_drops_the_receiver() { - let ch = FakeChannel::new(); - ch.reader.close(); // caller gone before anything is written - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - pump(rx, &ch.writer, Duration::from_secs(30)).await; // returns immediately - // the receiver was consumed and dropped by pump → upstream send fails - assert!(tx.send(done_event()).await.is_err()); + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); } } diff --git a/provider-anthropic/src/upstream.rs b/provider-anthropic/src/upstream.rs index fd28dca51..c45d416ad 100644 --- a/provider-anthropic/src/upstream.rs +++ b/provider-anthropic/src/upstream.rs @@ -7,6 +7,9 @@ use crate::sse::{ PartialState, }; use futures::StreamExt; +use llm_router::provider_scaffold::sse_transport::{ + append_utf8_chunk, drain_sse_blocks, error_chain, +}; use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; use serde_json::Value; use tokio::sync::mpsc; @@ -31,90 +34,19 @@ pub fn spawn_upstream( rx } -/// Flatten an error and its `source()` chain into one string. reqwest's -/// top-level Display for a builder error is just "builder error"; the real -/// cause (invalid header value, bad URL) lives in the source chain, so without -/// this the message is undiagnosable. -fn error_chain(e: &dyn std::error::Error) -> String { - let mut msg = e.to_string(); - let mut src = e.source(); - while let Some(s) = src { - let next = s.to_string(); - if !msg.ends_with(&next) { - msg.push_str(": "); - msg.push_str(&next); - } - src = s.source(); - } - msg -} - -/// Append chunk bytes to `text`, retaining any trailing incomplete UTF-8 sequence. -fn append_utf8_chunk(byte_buf: &mut Vec, text: &mut String, chunk: &[u8]) { - byte_buf.extend_from_slice(chunk); - let mut consumed = 0usize; - loop { - match std::str::from_utf8(&byte_buf[consumed..]) { - Ok(s) => { - text.push_str(s); - byte_buf.clear(); - return; - } - Err(e) => { - let valid = e.valid_up_to(); - if valid > 0 { - // SAFETY: valid_up_to guarantees valid UTF-8 in this prefix. - text.push_str(unsafe { - std::str::from_utf8_unchecked(&byte_buf[consumed..consumed + valid]) - }); - consumed += valid; - } - match e.error_len() { - Some(invalid) => { - byte_buf.drain(..consumed + invalid); - text.push('\u{FFFD}'); - consumed = 0; - } - None => { - if consumed > 0 { - byte_buf.drain(..consumed); - } - return; - } - } - } - } - } -} - -fn normalize_crlf(text: &mut String) { - if text.contains("\r\n") { - *text = text.replace("\r\n", "\n"); - } -} - -/// Drain complete `\n\n`-delimited SSE blocks from `text`. Returns true when a -/// terminal event was forwarded. -async fn drain_sse_blocks( +/// Drain complete `\n\n`-delimited SSE blocks from `text`, decoding each +/// through this provider's SSE state machine. Returns true when a terminal +/// event was forwarded. +async fn drain_blocks( text: &mut String, state: &mut PartialState, model: &str, tx: &mpsc::Sender, ) -> bool { - normalize_crlf(text); - while let Some(idx) = text.find("\n\n") { - let block: String = text.drain(..idx + 2).collect(); - for ev in handle_sse_event(&block, state, model) { - let terminal = ev.is_terminal(); - if tx.send(ev).await.is_err() { - return true; - } - if terminal { - return true; - } - } - } - false + drain_sse_blocks(text, tx, &mut |block: &str| { + handle_sse_event(block, state, model) + }) + .await } async fn run_upstream( @@ -185,7 +117,7 @@ async fn run_upstream( } }; append_utf8_chunk(&mut byte_buf, &mut text, &chunk); - if drain_sse_blocks(&mut text, &mut state, &args.model, &tx).await { + if drain_blocks(&mut text, &mut state, &args.model, &tx).await { return; } } @@ -196,7 +128,7 @@ async fn run_upstream( } if !text.trim().is_empty() { let remainder = std::mem::take(&mut text); - let _ = drain_sse_blocks(&mut (remainder + "\n\n"), &mut state, &args.model, &tx).await; + let _ = drain_blocks(&mut (remainder + "\n\n"), &mut state, &args.model, &tx).await; } if state.saw_message_stop { diff --git a/provider-anthropic/src/wire/names.rs b/provider-anthropic/src/wire/names.rs index 9440804a0..9d1b5179c 100644 --- a/provider-anthropic/src/wire/names.rs +++ b/provider-anthropic/src/wire/names.rs @@ -1,36 +1,5 @@ //! iii function ids ↔ Anthropic tool names. Anthropic enforces -//! `^[a-zA-Z0-9_-]{1,128}$`; bus ids use `::` separators. +//! `^[a-zA-Z0-9_-]{1,128}$`; bus ids use `::` separators. Shared codec +//! (and its tests) live in `llm_router::provider_scaffold::names`. -pub fn encode_tool_name(name: &str) -> String { - name.replace("::", "__") -} - -/// Inverse of `encode_tool_name`. Lossy precondition: an id containing a -/// literal `__` decodes to `::` — such ids are not in use on the bus today. -pub fn decode_tool_name(name: &str) -> String { - name.replace("__", "::") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trips_bus_ids() { - assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); - assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); - assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); - } - - #[test] - fn plain_names_pass_through() { - assert_eq!(encode_tool_name("submit_result"), "submit_result"); - assert_eq!(decode_tool_name("submit_result"), "submit_result"); - } - - #[test] - fn literal_double_underscore_is_the_accepted_lossy_case() { - // known limitation: a literal `__` in an id decodes to `::` - assert_eq!(decode_tool_name("a__b"), "a::b"); - } -} +pub use llm_router::provider_scaffold::names::{decode_tool_name, encode_tool_name}; diff --git a/provider-llamacpp/src/discovery.rs b/provider-llamacpp/src/discovery.rs index 668b99ade..1322e920f 100644 --- a/provider-llamacpp/src/discovery.rs +++ b/provider-llamacpp/src/discovery.rs @@ -7,11 +7,11 @@ //! with no `--api-key` at all — and the catalog is only pruned on an actual //! 401/403 from the server (an operator-configured key we don't have). use crate::config::{credential_parts, DEFAULT_API_URL, DEFAULT_MAX_TOKENS}; -use crate::errors::error_chain; use crate::{router_client, state, PROVIDER_ID}; use futures::future::BoxFuture; use iii_sdk::errors::Error; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::sse_transport::error_chain; use llm_router::types::model::Model; use llm_router::types::router::{RefreshModelsRequest, RefreshModelsResponse}; use serde_json::Value; diff --git a/provider-llamacpp/src/embed.rs b/provider-llamacpp/src/embed.rs index e59137389..3b49e6db9 100644 --- a/provider-llamacpp/src/embed.rs +++ b/provider-llamacpp/src/embed.rs @@ -8,11 +8,14 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::types::events::ErrorKind; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use crate::config::{config_from_resolve, ConfigError}; -use crate::{router_client, state}; +use crate::errors::{classify, classify_bus_error}; +use crate::state; /// Embeddings requests are bounded and non-streaming; a tight budget keeps /// hook callers (memory recall) inside their own timeouts. @@ -72,6 +75,7 @@ fn embed_url(chat_api_url: &str) -> String { pub async fn handle( iii: &IIIClient, http: &reqwest::Client, + cache: &ScaffoldCache, req: EmbedRequest, ) -> Result { if req.input.is_empty() || req.input.len() > 512 { @@ -84,8 +88,19 @@ pub async fn handle( .filter(|m| !m.trim().is_empty()) .unwrap_or_else(|| "default".to_string()); - let token = state::load_token(iii).await; - let resolved = router_client::resolve(iii, token.as_deref()).await?; + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + .inspect_err(|e| { + if classify_bus_error(e) == ErrorKind::AuthExpired { + cache.invalidate(); + } + })?; let cfg = config_from_resolve(&model, None, &resolved).map_err(|e| match e { ConfigError::InvalidApiUrl(u) => { Error::Handler(format!("provider/config: invalid endpoint url {u:?}")) @@ -109,6 +124,11 @@ pub async fn handle( let status = response.status(); if !status.is_success() { let body = response.text().await.unwrap_or_default(); + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if classify(Some(status.as_u16()), &body) == ErrorKind::AuthExpired { + cache.invalidate(); + } let excerpt: String = body.chars().take(300).collect(); return Err(Error::Handler(format!( "provider/upstream_status: {status}: {excerpt} \ diff --git a/provider-llamacpp/src/errors.rs b/provider-llamacpp/src/errors.rs index b3baca42f..698c367c8 100644 --- a/provider-llamacpp/src/errors.rs +++ b/provider-llamacpp/src/errors.rs @@ -106,26 +106,6 @@ pub fn invalid_request_from_serde(e: serde_json::Error) -> Error { invalid_request(format!("bad ProviderStreamInput: {e}")) } -/// Flatten an error and its `source()` chain into one string. reqwest's -/// top-level Display is often opaque ("builder error", "error sending request -/// for url (…)"); the real cause (invalid header value, connection refused, -/// "No route to host") lives in the source chain, so without this the message -/// is undiagnosable. -pub(crate) fn error_chain(e: &dyn std::error::Error) -> String { - let mut msg = e.to_string(); - let mut src = e.source(); - while let Some(s) = src { - let next = s.to_string(); - // reqwest sometimes nests the same text; skip exact repeats. - if !msg.ends_with(&next) { - msg.push_str(": "); - msg.push_str(&next); - } - src = s.source(); - } - msg -} - #[cfg(test)] mod tests { use super::*; diff --git a/provider-llamacpp/src/register.rs b/provider-llamacpp/src/register.rs index e3c5ca877..fe79aa0e0 100644 --- a/provider-llamacpp/src/register.rs +++ b/provider-llamacpp/src/register.rs @@ -9,6 +9,7 @@ use crate::{router_client, state, PROVIDER_ID}; use iii_sdk::errors::Error; use iii_sdk::protocol::RegisterTriggerInput; use iii_sdk::{IIIClient, RegisterFunction}; +use llm_router::provider_scaffold::cache::ScaffoldCache; use llm_router::types::router::{ ProviderDeclaration, ProviderDefaults, ProviderReadyAck, RouterReadyEvent, }; @@ -111,6 +112,11 @@ pub async fn declare_and_refresh(iii: IIIClient, http: reqwest::Client) { } pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { + // Shared per-process cache for the registration token and the resolve + // response (see llm_router::provider_scaffold::cache). Invalidated on + // router::ready — a restarted router may carry new config and reissues + // declare/refresh anyway — and on upstream auth errors (stream_fn). + let cache = ScaffoldCache::new(); // Streaming uses no total timeout (the router owns stream budgets); // connect failures surface fast. let http = reqwest::Client::builder() @@ -121,7 +127,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { iii.register_function( surface::STREAM_ID, RegisterFunction::new_async_with_bad_request( - make_stream(iii.clone(), http.clone()), + make_stream(iii.clone(), http.clone(), cache.clone()), invalid_request_from_serde, ) .description(surface::STREAM_DESC), @@ -136,11 +142,13 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { { let iii_ready = iii.clone(); let http_ready = http.clone(); + let cache_ready = cache.clone(); iii.register_function( surface::ON_ROUTER_READY_ID, RegisterFunction::new_async(move |_event: RouterReadyEvent| { let iii = iii_ready.clone(); let http = http_ready.clone(); + cache_ready.invalidate(); async move { tokio::spawn(declare_and_refresh(iii, http)); Ok::<_, Error>(ProviderReadyAck { ok: true }) @@ -153,11 +161,13 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { { let iii_embed = iii.clone(); let http_embed = http.clone(); + let cache_embed = cache.clone(); iii.register_function( surface::EMBED_ID, RegisterFunction::new_async(move |req: crate::embed::EmbedRequest| { - let (iii, http) = (iii_embed.clone(), http_embed.clone()); - async move { crate::embed::handle(&iii, &http, req).await } + let (iii, http, cache) = + (iii_embed.clone(), http_embed.clone(), cache_embed.clone()); + async move { crate::embed::handle(&iii, &http, &cache, req).await } }) .description(surface::EMBED_DESC) .metadata(json!({ "internal": true })), diff --git a/provider-llamacpp/src/router_client.rs b/provider-llamacpp/src/router_client.rs index 373c350e5..175f9e844 100644 --- a/provider-llamacpp/src/router_client.rs +++ b/provider-llamacpp/src/router_client.rs @@ -1,38 +1,20 @@ -//! Thin wrappers over the router's provider-protocol functions. All calls -//! carry the registration token (identity binding, spec adaptation #1). +//! Provider-scoped shims over the shared router-protocol client +//! (`llm_router::provider_scaffold::router_client`): every call binds this +//! crate's `PROVIDER_ID` and carries the registration token. use crate::PROVIDER_ID; use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::router_client as scaffold; use llm_router::types::model::Model; use llm_router::types::router::ProviderResolveResponse; -use serde_json::{json, Value}; - -async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { - iii.trigger(TriggerRequest { - function_id: function_id.into(), - payload, - action: None, - timeout_ms: Some(15_000), - }) - .await -} +use serde_json::Value; /// `router::provider::resolve` — credential + effective settings. pub async fn resolve( iii: &IIIClient, token: Option<&str>, ) -> Result { - let mut payload = json!({ "id": PROVIDER_ID }); - if let Some(t) = token { - payload["token"] = json!(t); - } - let raw = call(iii, "router::provider::resolve", payload).await?; - serde_json::from_value(raw).map_err(|e| Error::Remote { - code: "provider/bad_resolve_response".into(), - message: e.to_string(), - stacktrace: None, - }) + scaffold::resolve(iii, PROVIDER_ID, token).await } /// `router::models::reconcile` — replace this provider's catalog slice. @@ -41,30 +23,15 @@ pub async fn reconcile( models: Vec, token: Option<&str>, ) -> Result<(), Error> { - let mut payload = json!({ - "provider": PROVIDER_ID, - "models": serde_json::to_value(models).expect("serializable models"), - }); - if let Some(t) = token { - payload["token"] = json!(t); - } - call(iii, "router::models::reconcile", payload).await?; - Ok(()) + scaffold::reconcile(iii, PROVIDER_ID, models, token).await } /// `router::models::get` — authoritative catalog record (None when absent). pub async fn models_get(iii: &IIIClient, model_id: &str) -> Option { - let raw = call( - iii, - "router::models::get", - json!({ "provider": PROVIDER_ID, "id": model_id }), - ) - .await - .ok()?; - serde_json::from_value(raw.get("model")?.clone()).ok() + scaffold::models_get(iii, PROVIDER_ID, model_id).await } /// `router::provider::register` — returns the registration token to persist. pub async fn register(iii: &IIIClient, declaration: Value) -> Result { - call(iii, "router::provider::register", declaration).await + scaffold::register(iii, declaration).await } diff --git a/provider-llamacpp/src/sse.rs b/provider-llamacpp/src/sse.rs index 375420677..c1f117d11 100644 --- a/provider-llamacpp/src/sse.rs +++ b/provider-llamacpp/src/sse.rs @@ -261,7 +261,7 @@ pub fn handle_chunk( } state.thinking.push_str(reasoning); events.push(AssistantMessageEvent::ThinkingDelta { - partial: build_partial(state, model), + partial: None, delta: reasoning.to_string(), }); } @@ -277,7 +277,7 @@ pub fn handle_chunk( } state.text.push_str(text); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta: text.to_string(), }); } @@ -316,7 +316,7 @@ pub fn handle_chunk( if !args.is_empty() { state.function_calls[index].args_json.push_str(args); events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta: args.to_string(), id: state.function_calls[index].id.clone(), }); @@ -368,6 +368,34 @@ mod tests { .collect() } + /// Contract pin (llm-router types::events): delta frames are slim — + /// no cumulative partial per chunk — while block-boundary frames carry + /// the authoritative snapshot (cumulative text here). Readers + /// reconstruct via llm_router::chat::accumulate. + #[test] + fn deltas_are_slim_and_boundary_snapshots_are_cumulative() { + let (_, events) = run(&[ + json!({"choices":[{"index":0,"delta":{"content":"He"}}]}), + json!({"choices":[{"index":0,"delta":{"content":"llo"}}]}), + json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}), + ]); + for ev in &events { + if let AssistantMessageEvent::TextDelta { partial, .. } = ev { + assert!(partial.is_none(), "delta frames must not carry partial"); + } + } + let Some(AssistantMessageEvent::TextEnd { partial }) = events + .iter() + .find(|e| matches!(e, AssistantMessageEvent::TextEnd { .. })) + else { + panic!("want a text_end frame"); + }; + assert!( + matches!(&partial.content[0], ContentBlock::Text { text } if text == "Hello"), + "the End snapshot must carry the cumulative block text" + ); + } + #[test] fn text_stream_produces_start_delta_end_and_final_content() { let (state, events) = run(&[ diff --git a/provider-llamacpp/src/state.rs b/provider-llamacpp/src/state.rs index 52824e9c7..ac3fcb493 100644 --- a/provider-llamacpp/src/state.rs +++ b/provider-llamacpp/src/state.rs @@ -1,34 +1,15 @@ -//! Registration-token persistence in iii-state (engine `state::*` functions, -//! binary-worker.md § 7). The raw token lives here, under the provider's own -//! scope; the router persists only its sha256 hash. +//! Registration-token persistence, scoped to this provider's worker id +//! (shared logic in `llm_router::provider_scaffold::state`). use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; -use serde_json::{json, Value}; +use llm_router::provider_scaffold::state as scaffold; pub const STATE_SCOPE: &str = "provider-llamacpp"; -const TOKEN_KEY: &str = "registration_token"; pub async fn load_token(iii: &IIIClient) -> Option { - let value = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY }), - action: None, - timeout_ms: None, - }) - .await - .ok()?; - value.as_str().map(String::from) + scaffold::load_token(iii, STATE_SCOPE).await } pub async fn store_token(iii: &IIIClient, token: &str) -> Result<(), Error> { - iii.trigger(TriggerRequest { - function_id: "state::set".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY, "value": Value::from(token) }), - action: None, - timeout_ms: None, - }) - .await?; - Ok(()) + scaffold::store_token(iii, STATE_SCOPE, token).await } diff --git a/provider-llamacpp/src/stream_fn.rs b/provider-llamacpp/src/stream_fn.rs index 1842b5b5e..1404e0dac 100644 --- a/provider-llamacpp/src/stream_fn.rs +++ b/provider-llamacpp/src/stream_fn.rs @@ -5,33 +5,31 @@ use crate::config::config_from_resolve; use crate::errors::classify_bus_error; use crate::request::{build_body, build_headers, BodyArgs}; use crate::sse::synthetic_error_event; +use crate::state; use crate::upstream::{spawn_upstream, UpstreamArgs}; -use crate::{router_client, state}; use futures::future::BoxFuture; use iii_sdk::errors::Error; use iii_sdk::IIIClient; use llm_router::channels::open_sink; use llm_router::chat::relay::FrameSink; -use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::provider_scaffold::pump::{pump, send_event, PING_INTERVAL}; +use llm_router::types::events::ErrorKind; use llm_router::types::router::{ProviderStreamInput, ProviderStreamOutput}; -use std::time::Duration; -use tokio::sync::mpsc; - -/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). -pub const PING_INTERVAL: Duration = Duration::from_secs(30); pub fn make_stream( iii: IIIClient, http: reqwest::Client, + cache: ScaffoldCache, ) -> impl Fn(ProviderStreamInput) -> BoxFuture<'static, Result> + Send + Sync + 'static { move |input: ProviderStreamInput| { - let (iii, http) = (iii.clone(), http.clone()); + let (iii, http, cache) = (iii.clone(), http.clone(), cache.clone()); Box::pin(async move { let sink = open_sink(&iii, &input.writer_ref).await?; - run_stream_call(&iii, http, input, sink.as_ref()).await; + run_stream_call(&iii, http, &cache, input, sink.as_ref()).await; sink.close(); // ProviderStreamOutput (spec § stream contract) Ok(ProviderStreamOutput { ok: true }) @@ -39,30 +37,37 @@ pub fn make_stream( } } -fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { - let frame = serde_json::to_string(ev).expect("serializable event"); - sink.send(&frame).map_err(|_| ()) -} - async fn run_stream_call( iii: &IIIClient, http: reqwest::Client, + cache: &ScaffoldCache, input: ProviderStreamInput, sink: &dyn FrameSink, ) { let model = input.model.clone(); let mut warnings = Vec::new(); - let token = state::load_token(iii).await; - let resolved = match router_client::resolve(iii, token.as_deref()).await { + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + { Ok(r) => r, Err(e) => { + let kind = classify_bus_error(&e); + if kind == ErrorKind::AuthExpired { + cache.invalidate(); + } let _ = send_event( sink, &synthetic_error_event( &format!("router::provider::resolve failed: {e}"), &model, - classify_bus_error(&e), + kind, ), ); return; @@ -117,134 +122,9 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; -} - -/// Forward upstream events to the sink; ping through silence; stop on the -/// terminal event or on a failed write (caller gone → dropping `rx` aborts -/// the upstream task and its in-flight HTTP request). -/// Verbatim copy of provider-anthropic's pump — shared extraction into -/// llm-router is a listed follow-up. -pub async fn pump( - mut rx: mpsc::Receiver, - sink: &dyn FrameSink, - ping_interval: Duration, -) { - loop { - match tokio::time::timeout(ping_interval, rx.recv()).await { - Ok(Some(ev)) => { - let terminal = ev.is_terminal(); - if send_event(sink, &ev).is_err() { - return; - } - if terminal { - return; - } - } - // Upstream task ended without a terminal (panic/abort): the - // router synthesizes the terminal frame — never two terminals. - Ok(None) => return, - // Silent stretch: heartbeat (also probes for a gone caller). - Err(_elapsed) => { - if send_event(sink, &AssistantMessageEvent::Ping).is_err() { - return; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::sse::empty_assistant; - use llm_router::chat::relay::RelayRead; - use llm_router::testkit::fake_channels::FakeChannel; - use serde_json::Value; - - fn done_event() -> AssistantMessageEvent { - AssistantMessageEvent::Done { - message: empty_assistant("llama-test"), - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn forwards_events_and_stops_at_terminal() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - tx.send(done_event()).await.unwrap(); - // a frame after the terminal must never be forwarded - tx.send(AssistantMessageEvent::Ping).await.unwrap(); - drop(tx); - - pump(rx, &ch.writer, Duration::from_secs(30)).await; - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - assert_eq!(frames.len(), 2); - let last: Value = serde_json::from_str(&frames[1]).unwrap(); - assert_eq!(last["type"], "done"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn pings_through_silence() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel::(8); - // hold tx open, send nothing for > 2 ping intervals, then terminate - let pump_task = { - let writer = ch.writer.clone(); - tokio::spawn(async move { pump(rx, &writer, Duration::from_millis(50)).await }) - }; - tokio::time::sleep(Duration::from_millis(140)).await; - tx.send(done_event()).await.unwrap(); - drop(tx); - pump_task.await.unwrap(); - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - let pings = frames - .iter() - .filter(|f| serde_json::from_str::(f).unwrap()["type"] == "ping") - .count(); - assert!( - pings >= 2, - "want >=2 pings through 140ms of silence, got {pings}" - ); - assert_eq!( - serde_json::from_str::(frames.last().unwrap()).unwrap()["type"], - "done" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn reader_close_stops_the_pump_and_drops_the_receiver() { - let ch = FakeChannel::new(); - ch.reader.close(); // caller gone before anything is written - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - pump(rx, &ch.writer, Duration::from_secs(30)).await; // returns immediately - // the receiver was consumed and dropped by pump → upstream send fails - assert!(tx.send(done_event()).await.is_err()); + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); } } diff --git a/provider-llamacpp/src/upstream.rs b/provider-llamacpp/src/upstream.rs index 5bfff3171..3a46bd7a6 100644 --- a/provider-llamacpp/src/upstream.rs +++ b/provider-llamacpp/src/upstream.rs @@ -2,9 +2,12 @@ //! mpsc. //! The receiver dropping aborts the upstream: every send error returns, //! which drops the reqwest response mid-body and closes the connection. -use crate::errors::{classify, error_chain}; +use crate::errors::classify; use crate::sse::{build_final, build_partial, handle_chunk, synthetic_error_event, PartialState}; use futures::StreamExt; +use llm_router::provider_scaffold::sse_transport::{ + append_utf8_chunk, drain_sse_blocks, error_chain, +}; use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; use serde_json::Value; use tokio::sync::mpsc; @@ -29,48 +32,6 @@ pub fn spawn_upstream( rx } -/// Append chunk bytes to `text`, retaining any trailing incomplete UTF-8 -/// sequence: transport chunk boundaries are arbitrary, so a multi-byte -/// character (a locally-hosted model's output may be CJK or other non-ASCII -/// text) can split across chunks — a -/// per-chunk lossy decode would corrupt it into U+FFFD. -fn append_utf8_chunk(byte_buf: &mut Vec, text: &mut String, chunk: &[u8]) { - byte_buf.extend_from_slice(chunk); - let mut consumed = 0usize; - loop { - match std::str::from_utf8(&byte_buf[consumed..]) { - Ok(s) => { - text.push_str(s); - byte_buf.clear(); - return; - } - Err(e) => { - let valid = e.valid_up_to(); - if valid > 0 { - // SAFETY: valid_up_to guarantees valid UTF-8 in this prefix. - text.push_str(unsafe { - std::str::from_utf8_unchecked(&byte_buf[consumed..consumed + valid]) - }); - consumed += valid; - } - match e.error_len() { - Some(invalid) => { - byte_buf.drain(..consumed + invalid); - text.push('\u{FFFD}'); - consumed = 0; - } - None => { - if consumed > 0 { - byte_buf.drain(..consumed); - } - return; - } - } - } - } - } -} - /// Last `data: ` payload in an SSE block, if any. fn data_line(block: &str) -> Option<&str> { block @@ -130,7 +91,35 @@ async fn run_upstream( let mut stream = resp.bytes_stream(); let mut buf = String::new(); + // Cross-chunk UTF-8 buffering: transport chunk boundaries are arbitrary, + // so a multi-byte character (a locally-hosted model's output may be CJK + // or other non-ASCII text) can split across chunks — a per-chunk lossy + // decode would corrupt it into U+FFFD. let mut byte_buf: Vec = Vec::new(); + // Block decoder: [DONE] closes the stream (Stop + Done, Done terminal); + // anything else parses and runs the chunk state machine. + let decode = + |block: &str, state: &mut PartialState, model: &str| -> Vec { + let Some(data) = data_line(block) else { + return vec![]; + }; + if data == "[DONE]" { + return vec![ + AssistantMessageEvent::Stop { + stop_reason: state.stop_reason(), + error_message: None, + error_kind: None, + }, + AssistantMessageEvent::Done { + message: build_final(state, model), + }, + ]; + } + let Ok(parsed) = serde_json::from_str::(data) else { + return vec![]; + }; + handle_chunk(&parsed, state, model) + }; while let Some(chunk) = stream.next().await { let chunk = match chunk { Ok(c) => c, @@ -146,38 +135,12 @@ async fn run_upstream( } }; append_utf8_chunk(&mut byte_buf, &mut buf, &chunk); - while let Some(idx) = buf.find("\n\n") { - let block: String = buf.drain(..idx + 2).collect(); - let Some(data) = data_line(&block) else { - continue; - }; - if data == "[DONE]" { - let _ = tx - .send(AssistantMessageEvent::Stop { - stop_reason: state.stop_reason(), - error_message: None, - error_kind: None, - }) - .await; - let _ = tx - .send(AssistantMessageEvent::Done { - message: build_final(&state, &args.model), - }) - .await; - return; - } - let Ok(parsed) = serde_json::from_str::(data) else { - continue; - }; - for ev in handle_chunk(&parsed, &mut state, &args.model) { - let terminal = ev.is_terminal(); - if tx.send(ev).await.is_err() { - return; // receiver dropped → abort upstream - } - if terminal { - return; // exactly one terminal event - } - } + if drain_sse_blocks(&mut buf, &tx, &mut |block: &str| { + decode(block, &mut state, &args.model) + }) + .await + { + return; // terminal forwarded, or receiver dropped → abort upstream } } // Stream ended without [DONE] (connection close framing): still terminal. @@ -334,24 +297,6 @@ mod tests { } } - #[test] - fn utf8_split_across_chunks_survives_reassembly() { - // "你" = 3 bytes; split it across two transport chunks. - let bytes = "data: 你好\n\n".as_bytes(); - let (a, b) = bytes.split_at(8); // mid-character - let (mut byte_buf, mut text) = (Vec::new(), String::new()); - append_utf8_chunk(&mut byte_buf, &mut text, a); - append_utf8_chunk(&mut byte_buf, &mut text, b); - assert_eq!(text, "data: 你好\n\n"); - assert!(!text.contains('\u{FFFD}'), "lossy corruption: {text:?}"); - assert!(byte_buf.is_empty()); - - // Truly invalid bytes still degrade to U+FFFD instead of stalling. - let (mut byte_buf, mut text) = (Vec::new(), String::new()); - append_utf8_chunk(&mut byte_buf, &mut text, &[b'a', 0xFF, b'b']); - assert_eq!(text, "a\u{FFFD}b"); - } - #[tokio::test(flavor = "multi_thread")] async fn warnings_arrive_on_the_final_message() { let url = stub(HAPPY).await; diff --git a/provider-llamacpp/src/wire/names.rs b/provider-llamacpp/src/wire/names.rs index be800a723..64e0b0d2c 100644 --- a/provider-llamacpp/src/wire/names.rs +++ b/provider-llamacpp/src/wire/names.rs @@ -1,29 +1,6 @@ //! iii function ids ↔ OpenAI-style tool function names. `::` is not a valid //! character in most chat-template tool-name grammars, so bus ids get -//! sanitized to `__` on the wire. +//! sanitized to `__` on the wire. Shared codec (and its tests) live in +//! `llm_router::provider_scaffold::names`. -pub fn encode_tool_name(name: &str) -> String { - name.replace("::", "__") -} - -pub fn decode_tool_name(name: &str) -> String { - name.replace("__", "::") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trips_bus_ids() { - assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); - assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); - assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); - } - - #[test] - fn plain_names_pass_through() { - assert_eq!(encode_tool_name("submit_result"), "submit_result"); - assert_eq!(decode_tool_name("submit_result"), "submit_result"); - } -} +pub use llm_router::provider_scaffold::names::{decode_tool_name, encode_tool_name}; diff --git a/provider-llamacpp/tests/integration.rs b/provider-llamacpp/tests/integration.rs index 1afc3acfe..e2a4c3f27 100644 --- a/provider-llamacpp/tests/integration.rs +++ b/provider-llamacpp/tests/integration.rs @@ -311,15 +311,41 @@ async fn configure_stub_no_key(router_iii: &IIIClient, stub_url: &str) { .expect("config set"); } +/// Call `provider::llamacpp::refresh_models`, retrying while the router's +/// in-memory configuration snapshot has not yet absorbed the test's +/// `configuration::set` (the engine delivers the configuration trigger +/// asynchronously, so an immediate refresh can still resolve the pre-set +/// default api_url and fail to connect). Unlike the keyed cloud providers — +/// whose refresh no-ops until a credential is visible — llamacpp's +/// credential-less refresh actually dials the default URL, so the +/// propagation window surfaces as a transient error here. The retry never +/// masks a real failure: discovery against the stub either succeeds within +/// the deadline or the last error is surfaced. +async fn refresh_models(provider_iii: &IIIClient) -> Value { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + match call( + provider_iii, + "provider::llamacpp::refresh_models", + json!({}), + ) + .await + { + Ok(res) => return res, + Err(e) => { + assert!( + Instant::now() < deadline, + "refresh_models kept failing (configuration never propagated?): {e:?}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + } +} + /// Discover the live catalog and wait until routing can see it. async fn refresh_and_wait(router_iii: &IIIClient, provider_iii: &IIIClient, expect_id: &str) { - let res = call( - provider_iii, - "provider::llamacpp::refresh_models", - json!({}), - ) - .await - .expect("refresh succeeds"); + let res = refresh_models(provider_iii).await; assert_eq!(res["ok"], true, "refresh response: {res}"); let deadline = Instant::now() + Duration::from_secs(10); loop { @@ -358,13 +384,7 @@ async fn provider_registers_with_persisted_token_and_discovers_catalog_without_a // Unlike every cloud provider here, no credential is required at all: // discovery succeeds and reconciles the stub's one model. - let res = call( - &provider_iii, - "provider::llamacpp::refresh_models", - json!({}), - ) - .await - .expect("refresh succeeds without a credential"); + let res = refresh_models(&provider_iii).await; assert_eq!(res["ok"], true, "refresh response: {res}"); assert_eq!( res["count"], 1, diff --git a/provider-openai-codex/Cargo.lock b/provider-openai-codex/Cargo.lock index 618328445..98902bc93 100644 --- a/provider-openai-codex/Cargo.lock +++ b/provider-openai-codex/Cargo.lock @@ -776,7 +776,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.1.1" +version = "1.1.3" dependencies = [ "async-trait", "clap", diff --git a/provider-openai-codex/src/register.rs b/provider-openai-codex/src/register.rs index 04262d56e..b83c7bd7d 100644 --- a/provider-openai-codex/src/register.rs +++ b/provider-openai-codex/src/register.rs @@ -11,6 +11,7 @@ use crate::{auth, router_client, state, PROVIDER_ID}; use iii_sdk::errors::Error; use iii_sdk::protocol::RegisterTriggerInput; use iii_sdk::{IIIClient, RegisterFunction}; +use llm_router::provider_scaffold::cache::ScaffoldCache; use llm_router::types::router::{ ProviderDeclaration, ProviderDefaults, ProviderReadyAck, RouterReadyEvent, }; @@ -125,6 +126,11 @@ fn read_timeout() -> Duration { } pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { + // Shared per-process cache for the registration token and the resolve + // response (see llm_router::provider_scaffold::cache). Invalidated on + // router::ready — a restarted router may carry new config and reissues + // declare/refresh anyway — and on upstream auth errors (stream_fn). + let cache = ScaffoldCache::new(); // Reads are silence-bounded: a stalled upstream otherwise pings the router // past its idle guard until the engine kills the call at stream_timeout. let http = reqwest::Client::builder() @@ -137,7 +143,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { iii.register_function( surface::STREAM_ID, RegisterFunction::new_async_with_bad_request( - make_stream(iii.clone(), http.clone()), + make_stream(iii.clone(), http.clone(), cache.clone()), invalid_request_from_serde, ) .description(surface::STREAM_DESC) @@ -158,6 +164,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { let iii_ready = iii.clone(); let http_ready = http.clone(); let refresh_state_ready = refresh_state.clone(); + let cache_ready = cache.clone(); iii.register_function( surface::ON_ROUTER_READY_ID, RegisterFunction::new_async(move |_event: RouterReadyEvent| { @@ -166,6 +173,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { http_ready.clone(), refresh_state_ready.clone(), ); + cache_ready.invalidate(); async move { tokio::spawn(declare_and_refresh(iii, http, refresh_state)); Ok::<_, Error>(ProviderReadyAck { ok: true }) diff --git a/provider-openai-codex/src/router_client.rs b/provider-openai-codex/src/router_client.rs index 6cf983c0b..b082daec8 100644 --- a/provider-openai-codex/src/router_client.rs +++ b/provider-openai-codex/src/router_client.rs @@ -1,13 +1,14 @@ -//! Thin wrappers over the router's provider-protocol functions (register / -//! resolve / reconcile — carrying the registration token) AND the +//! Provider-scoped shims over the shared router-protocol client +//! (`llm_router::provider_scaffold::router_client` — every call binds this +//! crate's `PROVIDER_ID` and carries the registration token) AND the //! `auth-credentials` vault + `oauth-openai-codex` refresh. Credentials come //! from the vault, never from the router config: this provider is a dumb token //! consumer (login + refresh live out-of-band). use crate::PROVIDER_ID; use iii_sdk::engine::EngineFunctions; use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::router_client::{self as scaffold, call}; use llm_router::types::model::Model; use llm_router::types::router::ProviderResolveResponse; use serde_json::{json, Value}; @@ -15,16 +16,6 @@ use serde_json::{json, Value}; const AUTH_GET_TOKEN_FN: &str = "auth::get_token"; const AUTH_SET_TOKEN_FN: &str = "auth::set_token"; -async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { - iii.trigger(TriggerRequest { - function_id: function_id.into(), - payload, - action: None, - timeout_ms: Some(15_000), - }) - .await -} - fn function_prefix(function_id: &str) -> &str { function_id .split_once("::") @@ -75,16 +66,7 @@ pub async fn resolve( iii: &IIIClient, token: Option<&str>, ) -> Result { - let mut payload = json!({ "id": PROVIDER_ID }); - if let Some(t) = token { - payload["token"] = json!(t); - } - let raw = call(iii, "router::provider::resolve", payload).await?; - serde_json::from_value(raw).map_err(|e| Error::Remote { - code: "provider/bad_resolve_response".into(), - message: e.to_string(), - stacktrace: None, - }) + scaffold::resolve(iii, PROVIDER_ID, token).await } /// `router::models::reconcile` — replace this provider's catalog slice. @@ -93,32 +75,17 @@ pub async fn reconcile( models: Vec, token: Option<&str>, ) -> Result<(), Error> { - let mut payload = json!({ - "provider": PROVIDER_ID, - "models": serde_json::to_value(models).expect("serializable models"), - }); - if let Some(t) = token { - payload["token"] = json!(t); - } - call(iii, "router::models::reconcile", payload).await?; - Ok(()) + scaffold::reconcile(iii, PROVIDER_ID, models, token).await } /// `router::models::get` — authoritative catalog record (None when absent). pub async fn models_get(iii: &IIIClient, model_id: &str) -> Option { - let raw = call( - iii, - "router::models::get", - json!({ "provider": PROVIDER_ID, "id": model_id }), - ) - .await - .ok()?; - serde_json::from_value(raw.get("model")?.clone()).ok() + scaffold::models_get(iii, PROVIDER_ID, model_id).await } /// `router::provider::register` — returns the registration token to persist. pub async fn register(iii: &IIIClient, declaration: Value) -> Result { - call(iii, "router::provider::register", declaration).await + scaffold::register(iii, declaration).await } // ── auth-credentials vault ────────────────────────────────────────────────── diff --git a/provider-openai-codex/src/sse.rs b/provider-openai-codex/src/sse.rs index 630d63c40..63ca686b7 100644 --- a/provider-openai-codex/src/sse.rs +++ b/provider-openai-codex/src/sse.rs @@ -296,7 +296,7 @@ pub fn handle_chunk( open_text(state, model, &mut events); state.text.push_str(&delta); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta, }); } @@ -308,7 +308,7 @@ pub fn handle_chunk( open_thinking(state, model, &mut events); state.thinking.push_str(&delta); events.push(AssistantMessageEvent::ThinkingDelta { - partial: build_partial(state, model), + partial: None, delta, }); } @@ -326,7 +326,7 @@ pub fn handle_chunk( last.args_json.push_str(&delta); } events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta, id: state .tool_calls @@ -394,6 +394,34 @@ mod tests { (state, out) } + /// Contract pin (llm-router types::events): delta frames are slim — + /// no cumulative partial per chunk — while block-boundary frames carry + /// the authoritative snapshot (cumulative text here). Readers + /// reconstruct via llm_router::chat::accumulate. + #[test] + fn deltas_are_slim_and_boundary_snapshots_are_cumulative() { + let (_, events) = run(&[ + json!({ "type": "response.output_text.delta", "delta": "He" }), + json!({ "type": "response.output_text.delta", "delta": "llo" }), + json!({ "type": "response.completed" }), + ]); + for ev in &events { + if let AssistantMessageEvent::TextDelta { partial, .. } = ev { + assert!(partial.is_none(), "delta frames must not carry partial"); + } + } + let Some(AssistantMessageEvent::TextEnd { partial }) = events + .iter() + .find(|e| matches!(e, AssistantMessageEvent::TextEnd { .. })) + else { + panic!("want a text_end frame"); + }; + assert!( + matches!(&partial.content[0], ContentBlock::Text { text } if text == "Hello"), + "the End snapshot must carry the cumulative block text" + ); + } + #[test] fn text_stream_yields_start_delta_and_completed_done() { let (state, events) = run(&[ diff --git a/provider-openai-codex/src/state.rs b/provider-openai-codex/src/state.rs index 30ccc37b5..e839edd45 100644 --- a/provider-openai-codex/src/state.rs +++ b/provider-openai-codex/src/state.rs @@ -1,34 +1,15 @@ -//! Registration-token persistence in iii-state (engine `state::*` functions, -//! binary-worker.md § 7). The raw token lives here, under the provider's own -//! scope; the router persists only its sha256 hash. +//! Registration-token persistence, scoped to this provider's worker id +//! (shared logic in `llm_router::provider_scaffold::state`). use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; -use serde_json::{json, Value}; +use llm_router::provider_scaffold::state as scaffold; pub const STATE_SCOPE: &str = "provider-openai-codex"; -const TOKEN_KEY: &str = "registration_token"; pub async fn load_token(iii: &IIIClient) -> Option { - let value = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY }), - action: None, - timeout_ms: None, - }) - .await - .ok()?; - value.as_str().map(String::from) + scaffold::load_token(iii, STATE_SCOPE).await } pub async fn store_token(iii: &IIIClient, token: &str) -> Result<(), Error> { - iii.trigger(TriggerRequest { - function_id: "state::set".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY, "value": Value::from(token) }), - action: None, - timeout_ms: None, - }) - .await?; - Ok(()) + scaffold::store_token(iii, STATE_SCOPE, token).await } diff --git a/provider-openai-codex/src/stream_fn.rs b/provider-openai-codex/src/stream_fn.rs index 6c2bb13f3..b4a2591b2 100644 --- a/provider-openai-codex/src/stream_fn.rs +++ b/provider-openai-codex/src/stream_fn.rs @@ -4,6 +4,7 @@ //! live in the oauth-openai-codex worker / auth-credentials vault — this //! provider only *triggers* a refresh when the token is near expiry. use crate::config::build_config; +use crate::errors::classify_bus_error; use crate::reasoning::{is_reasoning_model, native_reasoning_effort, reasoning_effort_for}; use crate::request::{build_body, build_headers, BodyArgs}; use crate::sse::synthetic_error_event; @@ -14,39 +15,32 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; use llm_router::channels::open_sink; use llm_router::chat::relay::FrameSink; -use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::provider_scaffold::pump::{pump, send_event, PING_INTERVAL}; +use llm_router::types::events::ErrorKind; use llm_router::types::router::{ CredentialSource, ProviderResolveResponse, ProviderStreamInput, ProviderStreamOutput, }; -use std::time::Duration; -use tokio::sync::mpsc; - -/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). -pub const PING_INTERVAL: Duration = Duration::from_secs(30); pub fn make_stream( iii: IIIClient, http: reqwest::Client, + cache: ScaffoldCache, ) -> impl Fn(ProviderStreamInput) -> BoxFuture<'static, Result> + Send + Sync + 'static { move |input: ProviderStreamInput| { - let (iii, http) = (iii.clone(), http.clone()); + let (iii, http, cache) = (iii.clone(), http.clone(), cache.clone()); Box::pin(async move { let sink = open_sink(&iii, &input.writer_ref).await?; - run_stream_call(&iii, http, input, sink.as_ref()).await; + run_stream_call(&iii, http, &cache, input, sink.as_ref()).await; sink.close(); Ok(ProviderStreamOutput { ok: true }) }) } } -fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { - let frame = serde_json::to_string(ev).expect("serializable event"); - sink.send(&frame).map_err(|_| ()) -} - fn default_resolve() -> ProviderResolveResponse { ProviderResolveResponse { configured: false, @@ -60,17 +54,33 @@ fn default_resolve() -> ProviderResolveResponse { async fn run_stream_call( iii: &IIIClient, http: reqwest::Client, + cache: &ScaffoldCache, input: ProviderStreamInput, sink: &dyn FrameSink, ) { let model = input.model.clone(); // router id (e.g. codex/gpt-5.5) let mut warnings = Vec::new(); - let token = state::load_token(iii).await; + // Token + resolve are cached (ScaffoldCache): zero engine round trips on + // the hot path within the TTL. The vault credential lookup below is NOT + // cached — the vault refreshes expiring OAuth tokens on its own resolve, + // so a cached access token could be served after expiry. + let token = cache.load_token(iii, state::STATE_SCOPE).await; // Effective settings from the router; tolerate a missing router (defaults). - let resolved = router_client::resolve(iii, token.as_deref()) + // An auth-classified failure drops the cache so the next attempt + // re-resolves fresh — retrying stays the router's job. + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) .await - .unwrap_or_else(|_| default_resolve()); + { + Ok(r) => r, + Err(e) => { + if classify_bus_error(&e) == ErrorKind::AuthExpired { + cache.invalidate(); + } + default_resolve() + } + }; let credential = auth::fetch_fresh_credential(iii).await; let cfg = match build_config( @@ -149,77 +159,10 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; -} - -/// Forward upstream events to the sink; ping through silence; stop on the -/// terminal event or a failed write (caller gone → dropping `rx` aborts the -/// upstream task and its in-flight HTTP request). -pub async fn pump( - mut rx: mpsc::Receiver, - sink: &dyn FrameSink, - ping_interval: Duration, -) { - loop { - match tokio::time::timeout(ping_interval, rx.recv()).await { - Ok(Some(ev)) => { - let terminal = ev.is_terminal(); - if send_event(sink, &ev).is_err() { - return; - } - if terminal { - return; - } - } - Ok(None) => return, - Err(_elapsed) => { - if send_event(sink, &AssistantMessageEvent::Ping).is_err() { - return; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::sse::empty_assistant; - use llm_router::chat::relay::RelayRead; - use llm_router::testkit::fake_channels::FakeChannel; - use serde_json::Value; - - fn done_event() -> AssistantMessageEvent { - AssistantMessageEvent::Done { - message: empty_assistant("codex/gpt-5.5"), - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn forwards_events_and_stops_at_terminal() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - tx.send(done_event()).await.unwrap(); - tx.send(AssistantMessageEvent::Ping).await.unwrap(); - drop(tx); - - pump(rx, &ch.writer, Duration::from_secs(30)).await; - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - assert_eq!(frames.len(), 2); - let last: Value = serde_json::from_str(&frames[1]).unwrap(); - assert_eq!(last["type"], "done"); + // An upstream auth terminal usually means the VAULT token expired (the + // existing vault flow handles refresh); invalidating the resolve/token + // cache is harmless — the next attempt just re-fetches both. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); } } diff --git a/provider-openai-codex/src/upstream.rs b/provider-openai-codex/src/upstream.rs index 882fbca25..925565b91 100644 --- a/provider-openai-codex/src/upstream.rs +++ b/provider-openai-codex/src/upstream.rs @@ -4,6 +4,9 @@ use crate::errors::classify; use crate::sse::{build_final, build_partial, handle_chunk, synthetic_error_event, PartialState}; use futures::StreamExt; +use llm_router::provider_scaffold::sse_transport::{ + append_utf8_chunk, drain_sse_blocks, error_chain, +}; use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; use serde_json::Value; use tokio::sync::mpsc; @@ -28,25 +31,6 @@ pub fn spawn_upstream( rx } -/// Flatten an error and its `source()` chain into one string. reqwest's -/// top-level Display for a builder error is just "builder error"; the real -/// cause (invalid header value, bad URL) lives in the source chain, so without -/// this the message is undiagnosable. -fn error_chain(e: &dyn std::error::Error) -> String { - let mut msg = e.to_string(); - let mut src = e.source(); - while let Some(s) = src { - let next = s.to_string(); - // reqwest sometimes nests the same text; skip exact repeats. - if !msg.ends_with(&next) { - msg.push_str(": "); - msg.push_str(&next); - } - src = s.source(); - } - msg -} - /// Last `data: ` payload in an SSE block, if any. fn data_line(block: &str) -> Option<&str> { block @@ -106,6 +90,33 @@ async fn run_upstream( let mut stream = resp.bytes_stream(); let mut buf = String::new(); + // Cross-chunk UTF-8 buffering: network chunks split multibyte + // codepoints, and a per-chunk lossy conversion corrupts them to U+FFFD. + let mut byte_buf = Vec::new(); + // Block decoder: [DONE] closes the stream (Stop + Done, Done terminal); + // anything else parses and runs the chunk state machine. + let decode = + |data_block: &str, state: &mut PartialState, model: &str| -> Vec { + let Some(data) = data_line(data_block) else { + return vec![]; + }; + if data == "[DONE]" { + return vec![ + AssistantMessageEvent::Stop { + stop_reason: state.stop_reason(), + error_message: None, + error_kind: None, + }, + AssistantMessageEvent::Done { + message: build_final(state, model), + }, + ]; + } + let Ok(parsed) = serde_json::from_str::(data) else { + return vec![]; + }; + handle_chunk(&parsed, state, model) + }; while let Some(chunk) = stream.next().await { let chunk = match chunk { Ok(c) => c, @@ -120,39 +131,13 @@ async fn run_upstream( return; } }; - buf.push_str(&String::from_utf8_lossy(&chunk)); - while let Some(idx) = buf.find("\n\n") { - let block: String = buf.drain(..idx + 2).collect(); - let Some(data) = data_line(&block) else { - continue; - }; - if data == "[DONE]" { - let _ = tx - .send(AssistantMessageEvent::Stop { - stop_reason: state.stop_reason(), - error_message: None, - error_kind: None, - }) - .await; - let _ = tx - .send(AssistantMessageEvent::Done { - message: build_final(&state, &args.model), - }) - .await; - return; - } - let Ok(parsed) = serde_json::from_str::(data) else { - continue; - }; - for ev in handle_chunk(&parsed, &mut state, &args.model) { - let terminal = ev.is_terminal(); - if tx.send(ev).await.is_err() { - return; // receiver dropped → abort upstream - } - if terminal { - return; // exactly one terminal event - } - } + append_utf8_chunk(&mut byte_buf, &mut buf, &chunk); + if drain_sse_blocks(&mut buf, &tx, &mut |block: &str| { + decode(block, &mut state, &args.model) + }) + .await + { + return; // terminal forwarded, or receiver dropped → abort upstream } } // Stream ended without [DONE]. With output accumulated this is diff --git a/provider-openai-codex/src/wire/names.rs b/provider-openai-codex/src/wire/names.rs index 7bc372f89..1801556bc 100644 --- a/provider-openai-codex/src/wire/names.rs +++ b/provider-openai-codex/src/wire/names.rs @@ -1,28 +1,5 @@ -//! iii function ids ↔ OpenAI function names. OpenAI enforces -//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. +//! iii function ids ↔ OpenAI tool names. OpenAI enforces +//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. Shared codec +//! (and its tests) live in `llm_router::provider_scaffold::names`. -pub fn encode_tool_name(name: &str) -> String { - name.replace("::", "__") -} - -pub fn decode_tool_name(name: &str) -> String { - name.replace("__", "::") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trips_bus_ids() { - assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); - assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); - assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); - } - - #[test] - fn plain_names_pass_through() { - assert_eq!(encode_tool_name("submit_result"), "submit_result"); - assert_eq!(decode_tool_name("submit_result"), "submit_result"); - } -} +pub use llm_router::provider_scaffold::names::{decode_tool_name, encode_tool_name}; diff --git a/provider-openai/Cargo.lock b/provider-openai/Cargo.lock index 375c3f914..371769073 100644 --- a/provider-openai/Cargo.lock +++ b/provider-openai/Cargo.lock @@ -776,7 +776,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.1.2" +version = "1.1.3" dependencies = [ "async-trait", "clap", diff --git a/provider-openai/src/embed.rs b/provider-openai/src/embed.rs index 02807e239..e94a3cee8 100644 --- a/provider-openai/src/embed.rs +++ b/provider-openai/src/embed.rs @@ -6,11 +6,14 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::types::events::ErrorKind; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use crate::config::{config_from_resolve, ConfigError}; -use crate::{router_client, state}; +use crate::errors::classify_bus_error; +use crate::state; pub const DEFAULT_EMBED_MODEL: &str = "text-embedding-3-small"; const EMBED_URL: &str = "https://api.openai.com/v1/embeddings"; @@ -71,6 +74,7 @@ struct WireResponse { pub async fn handle( iii: &IIIClient, http: &reqwest::Client, + cache: &ScaffoldCache, req: EmbedRequest, ) -> Result { if req.input.is_empty() || req.input.len() > 512 { @@ -83,8 +87,23 @@ pub async fn handle( .filter(|m| !m.trim().is_empty()) .unwrap_or_else(|| DEFAULT_EMBED_MODEL.to_string()); - let token = state::load_token(iii).await; - let resolved = router_client::resolve(iii, token.as_deref()).await?; + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + { + Ok(r) => r, + Err(e) => { + if classify_bus_error(&e) == ErrorKind::AuthExpired { + cache.invalidate(); + } + return Err(e); + } + }; // Reuse the chat-path credential validation; the api_url from resolve // targets the chat endpoint, so embeddings derive their sibling URL // from it (same host and version prefix). diff --git a/provider-openai/src/register.rs b/provider-openai/src/register.rs index ef5364b17..0c3501b01 100644 --- a/provider-openai/src/register.rs +++ b/provider-openai/src/register.rs @@ -9,6 +9,7 @@ use crate::{router_client, state, PROVIDER_ID}; use iii_sdk::errors::Error; use iii_sdk::protocol::RegisterTriggerInput; use iii_sdk::{IIIClient, RegisterFunction}; +use llm_router::provider_scaffold::cache::ScaffoldCache; use llm_router::types::router::{ ProviderDeclaration, ProviderDefaults, ProviderReadyAck, RouterReadyEvent, }; @@ -118,6 +119,11 @@ fn read_timeout() -> Duration { } pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { + // Shared per-process cache for the registration token and the resolve + // response (see llm_router::provider_scaffold::cache). Invalidated on + // router::ready — a restarted router may carry new config and reissues + // declare/refresh anyway — and on upstream auth errors (stream_fn). + let cache = ScaffoldCache::new(); // Streaming uses no total timeout (the router owns stream budgets), but // reads are silence-bounded: a stalled upstream otherwise pings the router // past its idle guard until the engine kills the call at stream_timeout. @@ -130,7 +136,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { iii.register_function( surface::STREAM_ID, RegisterFunction::new_async_with_bad_request( - make_stream(iii.clone(), http.clone()), + make_stream(iii.clone(), http.clone(), cache.clone()), invalid_request_from_serde, ) .description(surface::STREAM_DESC) @@ -145,11 +151,13 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { { let iii_embed = iii.clone(); let http_embed = http.clone(); + let cache_embed = cache.clone(); iii.register_function( surface::EMBED_ID, RegisterFunction::new_async(move |req: crate::embed::EmbedRequest| { - let (iii, http) = (iii_embed.clone(), http_embed.clone()); - async move { crate::embed::handle(&iii, &http, req).await } + let (iii, http, cache) = + (iii_embed.clone(), http_embed.clone(), cache_embed.clone()); + async move { crate::embed::handle(&iii, &http, &cache, req).await } }) .description(surface::EMBED_DESC) .metadata(json!({ "internal": true })), @@ -160,10 +168,12 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { { let iii_ready = iii.clone(); let http_ready = http.clone(); + let cache_ready = cache.clone(); iii.register_function( surface::ON_ROUTER_READY_ID, RegisterFunction::new_async(move |_event: RouterReadyEvent| { let (iii, http) = (iii_ready.clone(), http_ready.clone()); + cache_ready.invalidate(); async move { tokio::spawn(declare_and_refresh(iii, http)); Ok::<_, Error>(ProviderReadyAck { ok: true }) diff --git a/provider-openai/src/router_client.rs b/provider-openai/src/router_client.rs index 373c350e5..175f9e844 100644 --- a/provider-openai/src/router_client.rs +++ b/provider-openai/src/router_client.rs @@ -1,38 +1,20 @@ -//! Thin wrappers over the router's provider-protocol functions. All calls -//! carry the registration token (identity binding, spec adaptation #1). +//! Provider-scoped shims over the shared router-protocol client +//! (`llm_router::provider_scaffold::router_client`): every call binds this +//! crate's `PROVIDER_ID` and carries the registration token. use crate::PROVIDER_ID; use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::router_client as scaffold; use llm_router::types::model::Model; use llm_router::types::router::ProviderResolveResponse; -use serde_json::{json, Value}; - -async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { - iii.trigger(TriggerRequest { - function_id: function_id.into(), - payload, - action: None, - timeout_ms: Some(15_000), - }) - .await -} +use serde_json::Value; /// `router::provider::resolve` — credential + effective settings. pub async fn resolve( iii: &IIIClient, token: Option<&str>, ) -> Result { - let mut payload = json!({ "id": PROVIDER_ID }); - if let Some(t) = token { - payload["token"] = json!(t); - } - let raw = call(iii, "router::provider::resolve", payload).await?; - serde_json::from_value(raw).map_err(|e| Error::Remote { - code: "provider/bad_resolve_response".into(), - message: e.to_string(), - stacktrace: None, - }) + scaffold::resolve(iii, PROVIDER_ID, token).await } /// `router::models::reconcile` — replace this provider's catalog slice. @@ -41,30 +23,15 @@ pub async fn reconcile( models: Vec, token: Option<&str>, ) -> Result<(), Error> { - let mut payload = json!({ - "provider": PROVIDER_ID, - "models": serde_json::to_value(models).expect("serializable models"), - }); - if let Some(t) = token { - payload["token"] = json!(t); - } - call(iii, "router::models::reconcile", payload).await?; - Ok(()) + scaffold::reconcile(iii, PROVIDER_ID, models, token).await } /// `router::models::get` — authoritative catalog record (None when absent). pub async fn models_get(iii: &IIIClient, model_id: &str) -> Option { - let raw = call( - iii, - "router::models::get", - json!({ "provider": PROVIDER_ID, "id": model_id }), - ) - .await - .ok()?; - serde_json::from_value(raw.get("model")?.clone()).ok() + scaffold::models_get(iii, PROVIDER_ID, model_id).await } /// `router::provider::register` — returns the registration token to persist. pub async fn register(iii: &IIIClient, declaration: Value) -> Result { - call(iii, "router::provider::register", declaration).await + scaffold::register(iii, declaration).await } diff --git a/provider-openai/src/sse.rs b/provider-openai/src/sse.rs index 689f40806..a1ad8409e 100644 --- a/provider-openai/src/sse.rs +++ b/provider-openai/src/sse.rs @@ -262,7 +262,7 @@ pub fn handle_chunk( } state.text.push_str(text); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta: text.to_string(), }); } @@ -295,7 +295,7 @@ pub fn handle_chunk( if !args.is_empty() { state.function_calls[index].args_json.push_str(args); events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta: args.to_string(), id: state.function_calls[index].id.clone(), }); @@ -400,7 +400,7 @@ fn handle_responses_event( } state.text.push_str(delta); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta: delta.to_string(), }); } @@ -418,7 +418,7 @@ fn handle_responses_event( } state.thinking.push_str(delta); events.push(AssistantMessageEvent::ThinkingDelta { - partial: build_partial(state, model), + partial: None, delta: delta.to_string(), }); } @@ -443,7 +443,7 @@ fn handle_responses_event( call.args_json.push_str(delta); let id = call.id.clone(); events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta: delta.to_string(), id, }); diff --git a/provider-openai/src/state.rs b/provider-openai/src/state.rs index 22c229fd0..741feacbb 100644 --- a/provider-openai/src/state.rs +++ b/provider-openai/src/state.rs @@ -1,34 +1,15 @@ -//! Registration-token persistence in iii-state (engine `state::*` functions, -//! binary-worker.md § 7). The raw token lives here, under the provider's own -//! scope; the router persists only its sha256 hash. +//! Registration-token persistence, scoped to this provider's worker id +//! (shared logic in `llm_router::provider_scaffold::state`). use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; -use serde_json::{json, Value}; +use llm_router::provider_scaffold::state as scaffold; pub const STATE_SCOPE: &str = "provider-openai"; -const TOKEN_KEY: &str = "registration_token"; pub async fn load_token(iii: &IIIClient) -> Option { - let value = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY }), - action: None, - timeout_ms: None, - }) - .await - .ok()?; - value.as_str().map(String::from) + scaffold::load_token(iii, STATE_SCOPE).await } pub async fn store_token(iii: &IIIClient, token: &str) -> Result<(), Error> { - iii.trigger(TriggerRequest { - function_id: "state::set".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY, "value": Value::from(token) }), - action: None, - timeout_ms: None, - }) - .await?; - Ok(()) + scaffold::store_token(iii, STATE_SCOPE, token).await } diff --git a/provider-openai/src/stream_fn.rs b/provider-openai/src/stream_fn.rs index f2d2583be..0484a66ac 100644 --- a/provider-openai/src/stream_fn.rs +++ b/provider-openai/src/stream_fn.rs @@ -13,13 +13,10 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; use llm_router::channels::open_sink; use llm_router::chat::relay::FrameSink; -use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::provider_scaffold::pump::{pump, send_event, PING_INTERVAL}; +use llm_router::types::events::ErrorKind; use llm_router::types::router::{ProviderStreamInput, ProviderStreamOutput}; -use std::time::Duration; -use tokio::sync::mpsc; - -/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). -pub const PING_INTERVAL: Duration = Duration::from_secs(30); fn compatible_reasoning_effort( api_mode: ApiMode, @@ -41,15 +38,16 @@ fn compatible_reasoning_effort( pub fn make_stream( iii: IIIClient, http: reqwest::Client, + cache: ScaffoldCache, ) -> impl Fn(ProviderStreamInput) -> BoxFuture<'static, Result> + Send + Sync + 'static { move |input: ProviderStreamInput| { - let (iii, http) = (iii.clone(), http.clone()); + let (iii, http, cache) = (iii.clone(), http.clone(), cache.clone()); Box::pin(async move { let sink = open_sink(&iii, &input.writer_ref).await?; - run_stream_call(&iii, http, input, sink.as_ref()).await; + run_stream_call(&iii, http, &cache, input, sink.as_ref()).await; sink.close(); // ProviderStreamOutput (spec § stream contract) Ok(ProviderStreamOutput { ok: true }) @@ -57,30 +55,37 @@ pub fn make_stream( } } -fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { - let frame = serde_json::to_string(ev).expect("serializable event"); - sink.send(&frame).map_err(|_| ()) -} - async fn run_stream_call( iii: &IIIClient, http: reqwest::Client, + cache: &ScaffoldCache, input: ProviderStreamInput, sink: &dyn FrameSink, ) { let model = input.model.clone(); let mut warnings = Vec::new(); - let token = state::load_token(iii).await; - let resolved = match router_client::resolve(iii, token.as_deref()).await { + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + { Ok(r) => r, Err(e) => { + let kind = classify_bus_error(&e); + if kind == ErrorKind::AuthExpired { + cache.invalidate(); + } let _ = send_event( sink, &synthetic_error_event( &format!("router::provider::resolve failed: {e}"), &model, - classify_bus_error(&e), + kind, ), ); return; @@ -159,56 +164,16 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; -} - -/// Forward upstream events to the sink; ping through silence; stop on the -/// terminal event or on a failed write (caller gone → dropping `rx` aborts -/// the upstream task and its in-flight HTTP request). -/// Verbatim copy of provider-anthropic's pump — shared extraction into -/// llm-router is a listed follow-up. -pub async fn pump( - mut rx: mpsc::Receiver, - sink: &dyn FrameSink, - ping_interval: Duration, -) { - loop { - match tokio::time::timeout(ping_interval, rx.recv()).await { - Ok(Some(ev)) => { - let terminal = ev.is_terminal(); - if send_event(sink, &ev).is_err() { - return; - } - if terminal { - return; - } - } - // Upstream task ended without a terminal (panic/abort): the - // router synthesizes the terminal frame — never two terminals. - Ok(None) => return, - // Silent stretch: heartbeat (also probes for a gone caller). - Err(_elapsed) => { - if send_event(sink, &AssistantMessageEvent::Ping).is_err() { - return; - } - } - } + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); } } #[cfg(test)] mod tests { use super::*; - use crate::sse::empty_assistant; - use llm_router::chat::relay::RelayRead; - use llm_router::testkit::fake_channels::FakeChannel; - use serde_json::Value; - - fn done_event() -> AssistantMessageEvent { - AssistantMessageEvent::Done { - message: empty_assistant("gpt-test"), - } - } #[test] fn luna_tools_disable_effort_only_on_chat_completions() { @@ -235,84 +200,4 @@ mod tests { (Some("high"), false) ); } - - #[tokio::test(flavor = "multi_thread")] - async fn forwards_events_and_stops_at_terminal() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - tx.send(done_event()).await.unwrap(); - // a frame after the terminal must never be forwarded - tx.send(AssistantMessageEvent::Ping).await.unwrap(); - drop(tx); - - pump(rx, &ch.writer, Duration::from_secs(30)).await; - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - assert_eq!(frames.len(), 2); - let last: Value = serde_json::from_str(&frames[1]).unwrap(); - assert_eq!(last["type"], "done"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn pings_through_silence() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel::(8); - // hold tx open, send nothing for > 2 ping intervals, then terminate - let pump_task = { - let writer = ch.writer.clone(); - tokio::spawn(async move { pump(rx, &writer, Duration::from_millis(50)).await }) - }; - tokio::time::sleep(Duration::from_millis(140)).await; - tx.send(done_event()).await.unwrap(); - drop(tx); - pump_task.await.unwrap(); - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - let pings = frames - .iter() - .filter(|f| serde_json::from_str::(f).unwrap()["type"] == "ping") - .count(); - assert!( - pings >= 2, - "want >=2 pings through 140ms of silence, got {pings}" - ); - assert_eq!( - serde_json::from_str::(frames.last().unwrap()).unwrap()["type"], - "done" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn reader_close_stops_the_pump_and_drops_the_receiver() { - let ch = FakeChannel::new(); - ch.reader.close(); // caller gone before anything is written - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - pump(rx, &ch.writer, Duration::from_secs(30)).await; // returns immediately - // the receiver was consumed and dropped by pump → upstream send fails - assert!(tx.send(done_event()).await.is_err()); - } } diff --git a/provider-openai/src/upstream.rs b/provider-openai/src/upstream.rs index fb5b06887..6752fc942 100644 --- a/provider-openai/src/upstream.rs +++ b/provider-openai/src/upstream.rs @@ -5,6 +5,9 @@ use crate::errors::classify; use crate::sse::{build_final, build_partial, handle_chunk, synthetic_error_event, PartialState}; use futures::StreamExt; +use llm_router::provider_scaffold::sse_transport::{ + append_utf8_chunk, drain_sse_blocks, error_chain, +}; use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; use serde_json::Value; use tokio::sync::mpsc; @@ -29,25 +32,6 @@ pub fn spawn_upstream( rx } -/// Flatten an error and its `source()` chain into one string. reqwest's -/// top-level Display for a builder error is just "builder error"; the real -/// cause (invalid header value, bad URL) lives in the source chain, so without -/// this the message is undiagnosable. -fn error_chain(e: &dyn std::error::Error) -> String { - let mut msg = e.to_string(); - let mut src = e.source(); - while let Some(s) = src { - let next = s.to_string(); - // reqwest sometimes nests the same text; skip exact repeats. - if !msg.ends_with(&next) { - msg.push_str(": "); - msg.push_str(&next); - } - src = s.source(); - } - msg -} - /// Last `data: ` payload in an SSE block, if any. fn data_line(block: &str) -> Option<&str> { block @@ -107,6 +91,33 @@ async fn run_upstream( let mut stream = resp.bytes_stream(); let mut buf = String::new(); + // Cross-chunk UTF-8 buffering: network chunks split multibyte + // codepoints, and a per-chunk lossy conversion corrupts them to U+FFFD. + let mut byte_buf = Vec::new(); + // Block decoder: [DONE] closes the stream (Stop + Done, Done terminal); + // anything else parses and runs the chunk state machine. + let decode = + |data_block: &str, state: &mut PartialState, model: &str| -> Vec { + let Some(data) = data_line(data_block) else { + return vec![]; + }; + if data == "[DONE]" { + return vec![ + AssistantMessageEvent::Stop { + stop_reason: state.stop_reason(), + error_message: None, + error_kind: None, + }, + AssistantMessageEvent::Done { + message: build_final(state, model), + }, + ]; + } + let Ok(parsed) = serde_json::from_str::(data) else { + return vec![]; + }; + handle_chunk(&parsed, state, model) + }; while let Some(chunk) = stream.next().await { let chunk = match chunk { Ok(c) => c, @@ -121,39 +132,13 @@ async fn run_upstream( return; } }; - buf.push_str(&String::from_utf8_lossy(&chunk)); - while let Some(idx) = buf.find("\n\n") { - let block: String = buf.drain(..idx + 2).collect(); - let Some(data) = data_line(&block) else { - continue; - }; - if data == "[DONE]" { - let _ = tx - .send(AssistantMessageEvent::Stop { - stop_reason: state.stop_reason(), - error_message: None, - error_kind: None, - }) - .await; - let _ = tx - .send(AssistantMessageEvent::Done { - message: build_final(&state, &args.model), - }) - .await; - return; - } - let Ok(parsed) = serde_json::from_str::(data) else { - continue; - }; - for ev in handle_chunk(&parsed, &mut state, &args.model) { - let terminal = ev.is_terminal(); - if tx.send(ev).await.is_err() { - return; // receiver dropped → abort upstream - } - if terminal { - return; // exactly one terminal event - } - } + append_utf8_chunk(&mut byte_buf, &mut buf, &chunk); + if drain_sse_blocks(&mut buf, &tx, &mut |block: &str| { + decode(block, &mut state, &args.model) + }) + .await + { + return; // terminal forwarded, or receiver dropped → abort upstream } } // Responses usually terminates with `response.completed`; compatible diff --git a/provider-openai/src/wire/names.rs b/provider-openai/src/wire/names.rs index 7bc372f89..5ab3150c1 100644 --- a/provider-openai/src/wire/names.rs +++ b/provider-openai/src/wire/names.rs @@ -1,28 +1,5 @@ -//! iii function ids ↔ OpenAI function names. OpenAI enforces -//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. +//! iii function ids ↔ OpenAI tool names. OpenAI enforces +//! `^[a-zA-Z0-9_-]{1,128}$`; bus ids use `::` separators. Shared codec +//! (and its tests) live in `llm_router::provider_scaffold::names`. -pub fn encode_tool_name(name: &str) -> String { - name.replace("::", "__") -} - -pub fn decode_tool_name(name: &str) -> String { - name.replace("__", "::") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trips_bus_ids() { - assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); - assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); - assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); - } - - #[test] - fn plain_names_pass_through() { - assert_eq!(encode_tool_name("submit_result"), "submit_result"); - assert_eq!(decode_tool_name("submit_result"), "submit_result"); - } -} +pub use llm_router::provider_scaffold::names::{decode_tool_name, encode_tool_name}; diff --git a/provider-xai/Cargo.lock b/provider-xai/Cargo.lock index f65f988d8..28722bc5f 100644 --- a/provider-xai/Cargo.lock +++ b/provider-xai/Cargo.lock @@ -776,7 +776,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.1.1" +version = "1.1.3" dependencies = [ "async-trait", "clap", diff --git a/provider-xai/src/register.rs b/provider-xai/src/register.rs index 932f971b5..b66482c1d 100644 --- a/provider-xai/src/register.rs +++ b/provider-xai/src/register.rs @@ -9,6 +9,7 @@ use crate::{router_client, state, PROVIDER_ID}; use iii_sdk::errors::Error; use iii_sdk::protocol::RegisterTriggerInput; use iii_sdk::{IIIClient, RegisterFunction}; +use llm_router::provider_scaffold::cache::ScaffoldCache; use llm_router::types::router::{ ProviderDeclaration, ProviderDefaults, ProviderReadyAck, RouterReadyEvent, }; @@ -118,6 +119,11 @@ fn read_timeout() -> Duration { } pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { + // Shared per-process cache for the registration token and the resolve + // response (see llm_router::provider_scaffold::cache). Invalidated on + // router::ready — a restarted router may carry new config and reissues + // declare/refresh anyway — and on upstream auth errors (stream_fn). + let cache = ScaffoldCache::new(); // Streaming uses no total timeout (the router owns stream budgets), but // reads are silence-bounded: a stalled upstream otherwise pings the router // past its idle guard until the engine kills the call at stream_timeout. @@ -163,7 +169,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { iii.register_function( surface::STREAM_ID, RegisterFunction::new_async_with_bad_request( - make_stream(iii.clone(), http.clone(), cell.clone()), + make_stream(iii.clone(), http.clone(), cell.clone(), cache.clone()), invalid_request_from_serde, ) .description(surface::STREAM_DESC) @@ -180,10 +186,12 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { { let iii_ready = iii.clone(); let http_ready = http.clone(); + let cache_ready = cache.clone(); iii.register_function( surface::ON_ROUTER_READY_ID, RegisterFunction::new_async(move |_event: RouterReadyEvent| { let (iii, http) = (iii_ready.clone(), http_ready.clone()); + cache_ready.invalidate(); async move { tokio::spawn(declare_and_refresh(iii, http)); Ok::<_, Error>(ProviderReadyAck { ok: true }) diff --git a/provider-xai/src/responses.rs b/provider-xai/src/responses.rs index 85b1b55ef..d7cbe85a4 100644 --- a/provider-xai/src/responses.rs +++ b/provider-xai/src/responses.rs @@ -172,7 +172,7 @@ pub fn handle_event( } state.thinking.push_str(delta); out.push(AssistantMessageEvent::ThinkingDelta { - partial: partial(state, model), + partial: None, delta: delta.to_string(), }); } @@ -195,7 +195,7 @@ pub fn handle_event( } state.text.push_str(delta); out.push(AssistantMessageEvent::TextDelta { - partial: partial(state, model), + partial: None, delta: delta.to_string(), }); } @@ -477,6 +477,41 @@ mod tests { .collect() } + /// Contract pin (llm-router types::events): delta frames are slim — + /// no cumulative partial per chunk — while block-boundary frames carry + /// the authoritative snapshot. Readers reconstruct via + /// llm_router::chat::accumulate. + #[test] + fn deltas_are_slim_and_boundary_snapshots_are_cumulative() { + let (_, evs) = run(&[ + ("response.output_text.delta", json!({"delta":"He"})), + ("response.output_text.delta", json!({"delta":"llo"})), + ("response.output_text.done", json!({})), + ]); + for ev in &evs { + match ev { + AssistantMessageEvent::TextDelta { partial, .. } + | AssistantMessageEvent::ThinkingDelta { partial, .. } => { + assert!(partial.is_none(), "delta frames must not carry partial"); + } + _ => {} + } + } + let Some(AssistantMessageEvent::TextEnd { partial }) = evs + .iter() + .find(|e| matches!(e, AssistantMessageEvent::TextEnd { .. })) + else { + panic!("want a text_end frame"); + }; + assert!( + matches!( + &partial.content[0], + ContentBlock::Text { text } if text == "Hello" + ), + "the End snapshot must carry the cumulative block text" + ); + } + #[test] fn reasoning_then_text_with_citation_and_usage() { let (state, evs) = run(&[ diff --git a/provider-xai/src/router_client.rs b/provider-xai/src/router_client.rs index 373c350e5..175f9e844 100644 --- a/provider-xai/src/router_client.rs +++ b/provider-xai/src/router_client.rs @@ -1,38 +1,20 @@ -//! Thin wrappers over the router's provider-protocol functions. All calls -//! carry the registration token (identity binding, spec adaptation #1). +//! Provider-scoped shims over the shared router-protocol client +//! (`llm_router::provider_scaffold::router_client`): every call binds this +//! crate's `PROVIDER_ID` and carries the registration token. use crate::PROVIDER_ID; use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::router_client as scaffold; use llm_router::types::model::Model; use llm_router::types::router::ProviderResolveResponse; -use serde_json::{json, Value}; - -async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { - iii.trigger(TriggerRequest { - function_id: function_id.into(), - payload, - action: None, - timeout_ms: Some(15_000), - }) - .await -} +use serde_json::Value; /// `router::provider::resolve` — credential + effective settings. pub async fn resolve( iii: &IIIClient, token: Option<&str>, ) -> Result { - let mut payload = json!({ "id": PROVIDER_ID }); - if let Some(t) = token { - payload["token"] = json!(t); - } - let raw = call(iii, "router::provider::resolve", payload).await?; - serde_json::from_value(raw).map_err(|e| Error::Remote { - code: "provider/bad_resolve_response".into(), - message: e.to_string(), - stacktrace: None, - }) + scaffold::resolve(iii, PROVIDER_ID, token).await } /// `router::models::reconcile` — replace this provider's catalog slice. @@ -41,30 +23,15 @@ pub async fn reconcile( models: Vec, token: Option<&str>, ) -> Result<(), Error> { - let mut payload = json!({ - "provider": PROVIDER_ID, - "models": serde_json::to_value(models).expect("serializable models"), - }); - if let Some(t) = token { - payload["token"] = json!(t); - } - call(iii, "router::models::reconcile", payload).await?; - Ok(()) + scaffold::reconcile(iii, PROVIDER_ID, models, token).await } /// `router::models::get` — authoritative catalog record (None when absent). pub async fn models_get(iii: &IIIClient, model_id: &str) -> Option { - let raw = call( - iii, - "router::models::get", - json!({ "provider": PROVIDER_ID, "id": model_id }), - ) - .await - .ok()?; - serde_json::from_value(raw.get("model")?.clone()).ok() + scaffold::models_get(iii, PROVIDER_ID, model_id).await } /// `router::provider::register` — returns the registration token to persist. pub async fn register(iii: &IIIClient, declaration: Value) -> Result { - call(iii, "router::provider::register", declaration).await + scaffold::register(iii, declaration).await } diff --git a/provider-xai/src/sse.rs b/provider-xai/src/sse.rs index da535addb..371f7ad99 100644 --- a/provider-xai/src/sse.rs +++ b/provider-xai/src/sse.rs @@ -259,7 +259,7 @@ pub fn handle_chunk( } state.thinking.push_str(reasoning); events.push(AssistantMessageEvent::ThinkingDelta { - partial: build_partial(state, model), + partial: None, delta: reasoning.to_string(), }); } @@ -275,7 +275,7 @@ pub fn handle_chunk( } state.text.push_str(text); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta: text.to_string(), }); } @@ -308,7 +308,7 @@ pub fn handle_chunk( if !args.is_empty() { state.function_calls[index].args_json.push_str(args); events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta: args.to_string(), id: state.function_calls[index].id.clone(), }); @@ -365,6 +365,37 @@ mod tests { .collect() } + /// Contract pin (llm-router types::events): delta frames are slim — + /// no cumulative partial per chunk — while block-boundary frames carry + /// the authoritative snapshot. Readers reconstruct via + /// llm_router::chat::accumulate. + #[test] + fn deltas_are_slim_and_boundary_snapshots_are_cumulative() { + let (_, events) = run(&[ + json!({"choices":[{"index":0,"delta":{"content":"He"}}]}), + json!({"choices":[{"index":0,"delta":{"content":"llo"}}]}), + json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}), + ]); + for ev in &events { + if let AssistantMessageEvent::TextDelta { partial, .. } = ev { + assert!(partial.is_none(), "delta frames must not carry partial"); + } + } + let Some(AssistantMessageEvent::TextEnd { partial }) = events + .iter() + .find(|e| matches!(e, AssistantMessageEvent::TextEnd { .. })) + else { + panic!("want a text_end frame"); + }; + assert!( + matches!( + &partial.content[0], + ContentBlock::Text { text } if text == "Hello" + ), + "the End snapshot must carry the cumulative block text" + ); + } + #[test] fn text_stream_produces_start_delta_end_and_final_content() { let (state, events) = run(&[ diff --git a/provider-xai/src/state.rs b/provider-xai/src/state.rs index 24b0d3e96..4ec98c2ad 100644 --- a/provider-xai/src/state.rs +++ b/provider-xai/src/state.rs @@ -1,34 +1,15 @@ -//! Registration-token persistence in iii-state (engine `state::*` functions, -//! binary-worker.md § 7). The raw token lives here, under the provider's own -//! scope; the router persists only its sha256 hash. +//! Registration-token persistence, scoped to this provider's worker id +//! (shared logic in `llm_router::provider_scaffold::state`). use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; -use serde_json::{json, Value}; +use llm_router::provider_scaffold::state as scaffold; pub const STATE_SCOPE: &str = "provider-xai"; -const TOKEN_KEY: &str = "registration_token"; pub async fn load_token(iii: &IIIClient) -> Option { - let value = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY }), - action: None, - timeout_ms: None, - }) - .await - .ok()?; - value.as_str().map(String::from) + scaffold::load_token(iii, STATE_SCOPE).await } pub async fn store_token(iii: &IIIClient, token: &str) -> Result<(), Error> { - iii.trigger(TriggerRequest { - function_id: "state::set".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY, "value": Value::from(token) }), - action: None, - timeout_ms: None, - }) - .await?; - Ok(()) + scaffold::store_token(iii, STATE_SCOPE, token).await } diff --git a/provider-xai/src/stream_fn.rs b/provider-xai/src/stream_fn.rs index bf22c8ef4..2809bc30d 100644 --- a/provider-xai/src/stream_fn.rs +++ b/provider-xai/src/stream_fn.rs @@ -17,28 +17,26 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; use llm_router::channels::open_sink; use llm_router::chat::relay::FrameSink; -use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::provider_scaffold::pump::{pump, send_event, PING_INTERVAL}; +use llm_router::types::events::ErrorKind; use llm_router::types::router::{ProviderStreamInput, ProviderStreamOutput}; -use std::time::Duration; -use tokio::sync::mpsc; - -/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). -pub const PING_INTERVAL: Duration = Duration::from_secs(30); pub fn make_stream( iii: IIIClient, http: reqwest::Client, cell: ConfigCell, + cache: ScaffoldCache, ) -> impl Fn(ProviderStreamInput) -> BoxFuture<'static, Result> + Send + Sync + 'static { move |input: ProviderStreamInput| { - let (iii, http, cell) = (iii.clone(), http.clone(), cell.clone()); + let (iii, http, cell, cache) = (iii.clone(), http.clone(), cell.clone(), cache.clone()); Box::pin(async move { let sink = open_sink(&iii, &input.writer_ref).await?; let wc = { cell.read().await.clone() }; - run_stream_call(&iii, http, &wc, input, sink.as_ref()).await; + run_stream_call(&iii, http, &cache, &wc, input, sink.as_ref()).await; sink.close(); // ProviderStreamOutput (spec § stream contract) Ok(ProviderStreamOutput { ok: true }) @@ -46,14 +44,10 @@ pub fn make_stream( } } -fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { - let frame = serde_json::to_string(ev).expect("serializable event"); - sink.send(&frame).map_err(|_| ()) -} - async fn run_stream_call( iii: &IIIClient, http: reqwest::Client, + cache: &ScaffoldCache, wc: &crate::config::WorkerConfig, input: ProviderStreamInput, sink: &dyn FrameSink, @@ -61,16 +55,27 @@ async fn run_stream_call( let model = input.model.clone(); let mut warnings = Vec::new(); - let token = state::load_token(iii).await; - let resolved = match router_client::resolve(iii, token.as_deref()).await { + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + { Ok(r) => r, Err(e) => { + let kind = classify_bus_error(&e); + if kind == ErrorKind::AuthExpired { + cache.invalidate(); + } let _ = send_event( sink, &synthetic_error_event( &format!("router::provider::resolve failed: {e}"), &model, - classify_bus_error(&e), + kind, ), ); return; @@ -132,7 +137,11 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); + } return; } @@ -185,134 +194,9 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; -} - -/// Forward upstream events to the sink; ping through silence; stop on the -/// terminal event or on a failed write (caller gone → dropping `rx` aborts -/// the upstream task and its in-flight HTTP request). -/// Verbatim copy of provider-anthropic's pump — shared extraction into -/// llm-router is a listed follow-up. -pub async fn pump( - mut rx: mpsc::Receiver, - sink: &dyn FrameSink, - ping_interval: Duration, -) { - loop { - match tokio::time::timeout(ping_interval, rx.recv()).await { - Ok(Some(ev)) => { - let terminal = ev.is_terminal(); - if send_event(sink, &ev).is_err() { - return; - } - if terminal { - return; - } - } - // Upstream task ended without a terminal (panic/abort): the - // router synthesizes the terminal frame — never two terminals. - Ok(None) => return, - // Silent stretch: heartbeat (also probes for a gone caller). - Err(_elapsed) => { - if send_event(sink, &AssistantMessageEvent::Ping).is_err() { - return; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::sse::empty_assistant; - use llm_router::chat::relay::RelayRead; - use llm_router::testkit::fake_channels::FakeChannel; - use serde_json::Value; - - fn done_event() -> AssistantMessageEvent { - AssistantMessageEvent::Done { - message: empty_assistant("grok-test"), - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn forwards_events_and_stops_at_terminal() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - tx.send(done_event()).await.unwrap(); - // a frame after the terminal must never be forwarded - tx.send(AssistantMessageEvent::Ping).await.unwrap(); - drop(tx); - - pump(rx, &ch.writer, Duration::from_secs(30)).await; - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - assert_eq!(frames.len(), 2); - let last: Value = serde_json::from_str(&frames[1]).unwrap(); - assert_eq!(last["type"], "done"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn pings_through_silence() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel::(8); - // hold tx open, send nothing for > 2 ping intervals, then terminate - let pump_task = { - let writer = ch.writer.clone(); - tokio::spawn(async move { pump(rx, &writer, Duration::from_millis(50)).await }) - }; - tokio::time::sleep(Duration::from_millis(140)).await; - tx.send(done_event()).await.unwrap(); - drop(tx); - pump_task.await.unwrap(); - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - let pings = frames - .iter() - .filter(|f| serde_json::from_str::(f).unwrap()["type"] == "ping") - .count(); - assert!( - pings >= 2, - "want >=2 pings through 140ms of silence, got {pings}" - ); - assert_eq!( - serde_json::from_str::(frames.last().unwrap()).unwrap()["type"], - "done" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn reader_close_stops_the_pump_and_drops_the_receiver() { - let ch = FakeChannel::new(); - ch.reader.close(); // caller gone before anything is written - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - pump(rx, &ch.writer, Duration::from_secs(30)).await; // returns immediately - // the receiver was consumed and dropped by pump → upstream send fails - assert!(tx.send(done_event()).await.is_err()); + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); } } diff --git a/provider-xai/src/upstream.rs b/provider-xai/src/upstream.rs index 9d801f5ea..52ec081ce 100644 --- a/provider-xai/src/upstream.rs +++ b/provider-xai/src/upstream.rs @@ -4,6 +4,9 @@ use crate::errors::classify; use crate::sse::{build_final, build_partial, handle_chunk, synthetic_error_event, PartialState}; use futures::StreamExt; +use llm_router::provider_scaffold::sse_transport::{ + append_utf8_chunk, drain_sse_blocks, error_chain, +}; use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; use serde_json::Value; use tokio::sync::mpsc; @@ -28,25 +31,6 @@ pub fn spawn_upstream( rx } -/// Flatten an error and its `source()` chain into one string. reqwest's -/// top-level Display for a builder error is just "builder error"; the real -/// cause (invalid header value, bad URL) lives in the source chain, so without -/// this the message is undiagnosable. -fn error_chain(e: &dyn std::error::Error) -> String { - let mut msg = e.to_string(); - let mut src = e.source(); - while let Some(s) = src { - let next = s.to_string(); - // reqwest sometimes nests the same text; skip exact repeats. - if !msg.ends_with(&next) { - msg.push_str(": "); - msg.push_str(&next); - } - src = s.source(); - } - msg -} - /// Last `data: ` payload in an SSE block, if any. fn data_line(block: &str) -> Option<&str> { block @@ -106,6 +90,33 @@ async fn run_upstream( let mut stream = resp.bytes_stream(); let mut buf = String::new(); + // Cross-chunk UTF-8 buffering: network chunks split multibyte + // codepoints, and a per-chunk lossy conversion corrupts them to U+FFFD. + let mut byte_buf = Vec::new(); + // Block decoder: [DONE] closes the stream (Stop + Done, Done terminal); + // anything else parses and runs the chunk state machine. + let decode = + |data_block: &str, state: &mut PartialState, model: &str| -> Vec { + let Some(data) = data_line(data_block) else { + return vec![]; + }; + if data == "[DONE]" { + return vec![ + AssistantMessageEvent::Stop { + stop_reason: state.stop_reason(), + error_message: None, + error_kind: None, + }, + AssistantMessageEvent::Done { + message: build_final(state, model), + }, + ]; + } + let Ok(parsed) = serde_json::from_str::(data) else { + return vec![]; + }; + handle_chunk(&parsed, state, model) + }; while let Some(chunk) = stream.next().await { let chunk = match chunk { Ok(c) => c, @@ -120,39 +131,13 @@ async fn run_upstream( return; } }; - buf.push_str(&String::from_utf8_lossy(&chunk)); - while let Some(idx) = buf.find("\n\n") { - let block: String = buf.drain(..idx + 2).collect(); - let Some(data) = data_line(&block) else { - continue; - }; - if data == "[DONE]" { - let _ = tx - .send(AssistantMessageEvent::Stop { - stop_reason: state.stop_reason(), - error_message: None, - error_kind: None, - }) - .await; - let _ = tx - .send(AssistantMessageEvent::Done { - message: build_final(&state, &args.model), - }) - .await; - return; - } - let Ok(parsed) = serde_json::from_str::(data) else { - continue; - }; - for ev in handle_chunk(&parsed, &mut state, &args.model) { - let terminal = ev.is_terminal(); - if tx.send(ev).await.is_err() { - return; // receiver dropped → abort upstream - } - if terminal { - return; // exactly one terminal event - } - } + append_utf8_chunk(&mut byte_buf, &mut buf, &chunk); + if drain_sse_blocks(&mut buf, &tx, &mut |block: &str| { + decode(block, &mut state, &args.model) + }) + .await + { + return; // terminal forwarded, or receiver dropped → abort upstream } } // Stream ended without [DONE] (connection close framing): still terminal. diff --git a/provider-xai/src/wire/names.rs b/provider-xai/src/wire/names.rs index e7fcb4079..456788d96 100644 --- a/provider-xai/src/wire/names.rs +++ b/provider-xai/src/wire/names.rs @@ -1,28 +1,5 @@ -//! iii function ids ↔ xAI function names. xAI enforces -//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. +//! iii function ids ↔ xAI tool names. xAI enforces +//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. Shared codec +//! (and its tests) live in `llm_router::provider_scaffold::names`. -pub fn encode_tool_name(name: &str) -> String { - name.replace("::", "__") -} - -pub fn decode_tool_name(name: &str) -> String { - name.replace("__", "::") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trips_bus_ids() { - assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); - assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); - assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); - } - - #[test] - fn plain_names_pass_through() { - assert_eq!(encode_tool_name("submit_result"), "submit_result"); - assert_eq!(decode_tool_name("submit_result"), "submit_result"); - } -} +pub use llm_router::provider_scaffold::names::{decode_tool_name, encode_tool_name}; diff --git a/provider-zai/Cargo.lock b/provider-zai/Cargo.lock index dcaceb0b0..30e98c8ff 100644 --- a/provider-zai/Cargo.lock +++ b/provider-zai/Cargo.lock @@ -776,7 +776,7 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "llm-router" -version = "1.1.1" +version = "1.1.3" dependencies = [ "async-trait", "clap", diff --git a/provider-zai/src/register.rs b/provider-zai/src/register.rs index 95b8ff196..dd074a7ba 100644 --- a/provider-zai/src/register.rs +++ b/provider-zai/src/register.rs @@ -9,6 +9,7 @@ use crate::{router_client, state, PROVIDER_ID}; use iii_sdk::errors::Error; use iii_sdk::protocol::RegisterTriggerInput; use iii_sdk::{IIIClient, RegisterFunction}; +use llm_router::provider_scaffold::cache::ScaffoldCache; use llm_router::types::router::{ ProviderDeclaration, ProviderDefaults, ProviderReadyAck, RouterReadyEvent, }; @@ -121,6 +122,11 @@ fn read_timeout() -> Duration { } pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { + // Shared per-process cache for the registration token and the resolve + // response (see llm_router::provider_scaffold::cache). Invalidated on + // router::ready — a restarted router may carry new config and reissues + // declare/refresh anyway — and on upstream auth errors (stream_fn). + let cache = ScaffoldCache::new(); // Streaming uses no total timeout (the router owns stream budgets), but // reads are silence-bounded: a stalled upstream otherwise pings the router // past its idle guard until the engine kills the call at stream_timeout. @@ -133,7 +139,7 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { iii.register_function( surface::STREAM_ID, RegisterFunction::new_async_with_bad_request( - make_stream(iii.clone(), http.clone()), + make_stream(iii.clone(), http.clone(), cache.clone()), invalid_request_from_serde, ) .description(surface::STREAM_DESC) @@ -149,10 +155,12 @@ pub async fn register_provider(iii: IIIClient) -> Result<(), Error> { // Re-declare when the router restarts: bind to the router::ready trigger type. { let iii_ready = iii.clone(); + let cache_ready = cache.clone(); iii.register_function( surface::ON_ROUTER_READY_ID, RegisterFunction::new_async(move |_event: RouterReadyEvent| { let iii = iii_ready.clone(); + cache_ready.invalidate(); async move { tokio::spawn(declare_and_refresh(iii)); Ok::<_, Error>(ProviderReadyAck { ok: true }) diff --git a/provider-zai/src/router_client.rs b/provider-zai/src/router_client.rs index 373c350e5..175f9e844 100644 --- a/provider-zai/src/router_client.rs +++ b/provider-zai/src/router_client.rs @@ -1,38 +1,20 @@ -//! Thin wrappers over the router's provider-protocol functions. All calls -//! carry the registration token (identity binding, spec adaptation #1). +//! Provider-scoped shims over the shared router-protocol client +//! (`llm_router::provider_scaffold::router_client`): every call binds this +//! crate's `PROVIDER_ID` and carries the registration token. use crate::PROVIDER_ID; use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; +use llm_router::provider_scaffold::router_client as scaffold; use llm_router::types::model::Model; use llm_router::types::router::ProviderResolveResponse; -use serde_json::{json, Value}; - -async fn call(iii: &IIIClient, function_id: &str, payload: Value) -> Result { - iii.trigger(TriggerRequest { - function_id: function_id.into(), - payload, - action: None, - timeout_ms: Some(15_000), - }) - .await -} +use serde_json::Value; /// `router::provider::resolve` — credential + effective settings. pub async fn resolve( iii: &IIIClient, token: Option<&str>, ) -> Result { - let mut payload = json!({ "id": PROVIDER_ID }); - if let Some(t) = token { - payload["token"] = json!(t); - } - let raw = call(iii, "router::provider::resolve", payload).await?; - serde_json::from_value(raw).map_err(|e| Error::Remote { - code: "provider/bad_resolve_response".into(), - message: e.to_string(), - stacktrace: None, - }) + scaffold::resolve(iii, PROVIDER_ID, token).await } /// `router::models::reconcile` — replace this provider's catalog slice. @@ -41,30 +23,15 @@ pub async fn reconcile( models: Vec, token: Option<&str>, ) -> Result<(), Error> { - let mut payload = json!({ - "provider": PROVIDER_ID, - "models": serde_json::to_value(models).expect("serializable models"), - }); - if let Some(t) = token { - payload["token"] = json!(t); - } - call(iii, "router::models::reconcile", payload).await?; - Ok(()) + scaffold::reconcile(iii, PROVIDER_ID, models, token).await } /// `router::models::get` — authoritative catalog record (None when absent). pub async fn models_get(iii: &IIIClient, model_id: &str) -> Option { - let raw = call( - iii, - "router::models::get", - json!({ "provider": PROVIDER_ID, "id": model_id }), - ) - .await - .ok()?; - serde_json::from_value(raw.get("model")?.clone()).ok() + scaffold::models_get(iii, PROVIDER_ID, model_id).await } /// `router::provider::register` — returns the registration token to persist. pub async fn register(iii: &IIIClient, declaration: Value) -> Result { - call(iii, "router::provider::register", declaration).await + scaffold::register(iii, declaration).await } diff --git a/provider-zai/src/sse.rs b/provider-zai/src/sse.rs index e41a199d5..cab444e8d 100644 --- a/provider-zai/src/sse.rs +++ b/provider-zai/src/sse.rs @@ -260,7 +260,7 @@ pub fn handle_chunk( } state.thinking.push_str(reasoning); events.push(AssistantMessageEvent::ThinkingDelta { - partial: build_partial(state, model), + partial: None, delta: reasoning.to_string(), }); } @@ -276,7 +276,7 @@ pub fn handle_chunk( } state.text.push_str(text); events.push(AssistantMessageEvent::TextDelta { - partial: build_partial(state, model), + partial: None, delta: text.to_string(), }); } @@ -315,7 +315,7 @@ pub fn handle_chunk( if !args.is_empty() { state.function_calls[index].args_json.push_str(args); events.push(AssistantMessageEvent::FunctioncallDelta { - partial: build_partial(state, model), + partial: None, delta: args.to_string(), id: state.function_calls[index].id.clone(), }); @@ -372,6 +372,37 @@ mod tests { .collect() } + /// Contract pin (llm-router types::events): delta frames are slim — + /// no cumulative partial per chunk — while block-boundary frames carry + /// the authoritative snapshot. Readers reconstruct via + /// llm_router::chat::accumulate. + #[test] + fn deltas_are_slim_and_boundary_snapshots_are_cumulative() { + let (_, events) = run(&[ + json!({"choices":[{"index":0,"delta":{"content":"He"}}]}), + json!({"choices":[{"index":0,"delta":{"content":"llo"}}]}), + json!({"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}), + ]); + for ev in &events { + if let AssistantMessageEvent::TextDelta { partial, .. } = ev { + assert!(partial.is_none(), "delta frames must not carry partial"); + } + } + let Some(AssistantMessageEvent::TextEnd { partial }) = events + .iter() + .find(|e| matches!(e, AssistantMessageEvent::TextEnd { .. })) + else { + panic!("want a text_end frame"); + }; + assert!( + matches!( + &partial.content[0], + ContentBlock::Text { text } if text == "Hello" + ), + "the End snapshot must carry the cumulative block text" + ); + } + #[test] fn text_stream_produces_start_delta_end_and_final_content() { let (state, events) = run(&[ diff --git a/provider-zai/src/state.rs b/provider-zai/src/state.rs index 7f49c245a..b83ff257d 100644 --- a/provider-zai/src/state.rs +++ b/provider-zai/src/state.rs @@ -1,34 +1,15 @@ -//! Registration-token persistence in iii-state (engine `state::*` functions, -//! binary-worker.md § 7). The raw token lives here, under the provider's own -//! scope; the router persists only its sha256 hash. +//! Registration-token persistence, scoped to this provider's worker id +//! (shared logic in `llm_router::provider_scaffold::state`). use iii_sdk::errors::Error; -use iii_sdk::protocol::TriggerRequest; use iii_sdk::IIIClient; -use serde_json::{json, Value}; +use llm_router::provider_scaffold::state as scaffold; pub const STATE_SCOPE: &str = "provider-zai"; -const TOKEN_KEY: &str = "registration_token"; pub async fn load_token(iii: &IIIClient) -> Option { - let value = iii - .trigger(TriggerRequest { - function_id: "state::get".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY }), - action: None, - timeout_ms: None, - }) - .await - .ok()?; - value.as_str().map(String::from) + scaffold::load_token(iii, STATE_SCOPE).await } pub async fn store_token(iii: &IIIClient, token: &str) -> Result<(), Error> { - iii.trigger(TriggerRequest { - function_id: "state::set".into(), - payload: json!({ "scope": STATE_SCOPE, "key": TOKEN_KEY, "value": Value::from(token) }), - action: None, - timeout_ms: None, - }) - .await?; - Ok(()) + scaffold::store_token(iii, STATE_SCOPE, token).await } diff --git a/provider-zai/src/stream_fn.rs b/provider-zai/src/stream_fn.rs index 1d42265c3..e6ae888f2 100644 --- a/provider-zai/src/stream_fn.rs +++ b/provider-zai/src/stream_fn.rs @@ -13,26 +13,24 @@ use iii_sdk::errors::Error; use iii_sdk::IIIClient; use llm_router::channels::open_sink; use llm_router::chat::relay::FrameSink; -use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; +use llm_router::provider_scaffold::cache::ScaffoldCache; +use llm_router::provider_scaffold::pump::{pump, send_event, PING_INTERVAL}; +use llm_router::types::events::ErrorKind; use llm_router::types::router::{ProviderStreamInput, ProviderStreamOutput}; -use std::time::Duration; -use tokio::sync::mpsc; - -/// Heartbeat cadence while the upstream is silent (spec: at least every 30s). -pub const PING_INTERVAL: Duration = Duration::from_secs(30); pub fn make_stream( iii: IIIClient, http: reqwest::Client, + cache: ScaffoldCache, ) -> impl Fn(ProviderStreamInput) -> BoxFuture<'static, Result> + Send + Sync + 'static { move |input: ProviderStreamInput| { - let (iii, http) = (iii.clone(), http.clone()); + let (iii, http, cache) = (iii.clone(), http.clone(), cache.clone()); Box::pin(async move { let sink = open_sink(&iii, &input.writer_ref).await?; - run_stream_call(&iii, http, input, sink.as_ref()).await; + run_stream_call(&iii, http, &cache, input, sink.as_ref()).await; sink.close(); // ProviderStreamOutput (spec § stream contract) Ok(ProviderStreamOutput { ok: true }) @@ -40,30 +38,37 @@ pub fn make_stream( } } -fn send_event(sink: &dyn FrameSink, ev: &AssistantMessageEvent) -> Result<(), ()> { - let frame = serde_json::to_string(ev).expect("serializable event"); - sink.send(&frame).map_err(|_| ()) -} - async fn run_stream_call( iii: &IIIClient, http: reqwest::Client, + cache: &ScaffoldCache, input: ProviderStreamInput, sink: &dyn FrameSink, ) { let model = input.model.clone(); let mut warnings = Vec::new(); - let token = state::load_token(iii).await; - let resolved = match router_client::resolve(iii, token.as_deref()).await { + // Token + resolve are cached (ScaffoldCache): zero engine round trips + // on the hot path within the TTL. An auth-classified resolve failure + // drops the cache so the next attempt re-resolves fresh — retrying + // stays the router's job. + let token = cache.load_token(iii, state::STATE_SCOPE).await; + let resolved = match cache + .resolve(iii, crate::PROVIDER_ID, token.as_deref()) + .await + { Ok(r) => r, Err(e) => { + let kind = classify_bus_error(&e); + if kind == ErrorKind::AuthExpired { + cache.invalidate(); + } let _ = send_event( sink, &synthetic_error_event( &format!("router::provider::resolve failed: {e}"), &model, - classify_bus_error(&e), + kind, ), ); return; @@ -138,134 +143,9 @@ async fn run_stream_call( warnings, }, ); - pump(rx, sink, PING_INTERVAL).await; -} - -/// Forward upstream events to the sink; ping through silence; stop on the -/// terminal event or on a failed write (caller gone → dropping `rx` aborts -/// the upstream task and its in-flight HTTP request). -/// Verbatim copy of provider-anthropic's pump — shared extraction into -/// llm-router is a listed follow-up. -pub async fn pump( - mut rx: mpsc::Receiver, - sink: &dyn FrameSink, - ping_interval: Duration, -) { - loop { - match tokio::time::timeout(ping_interval, rx.recv()).await { - Ok(Some(ev)) => { - let terminal = ev.is_terminal(); - if send_event(sink, &ev).is_err() { - return; - } - if terminal { - return; - } - } - // Upstream task ended without a terminal (panic/abort): the - // router synthesizes the terminal frame — never two terminals. - Ok(None) => return, - // Silent stretch: heartbeat (also probes for a gone caller). - Err(_elapsed) => { - if send_event(sink, &AssistantMessageEvent::Ping).is_err() { - return; - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::sse::empty_assistant; - use llm_router::chat::relay::RelayRead; - use llm_router::testkit::fake_channels::FakeChannel; - use serde_json::Value; - - fn done_event() -> AssistantMessageEvent { - AssistantMessageEvent::Done { - message: empty_assistant("glm-test"), - } - } - - #[tokio::test(flavor = "multi_thread")] - async fn forwards_events_and_stops_at_terminal() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - tx.send(done_event()).await.unwrap(); - // a frame after the terminal must never be forwarded - tx.send(AssistantMessageEvent::Ping).await.unwrap(); - drop(tx); - - pump(rx, &ch.writer, Duration::from_secs(30)).await; - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - assert_eq!(frames.len(), 2); - let last: Value = serde_json::from_str(&frames[1]).unwrap(); - assert_eq!(last["type"], "done"); - } - - #[tokio::test(flavor = "multi_thread")] - async fn pings_through_silence() { - let ch = FakeChannel::new(); - let (tx, rx) = mpsc::channel::(8); - // hold tx open, send nothing for > 2 ping intervals, then terminate - let pump_task = { - let writer = ch.writer.clone(); - tokio::spawn(async move { pump(rx, &writer, Duration::from_millis(50)).await }) - }; - tokio::time::sleep(Duration::from_millis(140)).await; - tx.send(done_event()).await.unwrap(); - drop(tx); - pump_task.await.unwrap(); - ch.writer.close(); - - let mut frames = Vec::new(); - let mut reader = ch.reader; - while let llm_router::chat::relay::ReadEvent::Msg(m) = - reader.next(Duration::from_millis(100)).await - { - frames.push(m); - } - let pings = frames - .iter() - .filter(|f| serde_json::from_str::(f).unwrap()["type"] == "ping") - .count(); - assert!( - pings >= 2, - "want >=2 pings through 140ms of silence, got {pings}" - ); - assert_eq!( - serde_json::from_str::(frames.last().unwrap()).unwrap()["type"], - "done" - ); - } - - #[tokio::test(flavor = "multi_thread")] - async fn reader_close_stops_the_pump_and_drops_the_receiver() { - let ch = FakeChannel::new(); - ch.reader.close(); // caller gone before anything is written - let (tx, rx) = mpsc::channel(8); - tx.send(AssistantMessageEvent::Start { - partial: empty_assistant("m"), - }) - .await - .unwrap(); - pump(rx, &ch.writer, Duration::from_secs(30)).await; // returns immediately - // the receiver was consumed and dropped by pump → upstream send fails - assert!(tx.send(done_event()).await.is_err()); + // An upstream auth terminal means the cached credential was rotated + // out from under us: drop the cache so the next attempt re-resolves. + if pump(rx, sink, PING_INTERVAL).await == Some(ErrorKind::AuthExpired) { + cache.invalidate(); } } diff --git a/provider-zai/src/upstream.rs b/provider-zai/src/upstream.rs index 6bdc7f369..46dccbb83 100644 --- a/provider-zai/src/upstream.rs +++ b/provider-zai/src/upstream.rs @@ -5,6 +5,9 @@ use crate::errors::classify; use crate::sse::{build_final, build_partial, handle_chunk, synthetic_error_event, PartialState}; use futures::StreamExt; +use llm_router::provider_scaffold::sse_transport::{ + append_utf8_chunk, drain_sse_blocks, error_chain, +}; use llm_router::types::events::{AssistantMessageEvent, ErrorKind}; use serde_json::Value; use tokio::sync::mpsc; @@ -29,66 +32,6 @@ pub fn spawn_upstream( rx } -/// Flatten an error and its `source()` chain into one string. reqwest's -/// top-level Display for a builder error is just "builder error"; the real -/// cause (invalid header value, bad URL) lives in the source chain, so without -/// this the message is undiagnosable. -fn error_chain(e: &dyn std::error::Error) -> String { - let mut msg = e.to_string(); - let mut src = e.source(); - while let Some(s) = src { - let next = s.to_string(); - // reqwest sometimes nests the same text; skip exact repeats. - if !msg.ends_with(&next) { - msg.push_str(": "); - msg.push_str(&next); - } - src = s.source(); - } - msg -} - -/// Append chunk bytes to `text`, retaining any trailing incomplete UTF-8 -/// sequence: transport chunk boundaries are arbitrary, so a multi-byte -/// character (GLM output is frequently CJK) can split across chunks — a -/// per-chunk lossy decode would corrupt it into U+FFFD. -fn append_utf8_chunk(byte_buf: &mut Vec, text: &mut String, chunk: &[u8]) { - byte_buf.extend_from_slice(chunk); - let mut consumed = 0usize; - loop { - match std::str::from_utf8(&byte_buf[consumed..]) { - Ok(s) => { - text.push_str(s); - byte_buf.clear(); - return; - } - Err(e) => { - let valid = e.valid_up_to(); - if valid > 0 { - // SAFETY: valid_up_to guarantees valid UTF-8 in this prefix. - text.push_str(unsafe { - std::str::from_utf8_unchecked(&byte_buf[consumed..consumed + valid]) - }); - consumed += valid; - } - match e.error_len() { - Some(invalid) => { - byte_buf.drain(..consumed + invalid); - text.push('\u{FFFD}'); - consumed = 0; - } - None => { - if consumed > 0 { - byte_buf.drain(..consumed); - } - return; - } - } - } - } - } -} - /// Last `data: ` payload in an SSE block, if any. fn data_line(block: &str) -> Option<&str> { block @@ -148,7 +91,34 @@ async fn run_upstream( let mut stream = resp.bytes_stream(); let mut buf = String::new(); - let mut byte_buf: Vec = Vec::new(); + // Cross-chunk UTF-8 buffering: transport chunk boundaries are arbitrary, + // so a multi-byte character (GLM output is frequently CJK) can split + // across chunks — a per-chunk lossy decode would corrupt it into U+FFFD. + let mut byte_buf = Vec::new(); + // Block decoder: [DONE] closes the stream (Stop + Done, Done terminal); + // anything else parses and runs the chunk state machine. + let decode = + |data_block: &str, state: &mut PartialState, model: &str| -> Vec { + let Some(data) = data_line(data_block) else { + return vec![]; + }; + if data == "[DONE]" { + return vec![ + AssistantMessageEvent::Stop { + stop_reason: state.stop_reason(), + error_message: None, + error_kind: None, + }, + AssistantMessageEvent::Done { + message: build_final(state, model), + }, + ]; + } + let Ok(parsed) = serde_json::from_str::(data) else { + return vec![]; + }; + handle_chunk(&parsed, state, model) + }; while let Some(chunk) = stream.next().await { let chunk = match chunk { Ok(c) => c, @@ -164,38 +134,12 @@ async fn run_upstream( } }; append_utf8_chunk(&mut byte_buf, &mut buf, &chunk); - while let Some(idx) = buf.find("\n\n") { - let block: String = buf.drain(..idx + 2).collect(); - let Some(data) = data_line(&block) else { - continue; - }; - if data == "[DONE]" { - let _ = tx - .send(AssistantMessageEvent::Stop { - stop_reason: state.stop_reason(), - error_message: None, - error_kind: None, - }) - .await; - let _ = tx - .send(AssistantMessageEvent::Done { - message: build_final(&state, &args.model), - }) - .await; - return; - } - let Ok(parsed) = serde_json::from_str::(data) else { - continue; - }; - for ev in handle_chunk(&parsed, &mut state, &args.model) { - let terminal = ev.is_terminal(); - if tx.send(ev).await.is_err() { - return; // receiver dropped → abort upstream - } - if terminal { - return; // exactly one terminal event - } - } + if drain_sse_blocks(&mut buf, &tx, &mut |block: &str| { + decode(block, &mut state, &args.model) + }) + .await + { + return; // terminal forwarded, or receiver dropped → abort upstream } } // Stream ended without [DONE] (connection close framing): still terminal. @@ -349,24 +293,6 @@ mod tests { } } - #[test] - fn utf8_split_across_chunks_survives_reassembly() { - // "你" = 3 bytes; split it across two transport chunks. - let bytes = "data: 你好\n\n".as_bytes(); - let (a, b) = bytes.split_at(8); // mid-character - let (mut byte_buf, mut text) = (Vec::new(), String::new()); - append_utf8_chunk(&mut byte_buf, &mut text, a); - append_utf8_chunk(&mut byte_buf, &mut text, b); - assert_eq!(text, "data: 你好\n\n"); - assert!(!text.contains('\u{FFFD}'), "lossy corruption: {text:?}"); - assert!(byte_buf.is_empty()); - - // Truly invalid bytes still degrade to U+FFFD instead of stalling. - let (mut byte_buf, mut text) = (Vec::new(), String::new()); - append_utf8_chunk(&mut byte_buf, &mut text, &[b'a', 0xFF, b'b']); - assert_eq!(text, "a\u{FFFD}b"); - } - #[tokio::test(flavor = "multi_thread")] async fn warnings_arrive_on_the_final_message() { let url = stub(HAPPY).await; diff --git a/provider-zai/src/wire/names.rs b/provider-zai/src/wire/names.rs index f944c183a..037e504c2 100644 --- a/provider-zai/src/wire/names.rs +++ b/provider-zai/src/wire/names.rs @@ -1,28 +1,5 @@ -//! iii function ids ↔ Z.AI function names. Z.AI enforces -//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. +//! iii function ids ↔ Z.AI tool names. Z.AI enforces +//! `^[a-zA-Z0-9_-]{1,64}$`; bus ids use `::` separators. Shared codec +//! (and its tests) live in `llm_router::provider_scaffold::names`. -pub fn encode_tool_name(name: &str) -> String { - name.replace("::", "__") -} - -pub fn decode_tool_name(name: &str) -> String { - name.replace("__", "::") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn round_trips_bus_ids() { - assert_eq!(encode_tool_name("web::fetch"), "web__fetch"); - assert_eq!(decode_tool_name("web__fetch"), "web::fetch"); - assert_eq!(decode_tool_name(&encode_tool_name("a::b::c")), "a::b::c"); - } - - #[test] - fn plain_names_pass_through() { - assert_eq!(encode_tool_name("submit_result"), "submit_result"); - assert_eq!(decode_tool_name("submit_result"), "submit_result"); - } -} +pub use llm_router::provider_scaffold::names::{decode_tool_name, encode_tool_name}; diff --git a/session-manager/src/functions/messages.rs b/session-manager/src/functions/messages.rs index 11a7205ea..6372ff367 100644 --- a/session-manager/src/functions/messages.rs +++ b/session-manager/src/functions/messages.rs @@ -14,6 +14,12 @@ pub struct MessagesRequest { pub limit: Option, /// Opaque pagination cursor from a previous response. pub cursor: Option, + /// Return only entries strictly after this entry id on the + /// requested path (incremental fetch from a known watermark). + /// Errors `session/invalid_cursor` when the entry is not on the + /// path — e.g. after a fork moved the active leaf — so callers + /// fall back to a full load. Ignored when `cursor` is set. + pub after_entry_id: Option, /// Only messages with these roles. Setting this also excludes /// `kind: "custom"` entries (it is an explicit narrowing to roles). pub roles: Option>, diff --git a/session-manager/src/service.rs b/session-manager/src/service.rs index 3c38f7c0f..d1c94b83c 100644 --- a/session-manager/src/service.rs +++ b/session-manager/src/service.rs @@ -786,7 +786,23 @@ impl SessionService { .collect(); let start = match &req.cursor { - None => 0, + // `after_entry_id` is a caller-held watermark: resume + // strictly after it, with the same not-on-path error the + // cursor path uses so callers share one fallback. + None => match &req.after_entry_id { + None => 0, + Some(after) => { + let position = filtered + .iter() + .position(|e| e.id() == after.as_str()) + .ok_or_else(|| { + SessionError::InvalidCursor(format!( + "after_entry_id {after} is not on the requested path" + )) + })?; + position + 1 + } + }, Some(raw) => { let cursor: MessagesCursor = decode_cursor(raw)?; let position = filtered diff --git a/session-manager/tests/features/messages.feature b/session-manager/tests/features/messages.feature index c8454387f..337337944 100644 --- a/session-manager/tests/features/messages.feature +++ b/session-manager/tests/features/messages.feature @@ -85,6 +85,65 @@ Feature: session::messages — load the active path, oldest first And the response field "messages.0.entry_id" is "e_005" And the response has no field "next_cursor" + # Prevents: the incremental-fetch watermark drifting from strict-after + # semantics — a steering check re-seeing the watermark entry itself + # would loop, and skipping one entry would miss a user interjection. + Scenario: after_entry_id returns only entries strictly after it + Given a user message "m1" appended to "s_001" + And a user message "m2" appended to "s_001" + And a user message "m3" appended to "s_001" + When I call "session::messages" with: + """ + { "session_id": "s_001", "after_entry_id": "e_001" } + """ + Then the response field "messages" has length 2 + And the response field "messages.0.entry_id" is "e_002" + And the response field "messages.1.entry_id" is "e_003" + When I call "session::messages" with: + """ + { "session_id": "s_001", "after_entry_id": "e_003" } + """ + Then the call succeeds + And the response field "messages" has length 0 + + # Prevents: the watermark silently paging from the path start when the + # entry left the active path (fork/set_active_leaf) — callers need the + # invalid-cursor signal to fall back to a full reload. + Scenario: after_entry_id off the requested path is rejected + Given a user message "m1" appended to "s_001" + When I call "session::messages" with: + """ + { "session_id": "s_001", "after_entry_id": "ghost" } + """ + Then the call fails with code "session/invalid_cursor" + + # Prevents: the watermark landing on a filtered-out custom entry and + # erroring — the steering check reads with include_custom and its + # watermark can be a bookkeeping entry. + Scenario: after_entry_id works with include_custom and paginates + Given a user message "m1" appended to "s_001" + And a custom entry of type "compaction" appended to "s_001" + And a user message "m2" appended to "s_001" + And a user message "m3" appended to "s_001" + When I call "session::messages" with: + """ + { "session_id": "s_001", "after_entry_id": "e_002", "include_custom": true, "limit": 1 } + """ + Then the response field "messages" has length 1 + And the response field "messages.0.entry_id" is "e_003" + And I alias the response field "next_cursor" as "C1" + When I call "session::messages" with: + """ + { "session_id": "s_001", "cursor": "${C1}", "include_custom": true } + """ + Then the response field "messages" has length 1 + And the response field "messages.0.entry_id" is "e_004" + When I call "session::messages" with: + """ + { "session_id": "s_001", "after_entry_id": "e_002" } + """ + Then the call fails with code "session/invalid_cursor" + # Prevents: accepting forged/stale cursors and returning garbage pages. Scenario: a malformed cursor is rejected Given a user message "m1" appended to "s_001" diff --git a/session-manager/tests/golden/schemas/session.messages.json b/session-manager/tests/golden/schemas/session.messages.json index bf18fdb8c..e9d6e960d 100644 --- a/session-manager/tests/golden/schemas/session.messages.json +++ b/session-manager/tests/golden/schemas/session.messages.json @@ -15,6 +15,13 @@ } }, "properties": { + "after_entry_id": { + "description": "Return only entries strictly after this entry id on the requested path (incremental fetch from a known watermark). Errors `session/invalid_cursor` when the entry is not on the path — e.g. after a fork moved the active leaf — so callers fall back to a full load. Ignored when `cursor` is set.", + "type": [ + "string", + "null" + ] + }, "cursor": { "description": "Opaque pagination cursor from a previous response.", "type": [ diff --git a/tech-specs/2026-06-agentic/README.md b/tech-specs/2026-06-agentic/README.md index cc5b1ab6a..200153974 100644 --- a/tech-specs/2026-06-agentic/README.md +++ b/tech-specs/2026-06-agentic/README.md @@ -340,8 +340,12 @@ when asked (`include_custom`). Neither reaches the model: providers only ever re ### Streaming events The discriminated union providers stream over an iii channel, relayed verbatim by `llm-router` and -the `harness`. Non-terminal content frames carry a `partial` accumulator (`usage`, `ping`, and -`stop` do not); `done`/`error` carry the final assembled message. `done` and `error` are terminal. +the `harness`. Block-boundary frames (`start` and the `*_start`/`*_end` frames) carry a required +cumulative `partial` snapshot — thinking signatures and finalized function-call arguments exist +only on those boundaries. The `*_delta` frames may omit `partial` (a delta that carries one is +honored as an authoritative snapshot); readers reconstruct the cumulative message as the last +boundary snapshot plus accumulated deltas. `usage`, `ping`, and `stop` carry no `partial`; +`done`/`error` carry the final assembled message. `done` and `error` are terminal. ```typescript type StopReason = "end" | "length" | "function_call" | "aborted" | "error"; @@ -356,13 +360,13 @@ type Usage = { type AssistantMessageEvent = | { type: "start"; partial: AssistantMessage } | { type: "text_start"; partial: AssistantMessage } - | { type: "text_delta"; partial: AssistantMessage; delta: string } + | { type: "text_delta"; partial?: AssistantMessage; delta: string } | { type: "text_end"; partial: AssistantMessage } | { type: "thinking_start"; partial: AssistantMessage } - | { type: "thinking_delta"; partial: AssistantMessage; delta: string } + | { type: "thinking_delta"; partial?: AssistantMessage; delta: string } | { type: "thinking_end"; partial: AssistantMessage } | { type: "functioncall_start";partial: AssistantMessage } - | { type: "functioncall_delta";partial: AssistantMessage; delta: string } + | { type: "functioncall_delta";partial?: AssistantMessage; delta: string } | { type: "functioncall_end"; partial: AssistantMessage } | { type: "usage"; usage: Usage } | { type: "ping" } // liveness heartbeat; consumers ignore diff --git a/tech-specs/2026-06-agentic/llm-router.md b/tech-specs/2026-06-agentic/llm-router.md index 24b03907f..6ef19e231 100644 --- a/tech-specs/2026-06-agentic/llm-router.md +++ b/tech-specs/2026-06-agentic/llm-router.md @@ -342,8 +342,8 @@ Example: } // channel frames (abbreviated) { "type": "start", "partial": { "role": "assistant", "content": [], "...": "..." } } -{ "type": "text_delta", "delta": "He", "partial": { "...": "..." } } -{ "type": "text_delta", "delta": "llo", "partial": { "...": "..." } } +{ "type": "text_delta", "delta": "He" } // slim deltas omit `partial`; readers accumulate onto the last boundary snapshot +{ "type": "text_delta", "delta": "llo" } { "type": "usage", "usage": { "input": 12, "output": 2 } } { "type": "done", "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }], "stop_reason": "end", "model": "claude-sonnet-4", "provider": "anthropic", "timestamp": 2 } } // response (after stream closes)