diff --git a/crates/openhuman-core/src/config/schema/context.rs b/crates/openhuman-core/src/config/schema/context.rs index f2a893567c3..b17eadf7a28 100644 --- a/crates/openhuman-core/src/config/schema/context.rs +++ b/crates/openhuman-core/src/config/schema/context.rs @@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize}; pub use tinyagents_harness::config::{CompactionSettings, CompactionStrategy, ContextConfig}; -pub use tinyagents_harness::summarization::ToolRendering; +use tinyagents_harness::summarization::ToolRendering; /// The whole `[context]` table: TinyAgents' [`ContextConfig`] and the /// compaction summarizer's OpenHuman-side settings. diff --git a/crates/openhuman-core/src/inference/provider/factory.rs b/crates/openhuman-core/src/inference/provider/factory.rs index e4076039a4d..21f96d7df46 100644 --- a/crates/openhuman-core/src/inference/provider/factory.rs +++ b/crates/openhuman-core/src/inference/provider/factory.rs @@ -13,10 +13,20 @@ pub(crate) mod access_gates; mod chat_model; +mod cloud_fleet; +mod cloud_http; mod cloud_slug; mod cloud_spec; +// The adapter is compiled and unit-tested before wiring any live BYOK route. +// No credential group is enabled in production until explicit factory support. +#[allow(dead_code)] +mod credential_fleet; mod credentials; mod inspection; +// Wave-1 scheduler foundation: remain isolated until the model-factory integration +// adopts it in the next stacked change. No inference behavior changes in this draft. +#[allow(dead_code)] +mod key_pool; mod local_runtime; mod managed_backend; mod routing; @@ -29,6 +39,13 @@ pub use chat_model::{ pub(crate) use chat_model::{resolves_to_managed_backend, OpenHumanModelFactory}; pub(crate) use cloud_slug::resolve_cloud_slug_for; pub(crate) use cloud_spec::{cloud_provider_spec, CloudCredentialResolver}; +/// Catalog discovery is a one-credential HTTP call. Never give its request +/// builder a delimited fleet: that would forward multiple secrets as a single +/// Authorization header. The native inference path owns per-call rotation. +pub(crate) fn byok_catalog_keys(raw: &str) -> Vec { + key_pool::parse_api_keys(raw) +} + pub use credentials::{auth_key_for_slug, lookup_key_for_slug}; pub(crate) use credentials::{ openai_bearer_is_oauth, redact_endpoint_origin, resolve_codex_routing, diff --git a/crates/openhuman-core/src/inference/provider/factory/cloud_fleet.rs b/crates/openhuman-core/src/inference/provider/factory/cloud_fleet.rs new file mode 100644 index 00000000000..2013c3d9b1f --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/cloud_fleet.rs @@ -0,0 +1,58 @@ +//! Minimal native-cloud adapter for multi-key BYOK requests. +//! The ordinary single-key model stays unchanged. Inspection constructs only +//! the first credential's model; it must not perform runtime failover. +use std::sync::Arc; + +use tinyinference_llm::model::ChatModel; + +use super::credential_fleet::CredentialFleetModel; +use super::key_pool::RotationStrategy; + +/// Health and cooldown belong to a destination, not merely a provider alias: +/// changing the endpoint must not carry its previous 429 or invalid-key state. +pub(super) fn rotation_pool_scope( + workspace: &std::path::Path, + slug: &str, + model: &str, + endpoint: &str, +) -> String { + format!( + "{}:{}:{}:{:016x}", + workspace.display(), + slug, + model, + super::key_pool::scope_fingerprint(endpoint), + ) +} + +pub(super) fn build_cloud_model_with_keys( + scope: String, + keys: Vec, + enable_rotation: bool, + mut build: impl FnMut(&str) -> anyhow::Result>>, +) -> anyhow::Result>> { + let keys = if keys.is_empty() { + vec![String::new()] + } else { + keys + }; + if !enable_rotation || keys.len() == 1 { + return build(&keys[0]); + } + + let models = keys + .iter() + .map(|key| build(key)) + .collect::>>()?; + Ok(Arc::new(CredentialFleetModel::new( + scope, + keys, + models, + None, + RotationStrategy::RoundRobin, + ))) +} + +#[cfg(test)] +#[path = "cloud_fleet_tests.rs"] +mod tests; diff --git a/crates/openhuman-core/src/inference/provider/factory/cloud_fleet_tests.rs b/crates/openhuman-core/src/inference/provider/factory/cloud_fleet_tests.rs new file mode 100644 index 00000000000..b43810d54e8 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/cloud_fleet_tests.rs @@ -0,0 +1,47 @@ +use super::*; +use std::cell::RefCell; + +struct NoNetworkModel; + +#[async_trait::async_trait] +impl ChatModel<()> for NoNetworkModel { + async fn invoke( + &self, + _: &(), + _: tinyinference_llm::model::ModelRequest, + ) -> tinyinference_llm::Result { + Err(tinyinference_llm::Error::Model("unused test route".into())) + } +} + +#[test] +fn inspection_binds_one_credential_without_enabling_failover() { + let built = RefCell::new(Vec::new()); + let result = build_cloud_model_with_keys( + "inspection-model".into(), + vec!["key-one".into(), "key-two".into()], + false, + |key| { + built.borrow_mut().push(key.to_owned()); + Ok(Arc::new(NoNetworkModel) as Arc>) + }, + ); + assert!(result.is_ok()); + assert_eq!(*built.borrow(), ["key-one"]); +} + +#[test] +fn runtime_builds_one_native_model_per_distinct_credential() { + let built = RefCell::new(Vec::new()); + let result = build_cloud_model_with_keys( + "runtime-model".into(), + vec!["key-one".into(), "key-two".into()], + true, + |key| { + built.borrow_mut().push(key.to_owned()); + Ok(Arc::new(NoNetworkModel) as Arc>) + }, + ); + assert!(result.is_ok()); + assert_eq!(*built.borrow(), ["key-one", "key-two"]); +} diff --git a/crates/openhuman-core/src/inference/provider/factory/cloud_http.rs b/crates/openhuman-core/src/inference/provider/factory/cloud_http.rs new file mode 100644 index 00000000000..6bda0b7b493 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/cloud_http.rs @@ -0,0 +1,31 @@ +//! HTTP-client policy shared by native BYOK construction. +//! Credentialed routes never follow redirects: a custom authorization header +//! must not travel to a different host or a plaintext HTTP endpoint. +use crate::config::schema::cloud_providers::AuthStyle; + +pub(super) fn build_cloud_http_client( + style: AuthStyle, + anthropic_native: bool, + ca_pem: Option<&str>, +) -> Result, String> { + let service_key = if anthropic_native { + "provider.anthropic" + } else { + "provider.compatible" + }; + let pem = ca_pem.filter(|value| !value.is_empty()); + if matches!(style, AuthStyle::Bearer | AuthStyle::Anthropic) { + let client = crate::util::tls::credentialed_client_builder(service_key, pem)? + .build() + .map_err(|_| "Could not build secure provider inference HTTP client".to_string())?; + return Ok(Some(client)); + } + match pem { + Some(pem) => crate::util::tls::client_with_ca_bundle(pem, service_key).map(Some), + None => Ok(None), + } +} + +#[cfg(test)] +#[path = "cloud_http_tests.rs"] +mod tests; diff --git a/crates/openhuman-core/src/inference/provider/factory/cloud_http_tests.rs b/crates/openhuman-core/src/inference/provider/factory/cloud_http_tests.rs new file mode 100644 index 00000000000..75e80f99636 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/cloud_http_tests.rs @@ -0,0 +1,24 @@ +use super::*; + +#[test] +fn unauthenticated_cloud_routes_keep_the_optional_http_client_behavior() { + assert!(build_cloud_http_client(AuthStyle::None, false, None) + .expect("no-auth client") + .is_none()); + assert!(build_cloud_http_client( + AuthStyle::None, + false, + Some(include_str!("../../../util/tls/test-ca.pem")), + ) + .expect("custom CA client") + .is_some()); +} + +#[test] +fn authenticated_cloud_routes_always_install_a_redirect_safe_client() { + for style in [AuthStyle::Bearer, AuthStyle::Anthropic] { + assert!(build_cloud_http_client(style, false, None) + .expect("authenticated client") + .is_some()); + } +} diff --git a/crates/openhuman-core/src/inference/provider/factory/cloud_slug.rs b/crates/openhuman-core/src/inference/provider/factory/cloud_slug.rs index e4714264bce..51a23124d19 100644 --- a/crates/openhuman-core/src/inference/provider/factory/cloud_slug.rs +++ b/crates/openhuman-core/src/inference/provider/factory/cloud_slug.rs @@ -77,6 +77,8 @@ pub(crate) fn resolve_cloud_slug_for<'a>( ) } +pub(super) use super::cloud_fleet::rotation_pool_scope; + /// The user-facing error for a failed credential lookup: the guidance, plus /// the expiry signal when the lookup reported one. The doctor's /// refresh-required check and the Codex reconnect classifier key on those @@ -661,45 +663,66 @@ fn try_create_cloud_slug_chat_model_from_string_with_native_tools_and_access( crate::security::egress::EgressDescriptor::inference(&slug, &effective_model, true), ); - let http = match config - .cloud_provider_ca_certs - .get(&slug) - .filter(|pem| !pem.is_empty()) - { - Some(pem) => { - let service_key = if anthropic_native { - "provider.anthropic" - } else { - "provider.compatible" - }; - match crate::util::tls::client_with_ca_bundle(pem, service_key) { - Ok(http) => Some(http), - Err(error) => return Some(Err(anyhow::anyhow!(error))), - } - } - None => None, + let http = match super::cloud_http::build_cloud_http_client( + entry.auth_style, + anthropic_native, + config + .cloud_provider_ca_certs + .get(&slug) + .map(String::as_str), + ) { + Ok(client) => client, + Err(error) => return Some(Err(anyhow::anyhow!(error))), }; let auth_style = cloud_spec::registry_auth_style(entry.auth_style) .expect("managed JWT entries returned above"); - let chat = cloud_spec::build_cloud_model(cloud_spec::CloudBuild { - spec: cloud_spec::cloud_provider_spec(entry, &effective_model), - auth_style, - key: key.as_str(), - temperature_override, - temperature_unsupported_models: config.temperature_unsupported_models.as_slice(), - native_tool_calling, - default_provider_options, - extra_headers: extra_headers.as_slice(), - extra_query_params: extra_query_params.as_slice(), - user_agent: user_agent.as_deref(), - codex: codex_route, - http, - }); + // Parse only actual BYOK credentials. OAuth, hosted and unauthenticated + // routes retain their original single-key identity and provider semantics. + let eligible = + !codex.using_oauth && matches!(entry.auth_style, AuthStyle::Bearer | AuthStyle::Anthropic); + let keys = if eligible { + super::key_pool::parse_api_keys(&key) + } else { + vec![key.clone()] + }; + if keys.len() > 1 && (crate::core::runtime::is_saas() || codex.account_id.is_some()) { + return Some(Err(anyhow::anyhow!( + "multi-key BYOK rotation is unavailable for shared or account-bound sessions" + ))); + } + let inspection_key = keys.first().cloned().unwrap_or_default(); + let scope = rotation_pool_scope( + &config.workspace_dir, + &slug, + &effective_model, + &codex.endpoint, + ); + let chat = super::cloud_fleet::build_cloud_model_with_keys( + scope, + keys, + eligible && credential_access == CredentialAccess::Runtime, + |credential| { + cloud_spec::build_cloud_model(cloud_spec::CloudBuild { + spec: cloud_spec::cloud_provider_spec(entry, &effective_model), + auth_style: auth_style.clone(), + key: credential, + temperature_override, + temperature_unsupported_models: config.temperature_unsupported_models.as_slice(), + native_tool_calling, + default_provider_options: default_provider_options.clone(), + extra_headers: extra_headers.as_slice(), + extra_query_params: extra_query_params.as_slice(), + user_agent: user_agent.as_deref(), + codex: codex_route, + http: http.clone(), + }) + }, + ); if let Some(identity) = inspection_identity { let credential = if entry.auth_style == AuthStyle::None { String::new() } else { - key.clone() + inspection_key.clone() }; let output_budget = inspection_output_budget_capability(&codex); // Doctor's live probe deduplicates on this identity, so it carries diff --git a/crates/openhuman-core/src/inference/provider/factory/credential_fleet.rs b/crates/openhuman-core/src/inference/provider/factory/credential_fleet.rs new file mode 100644 index 00000000000..ea01b0af076 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/credential_fleet.rs @@ -0,0 +1,451 @@ +//! Per-call credential selection and multimodal endpoint routing for cloud LLMs. +//! +//! A fleet wraps native models, never copies network adapters. Every call picks a +//! healthy credential and retries only failures known to be safe to reroute. +//! A live stream is NEVER replayed after emitting content; its terminal event +//! still updates the pool and drop releases the key lease. +use std::sync::Arc; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use tinyinference_llm::message::{ContentBlock, Message}; +use tinyinference_llm::model::{ + ChatModel, DeferredHandle, DeferredStatus, InputModality, InputSource, ModelProfile, + ModelRequest, ModelResponse, ModelStream, ModelStreamItem, ProviderError, +}; + +use super::key_pool::{self, KeyFailure, RotationStrategy}; + +// Guard against amplifying a provider-wide 401/429 incident across the entire +// configured key fleet. Four total attempts (initial + three fallbacks) per call. +const MAX_KEY_ATTEMPTS_PER_CALL: usize = 4; + +/// Detect actual image input (not merely a vision-capable model). +pub(super) fn request_has_image(request: &ModelRequest) -> bool { + request.messages.iter().any(|message| { + let blocks: &[ContentBlock] = match message { + Message::User(m) => &m.content, + Message::Assistant(m) => &m.content, + Message::Tool(m) => &m.content, + _ => &[], + }; + blocks + .iter() + .any(|block| matches!(block, ContentBlock::Image(_))) + }) +} + +#[derive(Clone, Copy)] +struct RetryDecision { + failure: KeyFailure, + after_ms: Option, +} +/// A 403 is ambiguous: providers also return it for unsupported models, +/// denied tools and malformed requests. Never evict a credential based on +/// loose substring matches (e.g. `invalid_model`, `unauthorized_model`). +/// Only documented credential-scoped code families permit another key. +fn forbidden_key_failure(code: &str) -> Option { + match code.trim().to_ascii_lowercase().as_str() { + "invalid_api_key" + | "invalid_api_key_error" + | "api_key_invalid" + | "invalid_token" + | "invalid_credentials" + | "authentication_error" + | "unauthorized" + | "unauthorized_client" + | "revoked" + | "key_revoked" + | "api_key_revoked" + | "token_revoked" => Some(KeyFailure::Authentication), + "quota_exhausted" + | "quota_exceeded" + | "insufficient_quota" + | "resource_exhausted" + | "rate_limit_exceeded" + | "rate_limit_error" + | "rate_limited" => Some(KeyFailure::RateLimited), + _ => None, + } +} +fn classify( + status: Option, + code: Option<&str>, + retry_after_ms: Option, +) -> Option { + let failure = match status { + Some(401) => KeyFailure::Authentication, + Some(429) => KeyFailure::RateLimited, + Some(403) => forbidden_key_failure(code?)?, + _ => return None, + }; + Some(RetryDecision { + failure, + after_ms: retry_after_ms, + }) +} +fn provider_decision(error: &ProviderError) -> Option { + // A partial answer belongs to the caller. Never replace it with another + // credential's replay, even if the terminal status normally permits retry. + if error.partial_message.is_some() { + return None; + } + classify(error.status, error.code.as_deref(), error.retry_after_ms).or_else(|| { + // Typed transient transport/server failures may try another credential + // before any content is exposed. They are not evidence of a bad key or + // exhausted quota, so preserve its eligibility for a later call. + (error.retryable && matches!(error.status, None | Some(408) | Some(500..=599))).then_some( + RetryDecision { + failure: KeyFailure::Other, + after_ms: error.retry_after_ms, + }, + ) + }) +} +fn model_decision(error: &tinyinference_llm::Error) -> Option { + match error { + tinyinference_llm::Error::Provider(e) => provider_decision(e), + _ => None, + } +} + +const FLEET_DEFERRED_METADATA_KEY: &str = "_openhuman_fleet_route_v1"; + +#[derive(Clone, Serialize, Deserialize)] +struct DeferredRoute { + version: u8, + key_fingerprint: String, + image: bool, +} + +/// A domain-separated digest preserves route identity across app restarts +/// without storing raw API keys in the provider-issued deferred handle. +fn credential_fingerprint(secret: &str) -> String { + let mut digest = Sha256::new(); + digest.update(b"openhuman-model-fleet-deferred-route-v1:"); + digest.update(secret.as_bytes()); + let bytes = digest.finalize(); + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} +fn stamp_deferred_handle(handle: &mut DeferredHandle, key: &str, image: bool) { + let route = DeferredRoute { + version: 1, + key_fingerprint: credential_fingerprint(key), + image, + }; + if let Ok(value) = serde_json::to_value(route) { + handle + .metadata + .insert(FLEET_DEFERRED_METADATA_KEY.to_owned(), value); + } +} +fn decode_deferred_handle(handle: &DeferredHandle) -> tinyinference_llm::Result { + let value = handle + .metadata + .get(FLEET_DEFERRED_METADATA_KEY) + .ok_or_else(|| { + tinyinference_llm::Error::Validation( + "deferred call is missing model-fleet routing metadata".into(), + ) + })?; + let route: DeferredRoute = serde_json::from_value(value.clone())?; + if route.version != 1 { + return Err(tinyinference_llm::Error::Validation( + "unsupported model-fleet deferred route version".into(), + )); + } + Ok(route) +} + +struct Lease { + scope: String, + key: String, +} +impl Lease { + /// The key-pool selector has already reserved this active lease atomically. + fn from_reserved(scope: &str, key: &str) -> Self { + Self { + scope: scope.to_owned(), + key: key.to_owned(), + } + } + fn record(&self, failure: Option) { + key_pool::record_result_for( + &self.scope, + &self.key, + failure.map(|failure| failure.failure), + failure.and_then(|failure| failure.after_ms), + ); + } +} +impl Drop for Lease { + fn drop(&mut self) { + key_pool::release(&self.scope, &self.key); + } +} + +/// Key-index alignment is fixed within this instance. The shared key_pool uses +/// raw secret identity for outcomes so reconfiguration in a parallel instance +/// cannot corrupt a different credential's health. +pub(super) struct CredentialFleetModel { + scope: String, + keys: Vec, + models: Vec>>, + image_models: Option>>>, + strategy: RotationStrategy, +} +impl CredentialFleetModel { + pub(super) fn new( + scope: String, + keys: Vec, + models: Vec>>, + image_models: Option>>>, + strategy: RotationStrategy, + ) -> Self { + assert!(!models.is_empty() && models.len() == keys.len()); + assert!(image_models + .as_ref() + .is_none_or(|m| m.len() == models.len())); + Self { + scope, + keys, + models, + image_models, + strategy, + } + } + fn route(&self, idx: usize, image: bool) -> &Arc> { + if image { + if let Some(ref group) = self.image_models { + return &group[idx]; + } + } + &self.models[idx] + } +} +#[async_trait] +impl ChatModel<()> for CredentialFleetModel { + fn profile(&self) -> Option<&ModelProfile> { + self.models.first().and_then(|model| model.profile()) + } + fn supports_input(&self, modality: InputModality, mime: &str, source: InputSource) -> bool { + // An alternate image route is only selected when the request actually + // carries an image. Audio/video/document inputs still use the primary + // models, so capability probes must reflect that same route. + let models = if modality == InputModality::Image { + self.image_models.as_ref().unwrap_or(&self.models) + } else { + &self.models + }; + models + .first() + .is_some_and(|model| model.supports_input(modality, mime, source)) + } + fn cache_identity(&self) -> Option { + self.models.first().and_then(|model| model.cache_identity()) + } + async fn invoke( + &self, + state: &(), + request: ModelRequest, + ) -> tinyinference_llm::Result { + let image = request_has_image(&request); + let mut tried = Vec::new(); + let mut last_error = None; + loop { + if tried.len() >= MAX_KEY_ATTEMPTS_PER_CALL { + break; + } + let Some((index, key)) = key_pool::choose_key_and_acquire_for_run( + &self.scope, + &self.keys, + &tried, + self.strategy, + request + .correlation + .as_ref() + .map(|correlation| correlation.run_id.as_str()), + ) else { + break; + }; + tried.push(index); + let lease = Lease::from_reserved(&self.scope, &key); + match self + .route(index, image) + .invoke(state, request.clone()) + .await + { + Ok(result) => { + lease.record(None); + return Ok(result); + } + Err(error) => { + let decision = model_decision(&error); + lease.record(Some(decision.unwrap_or(RetryDecision { + failure: KeyFailure::Other, + after_ms: None, + }))); + if decision.is_none() { + return Err(error); + } + last_error = Some(error); + } + } + } + Err(last_error.unwrap_or_else(|| tinyinference_llm::Error::Model( + "No available API keys for this provider; check rate limits or replace invalid keys".to_string() + ))) + } + async fn stream( + &self, + state: &(), + request: ModelRequest, + ) -> tinyinference_llm::Result { + use futures::StreamExt; + + let image = request_has_image(&request); + let mut tried = Vec::new(); + let mut last_error = None; + loop { + if tried.len() >= MAX_KEY_ATTEMPTS_PER_CALL { + break; + } + let Some((index, key)) = key_pool::choose_key_and_acquire_for_run( + &self.scope, + &self.keys, + &tried, + self.strategy, + request + .correlation + .as_ref() + .map(|correlation| correlation.run_id.as_str()), + ) else { + break; + }; + tried.push(index); + let lease = Lease::from_reserved(&self.scope, &key); + let mut live = match self + .route(index, image) + .stream(state, request.clone()) + .await + { + Ok(stream) => stream, + Err(error) => { + let decision = model_decision(&error); + lease.record(Some(decision.unwrap_or(RetryDecision { + failure: KeyFailure::Other, + after_ms: None, + }))); + if decision.is_none() { + return Err(error); + } + last_error = Some(error); + continue; + } + }; + + // OpenAI-compatible stream adapters can return a stream successfully + // and then emit an immediate 401/429 ProviderFailed event. Peek at + // the pre-content preamble so callers never see an attempt that + // should have been routed to another healthy credential. Never + // replay once content, usage or a deferred handle is exposed. + let metadata = live.metadata().clone(); + let mut prefix = Vec::new(); + loop { + match live.next().await { + Some(ModelStreamItem::Started) if prefix.len() < 4 => { + prefix.push(ModelStreamItem::Started); + } + Some(ModelStreamItem::ProviderFailed(error)) + if provider_decision(&error).is_some() => + { + lease.record(provider_decision(&error)); + last_error = Some(tinyinference_llm::Error::Provider(Box::new(error))); + break; + } + Some(item) => { + prefix.push(item); + let joined = + ModelStream::new(Box::pin(futures::stream::iter(prefix).chain(live))) + .with_metadata(metadata); + // The lease lives in the mapped stream. Dropping an + // unfinished stream still releases its concurrency slot. + let mut recorded = false; + return Ok(joined.map_items(move |mut item| { + if !recorded { + match &mut item { + ModelStreamItem::Completed(_) => { + lease.record(None); + recorded = true; + } + ModelStreamItem::Deferred(handle) => { + stamp_deferred_handle(handle, &lease.key, image); + lease.record(None); + recorded = true; + } + ModelStreamItem::ProviderFailed(error) => { + lease.record(Some(provider_decision(error).unwrap_or( + RetryDecision { + failure: KeyFailure::Other, + after_ms: None, + }, + ))); + recorded = true; + } + ModelStreamItem::Failed(_) => { + lease.record(Some(RetryDecision { + failure: KeyFailure::Other, + after_ms: None, + })); + recorded = true; + } + _ => {} + } + } + item + })); + } + None => { + lease.record(Some(RetryDecision { + failure: KeyFailure::Other, + after_ms: None, + })); + return Err(tinyinference_llm::Error::Model( + "provider stream ended before returning content or a terminal result" + .into(), + )); + } + } + } + } + Err(last_error.unwrap_or_else(|| tinyinference_llm::Error::Model( + "No available API keys for this provider; check rate limits or replace invalid keys".into() + ))) + } + + async fn fetch_deferred( + &self, + handle: &DeferredHandle, + ) -> tinyinference_llm::Result { + let route = decode_deferred_handle(handle)?; + let index = self + .keys + .iter() + .position(|secret| credential_fingerprint(secret) == route.key_fingerprint) + .ok_or_else(|| { + tinyinference_llm::Error::Validation( + "the credential used for this deferred call has been removed".into(), + ) + })?; + let mut original = handle.clone(); + original.metadata.remove(FLEET_DEFERRED_METADATA_KEY); + self.route(index, route.image) + .fetch_deferred(&original) + .await + } +} + +#[cfg(test)] +#[path = "credential_fleet_tests.rs"] +mod tests; diff --git a/crates/openhuman-core/src/inference/provider/factory/credential_fleet_tests.rs b/crates/openhuman-core/src/inference/provider/factory/credential_fleet_tests.rs new file mode 100644 index 00000000000..a597f6570dd --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/credential_fleet_tests.rs @@ -0,0 +1,464 @@ +use super::*; +use tinyinference_llm::message::{ImageRef, UserMessage}; +#[test] +fn image_endpoints_only_apply_when_request_contains_images() { + let text = ModelRequest { + messages: vec![Message::user("text")], + ..Default::default() + }; + assert!(!request_has_image(&text)); + let image = ModelRequest { + messages: vec![Message::User(UserMessage { + content: vec![ + ContentBlock::Text("What's here?".into()), + ContentBlock::Image(ImageRef { + url: "data:image/png;base64,AAAA".into(), + mime_type: Some("image/png".into()), + }), + ], + })], + ..Default::default() + }; + assert!(request_has_image(&image)); +} +struct CapabilityStubModel { + modality: InputModality, +} +#[async_trait] +impl ChatModel<()> for CapabilityStubModel { + fn supports_input(&self, modality: InputModality, _: &str, _: InputSource) -> bool { + modality == self.modality + } + + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + Err(tinyinference_llm::Error::Model( + "capability probe only".into(), + )) + } +} + +#[test] +fn fleet_reports_non_image_capabilities_of_the_model_that_receives_the_request() { + let fleet = CredentialFleetModel::new( + "fleet-capability-route".into(), + vec!["credential".into()], + vec![Arc::new(CapabilityStubModel { + modality: InputModality::Audio, + })], + Some(vec![Arc::new(CapabilityStubModel { + modality: InputModality::Image, + })]), + RotationStrategy::RoundRobin, + ); + + assert!(fleet.supports_input(InputModality::Image, "image/png", InputSource::Base64)); + assert!(fleet.supports_input(InputModality::Audio, "audio/wav", InputSource::Url)); + assert!(!fleet.supports_input(InputModality::Document, "application/pdf", InputSource::Url)); + assert!(!fleet.supports_input(InputModality::Video, "video/mp4", InputSource::Url)); +} + +struct StubStreamModel { + fail_auth: bool, +} +#[async_trait] +impl ChatModel<()> for StubStreamModel { + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + Err(tinyinference_llm::Error::Model( + "not used by stream test".into(), + )) + } + async fn stream(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + let events = if self.fail_auth { + vec![ + ModelStreamItem::Started, + ModelStreamItem::ProviderFailed(ProviderError { + status: Some(429), + ..Default::default() + }), + ] + } else { + vec![ + ModelStreamItem::Started, + ModelStreamItem::MessageDelta(tinyinference_llm::message::MessageDelta { + text: "selected-second-key".into(), + reasoning: String::new(), + tool_call: None, + }), + ModelStreamItem::Failed("test stream finished".into()), + ] + }; + Ok(ModelStream::new(Box::pin(futures::stream::iter(events)))) + } +} +#[tokio::test] +async fn pre_content_rate_limit_rotates_stream_before_exposing_error() { + use futures::StreamExt; + let scope = "fleet-stream-first-429"; + let fleet = CredentialFleetModel::new( + scope.into(), + vec!["first".into(), "second".into()], + vec![ + Arc::new(StubStreamModel { fail_auth: true }), + Arc::new(StubStreamModel { fail_auth: false }), + ], + None, + RotationStrategy::RoundRobin, + ); + let stream = fleet.stream(&(), ModelRequest::default()).await.unwrap(); + let events: Vec = stream.collect().await; + assert!(matches!( + events.get(1), + Some(ModelStreamItem::MessageDelta(_)) + )); + assert_eq!(key_pool::snapshot(scope)[0].status, "cooldown"); +} + +#[test] +fn deferred_route_is_keyed_to_credential_not_array_position() { + let mut handle = tinyinference_llm::model::DeferredHandle::new("openai", "resp_123"); + stamp_deferred_handle(&mut handle, "secret-one", true); + assert!(serde_json::to_string(&handle) + .unwrap() + .find("secret-one") + .is_none()); + let route = decode_deferred_handle(&handle).unwrap(); + assert_eq!(route.key_fingerprint, credential_fingerprint("secret-one")); + assert!(route.image); +} + +#[test] +fn provider_failures_are_rerouted_conservatively() { + assert_eq!( + classify(Some(401), None, None).unwrap().failure, + KeyFailure::Authentication + ); + assert_eq!( + classify(Some(429), None, Some(5000)).unwrap().failure, + KeyFailure::RateLimited + ); + assert_eq!( + classify(Some(403), Some("quota_exhausted"), None) + .unwrap() + .failure, + KeyFailure::RateLimited + ); + assert!(classify(Some(400), None, None).is_none()); + assert!(classify(Some(403), Some("forbidden"), None).is_none()); + // A provider can use HTTP 403 for an invalid model or request. A broad + // substring match for "invalid" must NOT evict an otherwise healthy API key. + for code in [ + "invalid_model", + "invalid_model_id", + "invalid_request", + "invalid_request_error", + "invalid_parameter", + "unauthorized_model", + "quota_configuration_error", + "rate_limit_policy_disabled", + ] { + assert!( + classify(Some(403), Some(code), None).is_none(), + "403 {code} is not evidence that another credential would succeed" + ); + } + for code in [ + "invalid_api_key", + "api_key_invalid", + "invalid_credentials", + "authentication_error", + "unauthorized", + "api_key_revoked", + "token_revoked", + ] { + assert_eq!( + classify(Some(403), Some(code), None).unwrap().failure, + KeyFailure::Authentication, + "{code} should evict only the rejected credential" + ); + } + for code in [ + "quota_exhausted", + "insufficient_quota", + "resource_exhausted", + "rate_limit_exceeded", + ] { + assert_eq!( + classify(Some(403), Some(code), None).unwrap().failure, + KeyFailure::RateLimited, + "{code} is a credential-scoped cooldown" + ); + } +} + +#[test] +fn typed_transient_failures_can_try_another_key_without_marking_it_invalid() { + for status in [None, Some(408), Some(503)] { + let error = ProviderError { + status, + retryable: true, + ..Default::default() + }; + assert_eq!( + provider_decision(&error).unwrap().failure, + KeyFailure::Other + ); + } + let permanent = ProviderError { + status: Some(503), + retryable: false, + ..Default::default() + }; + assert!(provider_decision(&permanent).is_none()); +} + +#[test] +fn a_partial_answer_is_never_rerouted_to_another_credential() { + let error = ProviderError { + status: Some(429), + retryable: true, + partial_message: Some(ModelResponse::assistant("partial answer").message), + ..Default::default() + }; + assert!(provider_decision(&error).is_none()); +} + +struct CountingForbiddenModel { + attempts: Arc, +} +#[async_trait] +impl ChatModel<()> for CountingForbiddenModel { + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + self.attempts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Err(tinyinference_llm::Error::Provider(Box::new( + ProviderError { + status: Some(403), + code: Some("invalid_model".into()), + ..Default::default() + }, + ))) + } + + async fn stream(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + self.attempts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(ModelStream::new(Box::pin(futures::stream::iter(vec![ + ModelStreamItem::Started, + ModelStreamItem::ProviderFailed(ProviderError { + status: Some(403), + code: Some("invalid_model".into()), + ..Default::default() + }), + ])))) + } +} +#[tokio::test] +async fn invalid_model_403_does_not_evict_api_keys_or_retry_another_credential() { + use futures::StreamExt; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let first = Arc::new(AtomicUsize::new(0)); + let second = Arc::new(AtomicUsize::new(0)); + let scope = "fleet-invalid-model-403-unary"; + let fleet = CredentialFleetModel::new( + scope.into(), + vec!["unary-first".into(), "unary-second".into()], + vec![ + Arc::new(CountingForbiddenModel { + attempts: Arc::clone(&first), + }), + Arc::new(CountingForbiddenModel { + attempts: Arc::clone(&second), + }), + ], + None, + RotationStrategy::RoundRobin, + ); + assert!(fleet.invoke(&(), ModelRequest::default()).await.is_err()); + assert_eq!(first.load(Ordering::SeqCst), 1); + assert_eq!(second.load(Ordering::SeqCst), 0); + assert_eq!(key_pool::snapshot(scope)[0].status, "healthy"); + key_pool::forget(scope); + + let first = Arc::new(AtomicUsize::new(0)); + let second = Arc::new(AtomicUsize::new(0)); + let scope = "fleet-invalid-model-403-stream"; + let fleet = CredentialFleetModel::new( + scope.into(), + vec!["stream-first".into(), "stream-second".into()], + vec![ + Arc::new(CountingForbiddenModel { + attempts: Arc::clone(&first), + }), + Arc::new(CountingForbiddenModel { + attempts: Arc::clone(&second), + }), + ], + None, + RotationStrategy::RoundRobin, + ); + let stream = fleet.stream(&(), ModelRequest::default()).await.unwrap(); + let output: Vec = stream.collect().await; + assert_eq!(first.load(Ordering::SeqCst), 1); + assert_eq!(second.load(Ordering::SeqCst), 0); + assert!(output + .iter() + .any(|event| matches!(event, ModelStreamItem::ProviderFailed(_)))); + assert_eq!(key_pool::snapshot(scope)[0].status, "healthy"); + key_pool::forget(scope); +} + +struct PostOutputRateLimitModel { + attempts: Arc, +} +#[async_trait] +impl ChatModel<()> for PostOutputRateLimitModel { + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + Err(tinyinference_llm::Error::Model("stream only".into())) + } + async fn stream(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + self.attempts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(ModelStream::new(Box::pin(futures::stream::iter(vec![ + ModelStreamItem::Started, + ModelStreamItem::MessageDelta(tinyinference_llm::message::MessageDelta::text( + "visible output", + )), + ModelStreamItem::ProviderFailed(ProviderError { + status: Some(429), + ..Default::default() + }), + ])))) + } +} + +#[tokio::test] +async fn post_content_upstream_429_cannot_replay_or_duplicate_visible_tool_output() { + use futures::StreamExt; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let first = Arc::new(AtomicUsize::new(0)); + let second = Arc::new(AtomicUsize::new(0)); + let scope = "fleet-postcontent-no-replay"; + let fleet = CredentialFleetModel::new( + scope.into(), + vec!["already-streamed".into(), "unused-backup".into()], + vec![ + Arc::new(PostOutputRateLimitModel { + attempts: Arc::clone(&first), + }), + Arc::new(PostOutputRateLimitModel { + attempts: Arc::clone(&second), + }), + ], + None, + RotationStrategy::RoundRobin, + ); + let events: Vec = fleet + .stream(&(), ModelRequest::default()) + .await + .unwrap() + .collect() + .await; + assert_eq!(first.load(Ordering::SeqCst), 1); + assert_eq!(second.load(Ordering::SeqCst), 0); + assert_eq!( + events + .iter() + .filter(|item| matches!(item, ModelStreamItem::MessageDelta(_))) + .count(), + 1, + "an emitted delta is never replayed across credentials" + ); + assert!( + matches!(events.last(), Some(ModelStreamItem::ProviderFailed(_))), + "the upstream failure must remain visible" + ); + assert_eq!(key_pool::snapshot(scope)[0].status, "cooldown"); + assert!(key_pool::snapshot(scope) + .iter() + .all(|entry| entry.active_leases == 0)); + key_pool::forget(scope); +} + +struct CountingRateLimitedModel { + attempts: Arc, +} + +#[async_trait] +impl ChatModel<()> for CountingRateLimitedModel { + async fn invoke(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + self.attempts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Err(tinyinference_llm::Error::Provider(Box::new( + ProviderError { + status: Some(429), + ..Default::default() + }, + ))) + } + + async fn stream(&self, _: &(), _: ModelRequest) -> tinyinference_llm::Result { + self.attempts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(ModelStream::new(Box::pin(futures::stream::iter(vec![ + ModelStreamItem::Started, + ModelStreamItem::ProviderFailed(ProviderError { + status: Some(429), + ..Default::default() + }), + ])))) + } +} + +#[tokio::test] +async fn fleet_fails_over_at_most_the_per_call_budget_for_unary_requests() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let scope = "fleet-bounded-429-unary"; + let attempts = Arc::new(AtomicUsize::new(0)); + let keys: Vec = (0..8).map(|index| format!("unary-key-{index}")).collect(); + let models: Vec>> = keys + .iter() + .map(|_| { + Arc::new(CountingRateLimitedModel { + attempts: Arc::clone(&attempts), + }) as Arc> + }) + .collect(); + let fleet = CredentialFleetModel::new( + scope.into(), + keys, + models, + None, + RotationStrategy::RoundRobin, + ); + assert!(fleet.invoke(&(), ModelRequest::default()).await.is_err()); + assert_eq!(attempts.load(Ordering::SeqCst), MAX_KEY_ATTEMPTS_PER_CALL); + key_pool::forget(scope); +} + +#[tokio::test] +async fn fleet_fails_over_at_most_the_per_call_budget_for_precontent_stream_failures() { + use std::sync::atomic::{AtomicUsize, Ordering}; + let scope = "fleet-bounded-429-stream"; + let attempts = Arc::new(AtomicUsize::new(0)); + let keys: Vec = (0..8).map(|index| format!("stream-key-{index}")).collect(); + let models: Vec>> = keys + .iter() + .map(|_| { + Arc::new(CountingRateLimitedModel { + attempts: Arc::clone(&attempts), + }) as Arc> + }) + .collect(); + let fleet = CredentialFleetModel::new( + scope.into(), + keys, + models, + None, + RotationStrategy::RoundRobin, + ); + assert!(fleet.stream(&(), ModelRequest::default()).await.is_err()); + assert_eq!(attempts.load(Ordering::SeqCst), MAX_KEY_ATTEMPTS_PER_CALL); + key_pool::forget(scope); +} diff --git a/crates/openhuman-core/src/inference/provider/factory/key_pool.rs b/crates/openhuman-core/src/inference/provider/factory/key_pool.rs new file mode 100644 index 00000000000..cd35f895ce3 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/key_pool.rs @@ -0,0 +1,557 @@ +//! Runtime credential pool shared by native BYOK model instances. +//! +//! Parsing mirrors Penguin's multi-key support; selection and outcomes are +//! maintained in Rust, where requests actually execute. This module never +//! exposes or logs secrets. Pool identity includes the local workspace. +use std::collections::{HashMap, HashSet}; +use std::hash::{BuildHasher, RandomState}; +use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +const DEFAULT_COOLDOWN: Duration = Duration::from_secs(60); +const MAX_COOLDOWN: Duration = Duration::from_secs(300); + +/// Progressive cooldown adopted from 9Router's account fallback design, while +/// preserving OpenHuman's existing 60s first-rate-limit policy. Provider +/// Retry-After wins when supplied and is capped to prevent years-long lockout. +fn quota_cooldown(streak: u32, retry_after_ms: Option) -> Duration { + if let Some(retry_ms) = retry_after_ms { + return Duration::from_millis(retry_ms).min(MAX_COOLDOWN); + } + let shift = streak.saturating_sub(1).min(3); + DEFAULT_COOLDOWN + .saturating_mul(1 << shift) + .min(MAX_COOLDOWN) +} +const SESSION_BINDING_TTL: Duration = Duration::from_secs(30 * 60); +const MAX_STICKY_RUN_ID_BYTES: usize = 256; + +#[derive(Clone)] +struct KeyHealth { + id: u64, + cooldown_until: Option, + evicted: bool, + successes: u64, + failures: u64, + active: usize, + last_used: Option, + consecutive_successes: u64, + consecutive_failures: u64, + /// Rate-limit-only streak for progressive quota backoff. General model + /// errors must not lengthen the next 429 cooldown. + consecutive_rate_limits: u32, +} +impl KeyHealth { + fn new(id: u64) -> Self { + Self { + id, + cooldown_until: None, + evicted: false, + successes: 0, + failures: 0, + active: 0, + last_used: None, + consecutive_successes: 0, + consecutive_failures: 0, + consecutive_rate_limits: 0, + } + } + fn available(&self, now: Instant) -> bool { + !self.evicted && self.cooldown_until.is_none_or(|until| until <= now) + } +} + +#[derive(Default)] +struct Pool { + keys: Vec, + next: usize, + sticky: HashMap, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum KeyFailure { + Authentication, + RateLimited, + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RotationStrategy { + #[default] + RoundRobin, + LeastBusy, + PriorityWeighted, +} +impl RotationStrategy { + pub fn parse(raw: &str) -> Self { + match raw { + "least_busy" | "least-leases" => Self::LeastBusy, + "priority_weighted" | "priority-weighted" => Self::PriorityWeighted, + _ => Self::RoundRobin, + } + } +} + +fn registry() -> &'static Mutex> { + static POOLS: OnceLock>> = OnceLock::new(); + POOLS.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// A process-salted SipHash identifier avoids retaining a second plaintext copy +/// of every credential in the health registry. The model instances themselves +/// retain only the keys they need to authenticate. +fn secret_id(value: &str) -> u64 { + static SALT: OnceLock = OnceLock::new(); + SALT.get_or_init(RandomState::new).hash_one(value) +} + +/// Process-local fingerprint for non-secret routing scope components. A randomized +/// process salt avoids persisting or exposing endpoint strings in registry keys. +/// Identical endpoints remain stable within one running process. +pub(super) fn scope_fingerprint(endpoint: &str) -> u64 { + secret_id(endpoint) +} + +/// Read a secret group from comma, semicolon, or newline-separated credentials. +/// Empty segments are discarded and identical values are kept once. +pub fn parse_api_keys(raw: &str) -> Vec { + let mut keys = Vec::new(); + for segment in raw.split([',', ';', '\n', '\r']) { + let value = segment.trim(); + if !value.is_empty() && !keys.iter().any(|key| key == value) { + keys.push(value.to_owned()); + } + } + keys +} + +fn sync(pool: &mut Pool, configured: &[String]) { + if pool.keys.len() == configured.len() + && pool + .keys + .iter() + .zip(configured) + .all(|(entry, value)| entry.id == secret_id(value)) + { + return; + } + let mut previous = std::mem::take(&mut pool.keys); + pool.keys = configured + .iter() + .map(|value| { + let id = secret_id(value); + match previous.iter().position(|entry| entry.id == id) { + Some(index) => previous.remove(index), + None => KeyHealth::new(id), + } + }) + .collect(); + pool.next = 0; + let ids: HashSet = pool.keys.iter().map(|key| key.id).collect(); + pool.sticky.retain(|_, (id, _)| ids.contains(id)); +} + +/// Select without returning a still-cooling or invalid key. +/// `excluded` avoids accidentally retrying the same key within one request. +/// Dynamic score from Penguin's WeightedKeyRotator: reward recent success, +/// penalize consecutive failure and active work, never give a healthy key +/// an entirely zero chance. Selection is proportional to these scores. +fn weighted_score(key: &KeyHealth) -> f64 { + (100.0 + key.consecutive_successes as f64 * 5.0 - key.consecutive_failures as f64 * 20.0) + .max(1.0) + / (1.0 + key.active as f64 * 2.0) +} + +pub fn choose_key( + scope: &str, + configured: &[String], + excluded: &[usize], + strategy: RotationStrategy, +) -> Option<(usize, String)> { + choose_key_for_run(scope, configured, excluded, strategy, None) +} + +/// Correlation/run affinity is TTL-bound and advisory: an evicted, cooling or +/// already-tried key will never be selected to preserve stickiness. +pub fn choose_key_for_run( + scope: &str, + configured: &[String], + excluded: &[usize], + strategy: RotationStrategy, + run_id: Option<&str>, +) -> Option<(usize, String)> { + choose_key_for_run_impl(scope, configured, excluded, strategy, run_id, false) +} + +/// Choose and reserve an active lease in one lock acquisition. +/// +/// Model Fleet callers must use this instead of a separate `choose_key` and +/// `acquire`: with least_busy, concurrent requests would otherwise all observe +/// zero active requests and stampede onto the same credential. +pub fn choose_key_and_acquire_for_run( + scope: &str, + configured: &[String], + excluded: &[usize], + strategy: RotationStrategy, + run_id: Option<&str>, +) -> Option<(usize, String)> { + choose_key_for_run_impl(scope, configured, excluded, strategy, run_id, true) +} + +fn choose_key_for_run_impl( + scope: &str, + configured: &[String], + excluded: &[usize], + strategy: RotationStrategy, + run_id: Option<&str>, + reserve_lease: bool, +) -> Option<(usize, String)> { + // Correlation IDs are user-influenced in JSON-RPC requests. They are + // advisory; unsafe or oversized IDs may still select a key, but must not + // persist unbounded strings in the 30-minute affinity cache. + let run_id = run_id.filter(|id| { + !id.is_empty() && id.len() <= MAX_STICKY_RUN_ID_BYTES && !id.chars().any(char::is_control) + }); + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let pool = pools.entry(scope.to_owned()).or_default(); + sync(pool, configured); + if pool.keys.is_empty() { + return None; + } + let now = Instant::now(); + pool.sticky.retain(|_, (_, expires)| *expires > now); + let candidates: Vec = (0..pool.keys.len()) + .filter(|index| !excluded.contains(index) && pool.keys[*index].available(now)) + .collect(); + if candidates.is_empty() { + return None; + } + let pinned = run_id + .filter(|id| !id.is_empty()) + .and_then(|id| pool.sticky.get(id)) + .and_then(|(stored, expires)| { + if *expires <= now { + return None; + } + candidates + .iter() + .copied() + .find(|idx| pool.keys[*idx].id == *stored) + }); + let selected = if let Some(pinned) = pinned { + pinned + } else { + match strategy { + RotationStrategy::RoundRobin => (0..pool.keys.len()) + .map(|offset| (pool.next + offset) % pool.keys.len()) + .find(|idx| candidates.contains(idx)), + RotationStrategy::LeastBusy => candidates.iter().copied().min_by_key(|idx| { + ( + pool.keys[*idx].active, + pool.keys[*idx].failures, + pool.keys[*idx].last_used, + ) + }), + RotationStrategy::PriorityWeighted => { + let weights: Vec = candidates + .iter() + .map(|idx| weighted_score(&pool.keys[*idx])) + .collect(); + let total: f64 = weights.iter().sum(); + let mut ticket = rand::random::() * total; + let mut winner = *candidates.last()?; + for (idx, weight) in candidates.iter().zip(weights.iter()) { + ticket -= weight; + if ticket <= 0.0 { + winner = *idx; + break; + } + } + Some(winner) + } + }? + }; + pool.next = (selected + 1) % pool.keys.len(); + pool.keys[selected].last_used = Some(now); + if reserve_lease { + // Do this before unlocking so the next least-busy selection observes + // the in-flight reservation; Lease::from_reserved owns its release. + pool.keys[selected].active += 1; + } + if let Some(run) = run_id.filter(|id| !id.is_empty()) { + // A bounded cache prevents unbounded correlation IDs in long-lived apps. + if pool.sticky.len() >= 512 && !pool.sticky.contains_key(run) { + // HashMap iteration order is arbitrary; evict the binding that + // expires first so active sessions retain their affinity. + if let Some(oldest) = pool + .sticky + .iter() + .min_by_key(|(_, (_, expires))| *expires) + .map(|(session, _)| session.clone()) + { + pool.sticky.remove(&oldest); + } + } + pool.sticky.insert( + run.to_owned(), + (pool.keys[selected].id, now + SESSION_BINDING_TTL), + ); + } + Some((selected, configured[selected].clone())) +} + +pub fn next_key_excluding( + scope: &str, + configured: &[String], + excluded: &[usize], +) -> Option<(usize, String)> { + choose_key(scope, configured, excluded, RotationStrategy::RoundRobin) +} +pub fn next_key(scope: &str, configured: &[String]) -> Option<(usize, String)> { + next_key_excluding(scope, configured, &[]) +} + +/// A lease is held through a full unary request or a stream's lifetime. +pub fn acquire(scope: &str, key: &str) { + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(entry) = pools + .get_mut(scope) + .and_then(|pool| pool.keys.iter_mut().find(|k| k.id == secret_id(key))) + { + entry.active += 1; + } +} +pub fn release(scope: &str, key: &str) { + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(entry) = pools + .get_mut(scope) + .and_then(|pool| pool.keys.iter_mut().find(|k| k.id == secret_id(key))) + { + entry.active = entry.active.saturating_sub(1); + } +} + +/// Attribute outcomes to secret identity, not a positional index that may have +/// changed after a settings update while the request was running. +pub fn record_result_for( + scope: &str, + key: &str, + failure: Option, + retry_after_ms: Option, +) { + record_result_by_id(scope, secret_id(key), failure, retry_after_ms); +} +fn record_result_by_id( + scope: &str, + id: u64, + failure: Option, + retry_after_ms: Option, +) { + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(entry) = pools + .get_mut(scope) + .and_then(|pool| pool.keys.iter_mut().find(|item| item.id == id)) + else { + return; + }; + match failure { + None => { + entry.successes += 1; + entry.consecutive_successes += 1; + entry.consecutive_failures = 0; + entry.consecutive_rate_limits = 0; + // A success from a parallel in-flight request is not evidence that + // another request's active 429 cooldown has been lifted. + entry.cooldown_until = entry.cooldown_until.filter(|until| *until > Instant::now()); + } + Some(KeyFailure::Authentication) => { + entry.failures += 1; + entry.consecutive_failures += 1; + entry.consecutive_successes = 0; + entry.consecutive_rate_limits = 0; + entry.evicted = true; + } + Some(KeyFailure::RateLimited) => { + entry.failures += 1; + entry.consecutive_failures += 1; + entry.consecutive_successes = 0; + entry.consecutive_rate_limits = entry.consecutive_rate_limits.saturating_add(1); + let delay = quota_cooldown(entry.consecutive_rate_limits, retry_after_ms); + let until = Instant::now() + delay; + // Competing 429 responses may specify different Retry-After delays. + // A later shorter delay must never preempt the longer deadline. + entry.cooldown_until = Some( + entry + .cooldown_until + .map_or(until, |previous| previous.max(until)), + ); + } + Some(KeyFailure::Other) => { + entry.failures += 1; + entry.consecutive_failures += 1; + entry.consecutive_successes = 0; + // Transport and request errors are not quota signals. + entry.consecutive_rate_limits = 0; + } + } +} +pub fn record_result(scope: &str, key_index: usize, failure: Option) { + let id = { + let pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + pools + .get(scope) + .and_then(|pool| pool.keys.get(key_index)) + .map(|key| key.id) + }; + if let Some(id) = id { + record_result_by_id(scope, id, failure, None); + } +} + +/// Safe summary used by settings/diagnostics; deliberately no masked or raw keys. +#[derive(Clone, Debug, serde::Serialize)] +pub struct KeyStats { + pub index: usize, + pub status: &'static str, + pub successes: u64, + pub failures: u64, + pub active_leases: usize, + pub cooldown_remaining_ms: u64, +} +pub fn snapshot(scope: &str) -> Vec { + let pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let now = Instant::now(); + pools + .get(scope) + .map(|pool| { + pool.keys + .iter() + .enumerate() + .map(|(i, key)| { + let cooldown = key + .cooldown_until + .and_then(|until| until.checked_duration_since(now)) + .unwrap_or_default() + .as_millis() as u64; + KeyStats { + index: i, + status: if key.evicted { + "evicted" + } else if cooldown > 0 { + "cooldown" + } else { + "healthy" + }, + successes: key.successes, + failures: key.failures, + active_leases: key.active, + cooldown_remaining_ms: cooldown, + } + }) + .collect() + }) + .unwrap_or_default() +} +pub fn health(scope: &str) -> Vec<(bool, bool, u64, u64)> { + snapshot(scope) + .iter() + .map(|s| { + ( + s.status == "evicted", + s.status == "cooldown", + s.successes, + s.failures, + ) + }) + .collect() +} +pub fn reset(scope: &str) { + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(pool) = pools.get_mut(scope) { + for entry in &mut pool.keys { + entry.evicted = false; + entry.cooldown_until = None; + entry.successes = 0; + entry.failures = 0; + entry.consecutive_failures = 0; + entry.consecutive_successes = 0; + entry.consecutive_rate_limits = 0; + } + } +} + +/// Health views for all currently initialized models in a provider group. +/// The model name is derived from the scope; no token identifier is exposed. +pub fn provider_snapshots(prefix: &str) -> Vec<(String, Vec)> { + let scopes: Vec = { + let pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + pools + .keys() + .filter(|name| name.starts_with(prefix)) + .cloned() + .collect() + }; + let mut snapshots: Vec<_> = scopes + .into_iter() + .map(|scope| { + let model = scope.strip_prefix(prefix).unwrap_or_default().to_owned(); + (model, snapshot(&scope)) + }) + .collect(); + snapshots.sort_by(|a, b| a.0.cmp(&b.0)); + snapshots +} + +/// Clearing a provider's state revives evicted keys without changing secrets. +pub fn reset_provider(prefix: &str) { + let scopes: Vec = { + let pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + pools + .keys() + .filter(|name| name.starts_with(prefix)) + .cloned() + .collect() + }; + for scope in scopes { + reset(&scope); + } +} + +/// Drop cached health for every model of a disconnected provider, without +/// disturbing other providers in the same workspace. +pub fn forget_provider(prefix: &str) { + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + pools.retain(|scope, _| !scope.starts_with(prefix)); +} + +/// Called when a model is removed or a provider is replaced. +pub fn forget(scope: &str) { + let mut pools = registry() + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + pools.remove(scope); +} + +#[cfg(test)] +#[path = "key_pool_tests.rs"] +mod tests; diff --git a/crates/openhuman-core/src/inference/provider/factory/key_pool_tests.rs b/crates/openhuman-core/src/inference/provider/factory/key_pool_tests.rs new file mode 100644 index 00000000000..fd0133cd6b0 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory/key_pool_tests.rs @@ -0,0 +1,340 @@ +use super::*; +#[test] +fn split_deduplicates_and_discards_blanks() { + assert_eq!( + parse_api_keys(" key-a, key-b;key-a\nkey-c\r\n"), + vec!["key-a", "key-b", "key-c"] + ); +} +#[test] +fn routes_around_cooldown_and_revoked_keys() { + let scope = "key-pool-test-routing"; + let keys = vec!["a".to_owned(), "b".to_owned()]; + let (i, first) = next_key(scope, &keys).unwrap(); + assert_eq!(first, "a"); + record_result(scope, i, Some(KeyFailure::RateLimited)); + assert_eq!(next_key(scope, &keys).unwrap().1, "b"); + record_result(scope, 1, Some(KeyFailure::Authentication)); + assert!(next_key(scope, &keys).is_none()); + assert_eq!(health(scope).len(), 2); +} +#[test] +fn exhaustion_never_reuses_an_attempted_key() { + let scope = "key-pool-test-exhausted"; + let keys = vec!["a".to_owned(), "b".to_owned()]; + let (first, _) = next_key_excluding(scope, &keys, &[]).unwrap(); + let (second, _) = next_key_excluding(scope, &keys, &[first]).unwrap(); + assert_ne!(first, second); + assert!(next_key_excluding(scope, &keys, &[first, second]).is_none()); +} +#[test] +fn outcome_is_attributed_by_secret_not_stale_array_position() { + let scope = "key-pool-test-rotated"; + let initial = vec!["a".to_owned(), "b".to_owned()]; + let (_, key) = next_key(scope, &initial).unwrap(); + let updated = vec!["b".to_owned(), "a".to_owned()]; + next_key(scope, &updated).unwrap(); + record_result_for(scope, &key, Some(KeyFailure::Authentication), None); + assert_eq!(next_key(scope, &updated).unwrap().1, "b"); +} +#[test] +fn updates_pool_without_losing_existing_health() { + let scope = "key-pool-test-update"; + let initial = vec!["one".to_owned(), "two".to_owned()]; + next_key(scope, &initial); + record_result(scope, 0, Some(KeyFailure::Authentication)); + let updated = vec!["one".to_owned(), "three".to_owned()]; + assert_eq!(next_key(scope, &updated).unwrap().1, "three"); + assert!(health(scope)[0].0); +} +#[test] +fn one_run_reuses_healthy_key_but_fails_over_after_rate_limit() { + let scope = "key-pool-test-session-affinity"; + let keys = vec!["key-a".to_owned(), "key-b".to_owned()]; + let (first, _) = choose_key_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some("run-1"), + ) + .unwrap(); + let (sticky, _) = choose_key_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some("run-1"), + ) + .unwrap(); + assert_eq!(first, sticky); + record_result(scope, first, Some(KeyFailure::RateLimited)); + let (fallback, _) = choose_key_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some("run-1"), + ) + .unwrap(); + assert_ne!(fallback, first); +} +#[test] +fn weighted_scoring_prefers_unleased_successful_credentials() { + let fresh = KeyHealth::new(secret_id("fresh")); + let mut loaded = fresh.clone(); + loaded.active = 2; + assert!(weighted_score(&fresh) > weighted_score(&loaded)); + loaded.active = 0; + loaded.consecutive_failures = 2; + assert!(weighted_score(&fresh) > weighted_score(&loaded)); +} +#[test] +fn pool_retains_only_salted_key_ids_not_raw_secrets() { + let scope = "key-pool-test-no-secrets"; + let configured = vec!["sensitive_example_token".to_owned()]; + next_key(scope, &configured); + let pools = registry().lock().unwrap(); + let key = &pools.get(scope).unwrap().keys[0]; + assert_eq!(key.id, secret_id("sensitive_example_token")); + assert_ne!(key.id, secret_id("different_token")); +} +#[test] +fn deleted_model_scope_does_not_remove_sibling_key_health() { + let prefix = "key-pool-test-model-prune:"; + let a = format!("{prefix}model-a"); + let b = format!("{prefix}model-b"); + next_key(&a, &["a".to_owned()]); + next_key(&b, &["b".to_owned()]); + forget(&a); + let remaining = provider_snapshots(prefix); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].0, "model-b"); + forget_provider(prefix); +} +#[test] +fn deleted_provider_scopes_are_pruned() { + let prefix = "key-pool-test-prune:"; + next_key("key-pool-test-prune:model-a", &["a".to_owned()]); + next_key("key-pool-test-prune:model-b", &["b".to_owned()]); + forget_provider(prefix); + assert!(provider_snapshots(prefix).is_empty()); +} +#[test] +fn least_busy_reserves_active_lease_before_returning_the_key() { + // The two calls simulate simultaneous requests that do not yield between + // choosing their credentials. The first selection must reserve its lease + // before the second request makes its choice, otherwise both select A. + let scope = "key-pool-test-atomic-reservation"; + let keys = vec!["first".to_owned(), "second".to_owned()]; + let (first_index, first) = + choose_key_and_acquire_for_run(scope, &keys, &[], RotationStrategy::LeastBusy, None) + .expect("first credential available"); + assert_eq!(first_index, 0); + assert_eq!(snapshot(scope)[first_index].active_leases, 1); + + let (second_index, second) = + choose_key_and_acquire_for_run(scope, &keys, &[], RotationStrategy::LeastBusy, None) + .expect("second credential available"); + assert_ne!( + first_index, second_index, + "lease reservation must influence the next selection" + ); + assert_eq!(snapshot(scope)[second_index].active_leases, 1); + + release(scope, &first); + release(scope, &second); + assert!(snapshot(scope).iter().all(|row| row.active_leases == 0)); + forget(scope); +} + +#[test] +fn reservations_honor_excluded_keys_and_run_affinity_without_double_acquire() { + let scope = "key-pool-test-atomic-exclusions"; + let keys = vec!["a".to_owned(), "b".to_owned()]; + let (selected, secret) = choose_key_and_acquire_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some("session"), + ) + .unwrap(); + let (sticky, again) = choose_key_and_acquire_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some("session"), + ) + .unwrap(); + assert_eq!(selected, sticky, "healthy session binding stays stable"); + assert_eq!(secret, again); + assert_eq!(snapshot(scope)[selected].active_leases, 2); + assert_eq!( + choose_key_and_acquire_for_run(scope, &keys, &[0, 1], RotationStrategy::RoundRobin, None), + None, + "exhausted candidate set cannot reserve another lease" + ); + + release(scope, &secret); + release(scope, &again); + assert_eq!(snapshot(scope)[selected].active_leases, 0); + forget(scope); +} + +#[test] +fn leases_are_released_without_underflow() { + let scope = "key-pool-test-lease"; + let keys = vec!["a".to_owned()]; + next_key(scope, &keys); + acquire(scope, "a"); + assert_eq!(snapshot(scope)[0].active_leases, 1); + release(scope, "a"); + release(scope, "a"); + assert_eq!(snapshot(scope)[0].active_leases, 0); +} + +#[test] +fn late_success_must_not_clear_an_active_rate_limit_from_another_lease() { + let scope = "key-pool-test-concurrent-429-success"; + let secret = "same-provider-credential"; + let keys = vec![secret.to_owned()]; + assert!(next_key(scope, &keys).is_some()); + record_result_for(scope, secret, Some(KeyFailure::RateLimited), Some(120_000)); + assert_eq!(snapshot(scope)[0].status, "cooldown"); + + // Another request was already in flight and returned 200 after the 429. + // This is NOT evidence that provider's cooldown has expired. + record_result_for(scope, secret, None, None); + assert_eq!(snapshot(scope)[0].status, "cooldown"); + assert!(next_key(scope, &keys).is_none()); + forget(scope); +} + +#[test] +fn parallel_rate_limit_outcomes_cannot_shorten_the_existing_deadline() { + let scope = "key-pool-test-parallel-429"; + let secret = "rate-limited-key"; + next_key(scope, &[secret.to_owned()]); + record_result_for(scope, secret, Some(KeyFailure::RateLimited), Some(240_000)); + let first_deadline = { + let pools = registry().lock().unwrap(); + pools.get(scope).unwrap().keys[0].cooldown_until.unwrap() + }; + record_result_for(scope, secret, Some(KeyFailure::RateLimited), Some(1000)); + let second_deadline = { + let pools = registry().lock().unwrap(); + pools.get(scope).unwrap().keys[0].cooldown_until.unwrap() + }; + assert!( + second_deadline >= first_deadline, + "an overlapping 429 shortened an active cooldown" + ); + forget(scope); +} + +#[test] +fn session_affinity_cache_evicts_earliest_expiring_binding() { + let scope = "key-pool-test-affinity-order"; + let keys = vec!["stable".to_owned()]; + next_key(scope, &keys); + { + let mut pools = registry().lock().unwrap(); + let pool = pools.get_mut(scope).unwrap(); + let base = Instant::now(); + let id = secret_id("stable"); + for i in 0..512u64 { + pool.sticky + .insert(format!("run-{i}"), (id, base + Duration::from_secs(60 + i))); + } + } + + choose_key_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some("run-new"), + ) + .unwrap(); + { + let pools = registry().lock().unwrap(); + let sticky = &pools.get(scope).unwrap().sticky; + assert_eq!(sticky.len(), 512); + assert!( + !sticky.contains_key("run-0"), + "evict oldest/earliest expiry, not arbitrary HashMap entry" + ); + assert!(sticky.contains_key("run-511")); + assert!(sticky.contains_key("run-new")); + } + forget(scope); +} + +#[test] +fn untrusted_correlation_ids_do_not_fill_session_affinity_with_unbounded_strings() { + let scope = "key-pool-test-large-correlation"; + let keys = vec!["safe".to_owned()]; + let long_id = "r".repeat(64 * 1024); + for run_id in [&long_id[..], "newline\nseparator", "control\u{7}bell"] { + assert!(choose_key_for_run( + scope, + &keys, + &[], + RotationStrategy::RoundRobin, + Some(run_id), + ) + .is_some()); + } + let pools = registry().lock().unwrap(); + assert!(pools.get(scope).unwrap().sticky.is_empty()); + drop(pools); + forget(scope); + // No test-only correlation scope may leak into the process-wide registry. + assert!(!registry().lock().unwrap().contains_key(scope)); +} + +#[test] +fn quota_cooldown_exponentially_backs_off_without_overriding_server_retry_after() { + assert_eq!(quota_cooldown(1, None), Duration::from_secs(60)); + assert_eq!(quota_cooldown(2, None), Duration::from_secs(120)); + assert_eq!(quota_cooldown(3, None), Duration::from_secs(240)); + assert_eq!(quota_cooldown(4, None), Duration::from_secs(300)); + assert_eq!(quota_cooldown(u32::MAX, None), Duration::from_secs(300)); + assert_eq!(quota_cooldown(3, Some(1_500)), Duration::from_millis(1_500)); + assert_eq!(quota_cooldown(3, Some(u64::MAX)), Duration::from_secs(300)); +} + +#[test] +fn repeated_rate_limits_increase_default_cooldown_without_exposing_secrets() { + let scope = "key-pool-test-progressive-429"; + let key = "never-log-this-secret".to_string(); + next_key(scope, std::slice::from_ref(&key)); + record_result_for(scope, &key, Some(KeyFailure::RateLimited), None); + let first = snapshot(scope)[0].cooldown_remaining_ms; + assert!(first > 55_000 && first <= 60_000); + record_result_for(scope, &key, Some(KeyFailure::RateLimited), None); + let second = snapshot(scope)[0].cooldown_remaining_ms; + assert!(second > 115_000 && second <= 120_000); + reset(scope); + record_result_for(scope, &key, Some(KeyFailure::RateLimited), None); + let after_reset = snapshot(scope)[0].cooldown_remaining_ms; + assert!(after_reset > 55_000 && after_reset <= 60_000); +} + +#[test] +fn explicit_retry_after_is_honored_and_concurrent_shorter_limits_do_not_erase_it() { + let scope = "key-pool-test-progressive-429-explicit"; + let key = "another-secret".to_string(); + next_key(scope, std::slice::from_ref(&key)); + record_result_for(scope, &key, Some(KeyFailure::RateLimited), Some(250_000)); + let first = snapshot(scope)[0].cooldown_remaining_ms; + record_result_for(scope, &key, Some(KeyFailure::RateLimited), Some(2_000)); + let second = snapshot(scope)[0].cooldown_remaining_ms; + assert!(first > 245_000 && first <= 250_000); + assert!( + second > 244_000, + "later 429 must not shorten existing cooldown" + ); +} diff --git a/crates/openhuman-core/src/inference/provider/factory_credential_fleet_tests.rs b/crates/openhuman-core/src/inference/provider/factory_credential_fleet_tests.rs new file mode 100644 index 00000000000..96b588cf718 --- /dev/null +++ b/crates/openhuman-core/src/inference/provider/factory_credential_fleet_tests.rs @@ -0,0 +1,104 @@ +use super::*; + +#[test] +fn changing_a_byok_endpoint_does_not_inherit_the_previous_credential_cooldown() { + let workspace = std::path::Path::new("/test/workspace"); + let original = cloud_slug::rotation_pool_scope( + workspace, + "user-gateway", + "same-model", + "https://first.example/v1", + ); + let changed = cloud_slug::rotation_pool_scope( + workspace, + "user-gateway", + "same-model", + "https://second.example/v1", + ); + assert_ne!( + original, changed, + "endpoint change must isolate pool health" + ); + assert_eq!( + original, + cloud_slug::rotation_pool_scope( + workspace, + "user-gateway", + "same-model", + "https://first.example/v1", + ) + ); + let keys = vec!["unchanged-secret".to_owned()]; + let (index, _) = key_pool::next_key(&original, &keys).unwrap(); + key_pool::record_result(&original, index, Some(key_pool::KeyFailure::RateLimited)); + assert!( + key_pool::next_key(&original, &keys).is_none(), + "original endpoint must remain cooled down" + ); + assert!( + key_pool::next_key(&changed, &keys).is_some(), + "a new destination may not inherit the old destination's rate limit" + ); + key_pool::forget(&original); + key_pool::forget(&changed); +} + +#[tokio::test] +async fn byok_multi_key_route_retries_pre_output_429_with_an_individual_credential() { + use tinyinference_llm::message::Message; + use tinyinference_llm::model::ModelRequest; + use wiremock::matchers::header; + + let _guard = crate::inference::inference_test_guard_async().await; + let _env = crate::config::TEST_ENV_LOCK.lock().await; + let server = MockServer::start().await; + let mut provider = openai_entry("p_rotation_test", "rotator-test"); + provider.endpoint = format!("{}/v1", server.uri()); + provider.default_model = Some("test-model".to_owned()); + + let tmp = TempDir::new().expect("isolated credential workspace"); + let mut config = config_with_providers_in_tempdir(&tmp, vec![provider]); + config.chat_provider = Some("rotator-test:test-model".to_owned()); + store_byo_key(&config, "rotator-test", "fleet-key-a;fleet-key-b"); + + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer fleet-key-a")) + .respond_with(ResponseTemplate::new(429).set_body_json(serde_json::json!({ + "error": {"message": "rate limited", "type": "rate_limit_error", "code": "rate_limit_exceeded"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer fleet-key-b")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "chatcmpl-rotation-test", + "object": "chat.completion", + "created": 1770000000, + "model": "test-model", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .expect(1) + .mount(&server) + .await; + + let model = + create_chat_model_with_model_id_in(&fixture_library_dispatch(), "chat", &config, 0.0) + .expect("create native provider model") + .0; + let result = model + .invoke(&(), ModelRequest::new(vec![Message::user("hello")])) + .await; + assert!( + result.is_ok(), + "429 should fail over before exposing output: {result:?}" + ); + server.verify().await; +} diff --git a/crates/openhuman-core/src/inference/provider/factory_tests.rs b/crates/openhuman-core/src/inference/provider/factory_tests.rs index 66c2b7c89e5..92c6ab87526 100644 --- a/crates/openhuman-core/src/inference/provider/factory_tests.rs +++ b/crates/openhuman-core/src/inference/provider/factory_tests.rs @@ -626,6 +626,8 @@ mod crate_native_diagnostics_tests; mod crate_native_jwt_tests; #[path = "factory_crate_native_tests.rs"] mod crate_native_tests; +#[path = "factory_credential_fleet_tests.rs"] +mod credential_fleet_tests; #[path = "factory_egress_fallback_tests.rs"] mod egress_fallback_tests; #[path = "factory_explicit_ctx_tests.rs"] diff --git a/crates/openhuman-core/src/inference/provider/models/catalog_listing.rs b/crates/openhuman-core/src/inference/provider/models/catalog_listing.rs index d4608100bd3..ae9f2096aee 100644 --- a/crates/openhuman-core/src/inference/provider/models/catalog_listing.rs +++ b/crates/openhuman-core/src/inference/provider/models/catalog_listing.rs @@ -53,24 +53,55 @@ fn catalog_request( ) -> Result { let service_key = "providers.list_models"; let slug = entry.map(|e| e.slug.trim()).unwrap_or_default(); + use crate::config::schema::cloud_providers::AuthStyle; let provider_key = match entry { Some(entry) => { let looked_up = crate::inference::provider::factory::lookup_key_for_slug(&entry.slug, config) .unwrap_or_default(); - resolve_local_runtime_key(&entry.slug, looked_up, config) + let raw = resolve_local_runtime_key(&entry.slug, looked_up, config); + // A model catalog request can authenticate with only one key. + // Never send the raw delimited fleet as one Authorization header. + // OAuth credentials are opaque and must not be interpreted as a fleet. + let byok = matches!(entry.auth_style, AuthStyle::Bearer | AuthStyle::Anthropic) + && !(entry.slug == "openai" + && crate::inference::provider::factory::openai_bearer_is_oauth(config)); + if byok { + let keys = crate::inference::provider::factory::byok_catalog_keys(&raw); + if keys.len() > 1 && crate::core::runtime::is_saas() { + return Err( + "multi-key BYOK catalog discovery is unavailable in shared mode".into(), + ); + } + keys.into_iter().next().unwrap_or_default() + } else { + raw + } } None => String::new(), }; - let client = match config + let pem = config .cloud_provider_ca_certs .get(slug) - .filter(|pem| !pem.is_empty()) - { - Some(pem) => { - crate::util::tls::client_with_ca_bundle_with_timeouts(pem, service_key, 30, 10)? + .map(String::as_str) + .filter(|value| !value.is_empty()); + let credentialed = entry + .is_some_and(|entry| matches!(entry.auth_style, AuthStyle::Bearer | AuthStyle::Anthropic)); + // Catalog discovery carries the same credential headers as inference. + // Its HTTP client must not forward a custom header over redirects. + let client = if credentialed { + crate::util::tls::credentialed_client_builder(service_key, pem)? + .timeout(std::time::Duration::from_secs(30)) + .connect_timeout(std::time::Duration::from_secs(10)) + .build() + .map_err(|_| "Could not build secure model catalog HTTP client".to_string())? + } else { + match pem { + Some(pem) => { + crate::util::tls::client_with_ca_bundle_with_timeouts(pem, service_key, 30, 10)? + } + None => crate::config::build_runtime_proxy_client_with_timeouts(service_key, 30, 10), } - None => crate::config::build_runtime_proxy_client_with_timeouts(service_key, 30, 10), }; Ok(crate::inference::registry::CatalogHostRequest { provider_key, diff --git a/crates/openhuman-core/src/inference/provider/models/catalog_listing_error_tests.rs b/crates/openhuman-core/src/inference/provider/models/catalog_listing_error_tests.rs index 138f9c5e8d5..b5d646c3a07 100644 --- a/crates/openhuman-core/src/inference/provider/models/catalog_listing_error_tests.rs +++ b/crates/openhuman-core/src/inference/provider/models/catalog_listing_error_tests.rs @@ -113,6 +113,33 @@ async fn openrouter_key_is_trimmed_for_validation_and_catalog_probe() { ); } +#[tokio::test] +async fn multi_key_catalog_uses_one_bearer_and_never_forwards_the_group() { + let tmp = tempfile::tempdir().expect("catalog workspace"); + let (endpoint, state) = spawn_openrouter_probe_server(StatusCode::OK).await; + let config = + configure_openrouter_workspace(&tmp, endpoint, "catalog-key-first;catalog-key-second") + .await; + + let outcome = catalog_listing::list_configured_models_from_config("openrouter", &config) + .await + .expect("one credential must be used for catalog discovery"); + assert_eq!(outcome.value["models"][0]["id"], "openrouter/test-model"); + + let key_headers = state + .key_authorization + .lock() + .unwrap_or_else(|e| e.into_inner()) + .clone(); + let model_headers = state + .model_authorization + .lock() + .unwrap_or_else(|e| e.into_inner()) + .clone(); + assert_eq!(key_headers, vec![Some("Bearer catalog-key-first".into())]); + assert_eq!(model_headers, key_headers); +} + #[test] fn test_sanitize_api_error_utf8() { let input = "🦀".repeat(MAX_API_ERROR_CHARS + 10); diff --git a/crates/openhuman-core/src/util/tls/mod.rs b/crates/openhuman-core/src/util/tls/mod.rs index ed468fe8108..e05ecfd7f01 100644 --- a/crates/openhuman-core/src/util/tls/mod.rs +++ b/crates/openhuman-core/src/util/tls/mod.rs @@ -56,6 +56,28 @@ fn ca_client_builder(pem: &str, service_key: &str) -> Result, +) -> Result { + let mut builder = crate::config::apply_runtime_proxy_to_builder( + tls_client_builder().redirect(reqwest::redirect::Policy::none()), + service_key, + ); + if let Some(pem) = ca_pem.filter(|value| !value.is_empty()) { + for certificate in parse_ca_bundle(pem)? { + builder = builder.add_root_certificate(certificate); + } + } + Ok(builder) +} + pub fn client_with_ca_bundle(pem: &str, service_key: &str) -> Result { let builder = ca_client_builder(pem, service_key)?; builder diff --git a/crates/openhuman-core/src/util/tls/mod_tests.rs b/crates/openhuman-core/src/util/tls/mod_tests.rs index 8036a729f91..c83808a3d0d 100644 --- a/crates/openhuman-core/src/util/tls/mod_tests.rs +++ b/crates/openhuman-core/src/util/tls/mod_tests.rs @@ -20,3 +20,46 @@ fn invalid_ca_bundle_and_private_key_material_are_rejected() { let private_key_marker = concat!("-----BEGIN ", "PRIVATE KEY-----"); assert!(parse_ca_bundle(private_key_marker).is_err()); } + +#[tokio::test] +async fn authenticated_http_client_never_forwards_x_api_key_to_a_redirect() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let destination = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/capture")) + .respond_with(ResponseTemplate::new(200)) + .mount(&destination) + .await; + + let origin = MockServer::start().await; + let redirected_to = format!("{}/capture", destination.uri()); + Mock::given(method("GET")) + .and(path("/model")) + .respond_with(ResponseTemplate::new(302).insert_header("Location", redirected_to.as_str())) + .expect(1) + .mount(&origin) + .await; + + let http = credentialed_client_builder("provider.anthropic", None) + .unwrap() + .build() + .unwrap(); + let result = http + .get(format!("{}/model", origin.uri())) + .header("x-api-key", "provider-secret") + .send() + .await + .unwrap(); + assert_eq!(result.status(), reqwest::StatusCode::FOUND); + assert!(destination.received_requests().await.unwrap().is_empty()); +} + +#[test] +fn authenticated_http_client_accepts_additional_certificate_roots() { + let builder = + credentialed_client_builder("provider.compatible", Some(include_str!("test-ca.pem"))) + .unwrap(); + assert!(builder.build().is_ok()); +} diff --git a/scripts/ci/process-statics.json b/scripts/ci/process-statics.json index ed07402fb93..6d8fff343b4 100644 --- a/scripts/ci/process-statics.json +++ b/scripts/ci/process-statics.json @@ -1,4 +1,24 @@ [ + { + "file": "crates/openhuman-core/src/agent_registry/host_definition.rs", + "scope": "", + "name": "GLOBAL", + "type": "OnceLock>", + "count": 1, + "category": "A", + "reason": "The process registry is operator configuration in single-user hosts; SaaS startup installs built-ins only and refuses workspace or home overrides, so it cannot carry one SaaS profile's definitions into another.", + "verification": "core/runtime/saas_tests.rs: boot_refuses_a_registry_with_workspace_definitions" + }, + { + "file": "crates/openhuman-core/src/agent_runs/context_breakdown.rs", + "scope": "", + "name": "REPORT_CACHE", + "type": "Lazy>>", + "count": 1, + "category": "B", + "reason": "Each cache key includes the dispatch tenant (profile and agent), requested agent id, workspace, and config path; the serialized-config fingerprint only validates freshness within that scope.", + "verification": "agent_runs/context_breakdown_tests.rs: report_cache_key_isolated_by_tenant_even_for_identical_config" + }, { "file": "crates/openhuman-core/src/agent/hooks.rs", "scope": "", @@ -49,6 +69,16 @@ "reason": "Each in-flight delivery lease key combines the calling profile key with its thread id; it stores only active per-thread exclusion keys and no completion contents.", "verification": "agent/orchestration/background_delivery_slots_tests.rs: one_profiles_busy_turn_does_not_block_anothers_thread" }, + { + "file": "crates/openhuman-core/src/agent/session_lease.rs", + "scope": "", + "name": "LEASES", + "type": "LazyLock>>>>", + "count": 1, + "category": "D", + "reason": "Process-local lease table keyed by the complete tenant, agent and canonical session identity; it holds only weak mutex handles that coordinate runners inside this process and stores no tenant data.", + "verification": "agent/session_lease_tests.rs: distinct_tenant_agent_or_session_keys_do_not_serialize" + }, { "file": "crates/openhuman-core/src/agent/tinyagents/turn_exec/root_hosted_harness.rs", "scope": "", @@ -59,26 +89,6 @@ "reason": "One unconfigured harness is process infrastructure; per-invocation models, tools, middleware, progress, and host capabilities are supplied through each invocation, so the singleton carries no tenant authority or data.", "verification": "agent/tinyagents/turn_exec_tests.rs: hosted_roots_share_only_an_unconfigured_process_harness" }, - { - "file": "crates/openhuman-core/src/agent_registry/host_definition.rs", - "scope": "", - "name": "GLOBAL", - "type": "OnceLock>", - "count": 1, - "category": "A", - "reason": "The process registry is operator configuration in single-user hosts; SaaS startup installs built-ins only and refuses workspace or home overrides, so it cannot carry one SaaS profile's definitions into another.", - "verification": "core/runtime/saas_tests.rs: boot_refuses_a_registry_with_workspace_definitions" - }, - { - "file": "crates/openhuman-core/src/agent_runs/context_breakdown.rs", - "scope": "", - "name": "REPORT_CACHE", - "type": "Lazy>>", - "count": 1, - "category": "B", - "reason": "Each cache key includes the dispatch tenant (profile and agent), requested agent id, workspace, and config path; the serialized-config fingerprint only validates freshness within that scope.", - "verification": "agent_runs/context_breakdown_tests.rs: report_cache_key_isolated_by_tenant_even_for_identical_config" - }, { "file": "crates/openhuman-core/src/artifacts/publisher_slot.rs", "scope": "", @@ -149,6 +159,16 @@ "reason": "Host-provided provider preset metadata is process configuration; runtime-specific additions are kept in CoreContext extensions, and no credentials are stored here.", "verification": "config/schema/cloud_providers_tests.rs: hosted_presets_and_classification_are_context_scoped" }, + { + "file": "crates/openhuman-core/src/core/observability_reporting.rs", + "scope": "", + "name": "LAST_FS_LIMIT_REPORT", + "type": "std::sync::LazyLock>, >", + "count": 1, + "category": "A", + "reason": "A process-wide rate limiter keyed by the content-free domain and operation names prevents repeated filesystem-limit reports from flooding observability. Entries older than the five-minute cooldown are pruned on every access, and the map stores no tenant or user data.", + "verification": "core/observability_reporting.rs: was_recently_reported" + }, { "file": "crates/openhuman-core/src/core/observability/telemetry/langfuse/environment.rs", "scope": "", @@ -199,6 +219,26 @@ "reason": "The module's shipped queues.json and pipeline.json, parsed once; immutable process configuration with no tenant data.", "verification": "inbound/settings_tests.rs: the_defaults_resolve_to_the_shipped_files_in_storage" }, + { + "file": "crates/openhuman-core/src/inference/provider/factory/key_pool.rs", + "scope": "", + "name": "POOLS", + "type": "OnceLock>>", + "count": 1, + "category": "B", + "reason": "Process-local credential-health registry, keyed by the caller-supplied scope. It stores process-salted secret IDs and bounded health/lease state, never plaintext keys. Not connected to live inference in this foundation PR; production wiring must pass canonical tenant, workspace and provider scope before enabling.", + "verification": "crates/openhuman-core/src/inference/provider/factory/key_pool_tests.rs: pool_retains_only_salted_key_ids_not_raw_secrets; deleted_provider_scopes_are_pruned; least_busy_reserves_active_lease_before_returning_the_key" + }, + { + "file": "crates/openhuman-core/src/inference/provider/factory/key_pool.rs", + "scope": "", + "name": "SALT", + "type": "OnceLock", + "count": 1, + "category": "A", + "reason": "One per-process randomized hash builder used solely to derive ephemeral credential identities for the in-memory pool, preventing use of stable or plaintext secret identifiers. No tenant payload, keys, or requests are stored.", + "verification": "crates/openhuman-core/src/inference/provider/factory/key_pool_tests.rs: pool_retains_only_salted_key_ids_not_raw_secrets" + }, { "file": "crates/openhuman-core/src/inference/provider/factory/routing.rs", "scope": "", @@ -469,6 +509,36 @@ "reason": "Temporary prompt-clock callbacks are process test/host wiring and are invoked only to obtain the current instant; runtime-specific registrations use CoreContext extensions.", "verification": "prompts/section_renderers_tests.rs: newest_prompt_clock_wins_and_non_lifo_drop_restores_the_remaining_provider" }, + { + "file": "crates/openhuman-core/src/storage/mod.rs", + "scope": "", + "name": "RUNTIME", + "type": "OnceLock>", + "count": 1, + "category": "A", + "reason": "The dedicated two-thread tokio runtime that storage drivers connect on, so their background tasks never share threads with callers parked on the blocking bridge. It is built once per process, holds no tenant, profile or per-agent state, and carries no data beyond the drivers' connection tasks.", + "verification": "storage/mod_tests.rs: bridged_calls_do_not_starve_a_capped_host_runtime; storage/mod_tests.rs: open_connects_on_the_storage_runtime" + }, + { + "file": "crates/openhuman-core/src/storage/mod.rs", + "scope": "", + "name": "SLOT_OWNER", + "type": "Mutex>", + "count": 1, + "category": "A", + "reason": "Test-only (cfg(test)): names the libtest thread that installed the storage slot, so a slot one test installs is hidden from tests on other threads in the same process (MAIN-57). It holds only a thread name, no records or tenant state, and is not compiled into any shipped build.", + "verification": "storage/mod_tests.rs: a_test_installed_slot_is_hidden_from_other_test_threads" + }, + { + "file": "crates/openhuman-core/src/storage/sessions/host_stores.rs", + "scope": "", + "name": "FALLBACKS", + "type": "LazyLock>>", + "count": 1, + "category": "D", + "reason": "Each entry routes one workspace path to its storage backend and DriverSessionStores provider; returned stores still require the caller's tenant session key, while authoritative session data remains in the backend.", + "verification": "storage/sessions/mod_tests.rs: each_workspace_keeps_its_sessions_in_its_own_database; storage/sessions/mod_tests.rs: two_profiles_default_agents_keep_separate_session_keys_and_stores; threads/ops/edit_tests.rs: a_profile_prefixed_key_round_trips_through_edit_and_regenerate" + }, { "file": "crates/openhuman-core/src/storage/sessions/mod.rs", "scope": "", @@ -489,16 +559,6 @@ "reason": "The process's registry of weak handles to the rebindable session-store providers a host installed, so a workspace switch (storage::rebind) can point each at the new backend. It holds no records and no tenant state, only Weak references; the stores' data lives in the storage backend. A test builds its own StoreRegistry, so rebinding one never moves another's stores.", "verification": "storage/sessions/rebound_tests.rs: a_registry_rebind_leaves_other_registries_alone" }, - { - "file": "crates/openhuman-core/src/storage/sessions/host_stores.rs", - "scope": "", - "name": "FALLBACKS", - "type": "LazyLock>>", - "count": 1, - "category": "D", - "reason": "Each entry routes one workspace path to its storage backend and DriverSessionStores provider; returned stores still require the caller's tenant session key, while authoritative session data remains in the backend.", - "verification": "storage/sessions/mod_tests.rs: each_workspace_keeps_its_sessions_in_its_own_database; storage/sessions/mod_tests.rs: two_profiles_default_agents_keep_separate_session_keys_and_stores; threads/ops/edit_tests.rs: a_profile_prefixed_key_round_trips_through_edit_and_regenerate" - }, { "file": "crates/openhuman-core/src/threads/artifact_rows.rs", "scope": "", @@ -579,56 +639,6 @@ "reason": "Host wiring: an id-keyed table of factories the embedding host registers at startup (core::slots::HookSlot). It holds code, never tenant data; a per-context override table lives on the CoreContext, not here.", "verification": "voice/live/registry_tests.rs: a_registered_provider_is_listed_after_the_built_in_ones" }, - { - "file": "crates/openhuman-core/src/core/observability_reporting.rs", - "scope": "", - "name": "LAST_FS_LIMIT_REPORT", - "type": "std::sync::LazyLock>, >", - "count": 1, - "category": "A", - "reason": "A process-wide rate limiter keyed by the content-free domain and operation names prevents repeated filesystem-limit reports from flooding observability. Entries older than the five-minute cooldown are pruned on every access, and the map stores no tenant or user data.", - "verification": "core/observability_reporting.rs: was_recently_reported" - }, - { - "file": "crates/openhuman-core/src/storage/mod.rs", - "scope": "", - "name": "RUNTIME", - "type": "OnceLock>", - "count": 1, - "category": "A", - "reason": "The dedicated two-thread tokio runtime that storage drivers connect on, so their background tasks never share threads with callers parked on the blocking bridge. It is built once per process, holds no tenant, profile or per-agent state, and carries no data beyond the drivers' connection tasks.", - "verification": "storage/mod_tests.rs: bridged_calls_do_not_starve_a_capped_host_runtime; storage/mod_tests.rs: open_connects_on_the_storage_runtime" - }, - { - "file": "crates/openhuman-core/src/storage/mod.rs", - "scope": "", - "name": "SLOT_OWNER", - "type": "Mutex>", - "count": 1, - "category": "A", - "reason": "Test-only (cfg(test)): names the libtest thread that installed the storage slot, so a slot one test installs is hidden from tests on other threads in the same process (MAIN-57). It holds only a thread name, no records or tenant state, and is not compiled into any shipped build.", - "verification": "storage/mod_tests.rs: a_test_installed_slot_is_hidden_from_other_test_threads" - }, - { - "file": "crates/openhuman-core/src/web_chat/ops/test_hooks.rs", - "scope": "", - "name": "TEST_PARK_AFTER_RUN", - "type": "std::sync::atomic::AtomicBool", - "count": 1, - "category": "A", - "reason": "A debug/test-only hook flag (`cfg(any(test, debug_assertions))`) that parks `run_chat_task` after its run so the idle guard can be exercised end to end; it holds no tenant data and is off unless a test sets it.", - "verification": "tests/agent_harness_e2e.rs: idle_guard_local::stalled_turn_is_stopped_by_the_idle_guard_and_closed_with_a_summary" - }, - { - "file": "crates/openhuman-core/src/agent/session_lease.rs", - "scope": "", - "name": "LEASES", - "type": "LazyLock>>>>", - "count": 1, - "category": "D", - "reason": "Process-local lease table keyed by the complete tenant, agent and canonical session identity; it holds only weak mutex handles that coordinate runners inside this process and stores no tenant data.", - "verification": "agent/session_lease_tests.rs: distinct_tenant_agent_or_session_keys_do_not_serialize" - }, { "file": "crates/openhuman-core/src/web_chat/event_bus_channel.rs", "scope": "", @@ -649,6 +659,16 @@ "reason": "Test-only #[cfg(test)] slot recording the last run_chat_task error string; it is not compiled into production builds and carries no tenant data.", "verification": "channels/bus/subscriber_tests.rs: reads the slot through last_run_chat_task_error_for_test" }, + { + "file": "crates/openhuman-core/src/web_chat/ops/test_hooks.rs", + "scope": "", + "name": "TEST_PARK_AFTER_RUN", + "type": "std::sync::atomic::AtomicBool", + "count": 1, + "category": "A", + "reason": "A debug/test-only hook flag (`cfg(any(test, debug_assertions))`) that parks `run_chat_task` after its run so the idle guard can be exercised end to end; it holds no tenant data and is off unless a test sets it.", + "verification": "tests/agent_harness_e2e.rs: idle_guard_local::stalled_turn_is_stopped_by_the_idle_guard_and_closed_with_a_summary" + }, { "file": "crates/openhuman-core/src/web_chat/proactive.rs", "scope": "",