diff --git a/Cargo.lock b/Cargo.lock index 36d139f0..7c5124e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1763,6 +1763,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "sha2", "tempfile", "thiserror", "tinyinference-core", diff --git a/crates/tinyinference-embeddings/src/cloud.rs b/crates/tinyinference-embeddings/src/cloud.rs index 87f75aef..5d6c1a63 100644 --- a/crates/tinyinference-embeddings/src/cloud.rs +++ b/crates/tinyinference-embeddings/src/cloud.rs @@ -1,5 +1,7 @@ //! Bearer-authenticated OpenAI-compatible cloud embedding model. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use async_trait::async_trait; @@ -13,7 +15,8 @@ pub const DEFAULT_CLOUD_MODEL: &str = "embedding-v1"; pub const DEFAULT_CLOUD_DIMENSIONS: usize = 1024; /// Resolves the current bearer token for each request. -pub type BearerResolver = Arc Result + Send + Sync>; +pub type BearerResolver = + Arc Pin> + Send>> + Send + Sync>; /// Host policy hook invoked immediately before an embedding batch may leave /// the process. This keeps privacy and audit policy in the host without @@ -81,7 +84,7 @@ impl CloudEmbeddingModel { /// /// [`Error::Validation`] when an input is empty or whitespace, when the /// bearer resolver fails, or when it resolves to a blank token. - fn delegate(&self, texts: &[String]) -> Result> { + async fn delegate(&self, texts: &[String]) -> Result> { if texts.is_empty() { return Ok(None); } @@ -95,7 +98,7 @@ impl CloudEmbeddingModel { if let Some(guard) = &self.egress_guard { guard(&self.model, texts.len())?; } - let bearer = (self.bearer)()?; + let bearer = (self.bearer)().await?; if bearer.trim().is_empty() { return Err(Error::Validation( "No backend session for cloud embeddings".into(), @@ -128,7 +131,7 @@ impl EmbeddingModel for CloudEmbeddingModel { } async fn embed(&self, texts: &[String]) -> Result>> { - let Some(delegate) = self.delegate(texts)? else { + let Some(delegate) = self.delegate(texts).await? else { return Ok(Vec::new()); }; delegate.embed(texts).await @@ -140,7 +143,7 @@ impl EmbeddingModel for CloudEmbeddingModel { &self, texts: &[String], ) -> Result<(Vec>, Option)> { - let Some(delegate) = self.delegate(texts)? else { + let Some(delegate) = self.delegate(texts).await? else { return Ok((Vec::new(), None)); }; delegate.embed_with_usage(texts).await diff --git a/crates/tinyinference-embeddings/src/cloud_tests.rs b/crates/tinyinference-embeddings/src/cloud_tests.rs index 3ff9dab2..e935e500 100644 --- a/crates/tinyinference-embeddings/src/cloud_tests.rs +++ b/crates/tinyinference-embeddings/src/cloud_tests.rs @@ -2,9 +2,11 @@ use super::*; fn missing_bearer() -> BearerResolver { Arc::new(|| { - Err(Error::Validation( - "No backend session for cloud embeddings".into(), - )) + Box::pin(async { + Err(Error::Validation( + "No backend session for cloud embeddings".into(), + )) + }) }) } @@ -49,7 +51,7 @@ async fn empty_and_blank_batches_skip_egress_guard_and_bearer_for_both_entrypoin 8, Arc::new(move || { bearer_count.fetch_add(1, Ordering::SeqCst); - Ok("token".into()) + Box::pin(async { Ok("token".into()) }) }), ) .with_egress_guard(Arc::new(move |_, _| { diff --git a/crates/tinyinference-embeddings/src/factory.rs b/crates/tinyinference-embeddings/src/factory.rs index 8c63793a..0bcf47c8 100644 --- a/crates/tinyinference-embeddings/src/factory.rs +++ b/crates/tinyinference-embeddings/src/factory.rs @@ -1,5 +1,7 @@ //! Provider construction and custom-endpoint validation. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use crate::{ @@ -12,8 +14,11 @@ use crate::{ pub type CredentialResolver = Arc String + Send + Sync>; /// Constructs the host-authenticated managed embedding model. -pub type ManagedModelFactory = - Arc Result> + Send + Sync>; +pub type ManagedModelFactory = Arc< + dyn Fn(&str, usize) -> Pin>> + Send>> + + Send + + Sync, +>; /// Host-independent inputs for selecting an embedding provider. #[derive(Clone, Debug)] @@ -30,7 +35,7 @@ pub struct EmbeddingFactorySettings { /// Builds a provider from persisted settings while keeping credential storage /// and managed authentication in host-supplied callbacks. -pub fn create_configured_embedding_model( +pub async fn create_configured_embedding_model( settings: &EmbeddingFactorySettings, credentials: &CredentialResolver, managed: &ManagedModelFactory, @@ -42,7 +47,7 @@ pub fn create_configured_embedding_model( None => (provider, None), }; if matches!(slug, "cloud" | "managed") { - return managed(&settings.model, settings.dimensions); + return managed(&settings.model, settings.dimensions).await; } let api_key = credentials(slug); create_embedding_model( @@ -57,14 +62,14 @@ pub fn create_configured_embedding_model( /// Builds the configured provider, falling back to the managed model when the /// selection is blank, invalid, or missing a required credential. -pub fn create_default_embedding_model( +pub async fn create_default_embedding_model( settings: &EmbeddingFactorySettings, credentials: &CredentialResolver, managed: &ManagedModelFactory, ) -> Result> { let provider = settings.provider.trim(); if provider.is_empty() || matches!(provider, "cloud" | "managed") { - return managed(DEFAULT_CLOUD_MODEL, DEFAULT_CLOUD_DIMENSIONS); + return managed(DEFAULT_CLOUD_MODEL, DEFAULT_CLOUD_DIMENSIONS).await; } let slug = if provider.starts_with("custom:") { @@ -79,14 +84,14 @@ pub fn create_default_embedding_model( provider = slug, "configured embedding provider has no credential; falling back to managed" ); - return managed(DEFAULT_CLOUD_MODEL, DEFAULT_CLOUD_DIMENSIONS); + return managed(DEFAULT_CLOUD_MODEL, DEFAULT_CLOUD_DIMENSIONS).await; } - match create_configured_embedding_model(settings, credentials, managed) { + match create_configured_embedding_model(settings, credentials, managed).await { Ok(model) => Ok(model), Err(error) => { tracing::warn!(provider = slug, %error, "configured embedding provider failed to build; falling back to managed"); - managed(DEFAULT_CLOUD_MODEL, DEFAULT_CLOUD_DIMENSIONS) + managed(DEFAULT_CLOUD_MODEL, DEFAULT_CLOUD_DIMENSIONS).await } } } diff --git a/crates/tinyinference-embeddings/src/verify.rs b/crates/tinyinference-embeddings/src/verify.rs index 5ac19a20..565df87a 100644 --- a/crates/tinyinference-embeddings/src/verify.rs +++ b/crates/tinyinference-embeddings/src/verify.rs @@ -200,6 +200,7 @@ pub async fn test_embedding_connection_detailed( }); } let model = create_configured_embedding_model(settings, credentials, managed) + .await .map_err(ConnectionTestError::Build)?; let vectors = model .embed(&["connection test".to_owned()]) diff --git a/crates/tinyinference-embeddings/src/verify_tests.rs b/crates/tinyinference-embeddings/src/verify_tests.rs index 07633312..414c3641 100644 --- a/crates/tinyinference-embeddings/src/verify_tests.rs +++ b/crates/tinyinference-embeddings/src/verify_tests.rs @@ -122,8 +122,11 @@ async fn connection_test_uses_host_managed_factory_and_returns_actual_width() { ollama_base_url: "http://localhost:11434".into(), }; let credentials: CredentialResolver = std::sync::Arc::new(|_| String::new()); - let managed: ManagedModelFactory = - std::sync::Arc::new(|_, dims| Ok(Box::new(crate::MockEmbeddingModel::new(dims)))); + let managed: ManagedModelFactory = std::sync::Arc::new(|_, dims| { + Box::pin(async move { + Ok(Box::new(crate::MockEmbeddingModel::new(dims)) as Box) + }) + }); let result = test_embedding_connection(&settings, &credentials, &managed, None) .await .unwrap(); @@ -139,7 +142,8 @@ async fn detailed_connection_test_identifies_model_build_failures() { ollama_base_url: "http://localhost:11434".into(), }; let credentials: CredentialResolver = std::sync::Arc::new(|_| String::new()); - let managed: ManagedModelFactory = std::sync::Arc::new(|_, _| unreachable!()); + let managed: ManagedModelFactory = + std::sync::Arc::new(|_, _| Box::pin(async { unreachable!() })); assert!(matches!( test_embedding_connection_detailed(&settings, &credentials, &managed, None).await, @@ -177,7 +181,9 @@ async fn detailed_connection_test_identifies_embedding_failures() { ollama_base_url: "http://localhost:11434".into(), }; let credentials: CredentialResolver = std::sync::Arc::new(|_| String::new()); - let managed: ManagedModelFactory = std::sync::Arc::new(|_, _| Ok(Box::new(FailingModel))); + let managed: ManagedModelFactory = std::sync::Arc::new(|_, _| { + Box::pin(async { Ok(Box::new(FailingModel) as Box) }) + }); assert!(matches!( test_embedding_connection_detailed(&settings, &credentials, &managed, None).await, @@ -194,7 +200,8 @@ async fn connection_test_preserves_noop_provider_success() { ollama_base_url: String::new(), }; let credentials: CredentialResolver = std::sync::Arc::new(|_| String::new()); - let managed: ManagedModelFactory = std::sync::Arc::new(|_, _| unreachable!()); + let managed: ManagedModelFactory = + std::sync::Arc::new(|_, _| Box::pin(async { unreachable!() })); let result = test_embedding_connection(&settings, &credentials, &managed, None) .await .unwrap(); @@ -331,7 +338,7 @@ async fn custom_connection_infers_stored_endpoint_and_resolves_custom_credential capture.lock().unwrap().push(slug.to_owned()); "key".into() }); - let managed: ManagedModelFactory = Arc::new(|_, _| unreachable!()); + let managed: ManagedModelFactory = Arc::new(|_, _| Box::pin(async { unreachable!() })); let result = test_embedding_connection(&settings, &credentials, &managed, None) .await .unwrap(); diff --git a/crates/tinyinference-image/src/openrouter_tests.rs b/crates/tinyinference-image/src/openrouter_tests.rs index c2d4a13d..d4f547ba 100644 --- a/crates/tinyinference-image/src/openrouter_tests.rs +++ b/crates/tinyinference-image/src/openrouter_tests.rs @@ -257,15 +257,26 @@ fn enveloped_png_reply(_call: usize) -> Response { #[tokio::test] async fn proxied_backend_base_url_and_bearer_resolver() { let fixture = fixture("/agent-integrations/openrouter", enveloped_png_reply, None).await; - let resolver: crate::BearerResolver = Arc::new(|| Ok("session-jwt-abcdefgh".to_owned())); + let resolver_calls = Arc::new(AtomicUsize::new(0)); + let resolver: crate::BearerResolver = { + let resolver_calls = Arc::clone(&resolver_calls); + Arc::new(move || { + let call = resolver_calls.fetch_add(1, Ordering::SeqCst); + Box::pin(async move { Ok(format!("session-jwt-{call}")) }) + }) + }; let generator = OpenRouterImageGenerator::with_transport( MediaTransport::new(MediaAuth::Bearer(resolver)).with_base_url(&fixture.base_url), - ); + ) + .with_capability_check(false); let response = generator.generate(ImageRequest::new("x")).await.unwrap(); assert_eq!(response.images.len(), 1); assert_eq!(response.cost_usd, Some(0.035)); - let headers = fixture.captured.headers.lock().unwrap()[0].clone(); - assert_eq!(headers["authorization"], "Bearer session-jwt-abcdefgh"); + generator.generate(ImageRequest::new("y")).await.unwrap(); + assert_eq!(resolver_calls.load(Ordering::SeqCst), 2); + let headers = fixture.captured.headers.lock().unwrap(); + assert_eq!(headers[0]["authorization"], "Bearer session-jwt-0"); + assert_eq!(headers[1]["authorization"], "Bearer session-jwt-1"); } #[tokio::test] @@ -288,7 +299,7 @@ async fn failed_envelope_is_an_error_even_with_a_2xx_status() { #[tokio::test] async fn blank_bearer_fails_without_a_request() { let fixture = fixture("/api/v1", png_reply, None).await; - let resolver: crate::BearerResolver = Arc::new(|| Ok(" ".to_owned())); + let resolver: crate::BearerResolver = Arc::new(|| Box::pin(async { Ok(" ".to_owned()) })); let generator = OpenRouterImageGenerator::with_transport( MediaTransport::new(MediaAuth::Bearer(resolver)).with_base_url(&fixture.base_url), ) diff --git a/crates/tinyinference-image/src/transport.rs b/crates/tinyinference-image/src/transport.rs index d86ef6e2..d572e434 100644 --- a/crates/tinyinference-image/src/transport.rs +++ b/crates/tinyinference-image/src/transport.rs @@ -18,6 +18,8 @@ //! connection after a submit is surfaced as-is, because resubmitting could //! bill a second generation. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -40,7 +42,8 @@ pub const DEFAULT_MAX_MEDIA_BYTES: usize = 512 * 1024 * 1024; const MAX_ERROR_BODY_BYTES: usize = 16 * 1024; /// Resolves the current bearer token for each request. -pub type BearerResolver = Arc Result + Send + Sync>; +pub type BearerResolver = + Arc Pin> + Send>> + Send + Sync>; /// How requests authenticate. #[derive(Clone)] @@ -64,10 +67,10 @@ impl MediaAuth { } } - fn token(&self) -> Result { + async fn token(&self) -> Result { let token = match self { Self::ApiKey(key) => key.clone(), - Self::Bearer(resolve) => resolve()?, + Self::Bearer(resolve) => resolve().await?, }; let token = token.trim().to_owned(); if token.is_empty() { @@ -197,7 +200,7 @@ impl MediaTransport { T: DeserializeOwned, { let body = serde_json::to_vec(body)?; - let response = self + let (response, _) = self .send(reqwest::Method::POST, path, Some(body), Billing::Billable) .await?; decode_json_with_limit(response, self.json_limit()).await @@ -210,7 +213,7 @@ impl MediaTransport { /// [`Error::Auth`], [`Error::Http`], [`Error::Transport`] or /// [`Error::Decode`] after retries are exhausted. pub async fn get_json(&self, path: &str) -> Result { - let response = self + let (response, _) = self .send(reqwest::Method::GET, path, None, Billing::Idempotent) .await?; decode_json_with_limit(response, self.json_limit()).await @@ -264,8 +267,7 @@ impl MediaTransport { } async fn get_bytes_once(&self, path: &str) -> Result<(Bytes, Option)> { - let token = self.auth.token()?; - let mut response = self + let (mut response, token) = self .send(reqwest::Method::GET, path, None, Billing::Idempotent) .await?; let content_type = response @@ -317,11 +319,11 @@ impl MediaTransport { path: &str, body: Option>, billing: Billing, - ) -> Result { + ) -> Result<(reqwest::Response, String)> { let url = self.url(path); let mut attempt = 0u32; loop { - let token = self.auth.token()?; + let token = self.auth.token().await?; let mut request = self .client .request(method.clone(), &url) @@ -340,7 +342,7 @@ impl MediaTransport { ); let outcome = request.send().await; let (retry, retry_after, error) = match outcome { - Ok(response) if response.status().is_success() => return Ok(response), + Ok(response) if response.status().is_success() => return Ok((response, token)), Ok(response) => { let status = response.status().as_u16(); let retry_after = response diff --git a/crates/tinyinference-llm/Cargo.toml b/crates/tinyinference-llm/Cargo.toml index 1663d490..f1e2acb7 100644 --- a/crates/tinyinference-llm/Cargo.toml +++ b/crates/tinyinference-llm/Cargo.toml @@ -28,6 +28,7 @@ futures = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +sha2.workspace = true thiserror = { workspace = true } tinyinference-providers = { version = "0.5.0", path = "../tinyinference-providers" } tinyinference-core = { version = "0.5.0", path = "../tinyinference-core" } diff --git a/crates/tinyinference-llm/src/providers/overrides.rs b/crates/tinyinference-llm/src/providers/overrides.rs index da169135..59f4e9a2 100644 --- a/crates/tinyinference-llm/src/providers/overrides.rs +++ b/crates/tinyinference-llm/src/providers/overrides.rs @@ -177,16 +177,16 @@ impl Clone for ProviderModelMetadata { } /// Inputs a host embedding factory uses to construct an embedding model. -pub struct EmbeddingRequest<'a, Ctx> { +pub struct EmbeddingRequest { /// Embedding model id selected by the host. - pub model: &'a str, + pub model: String, /// Vector dimensionality selected by the host. pub dimensions: usize, - /// Borrowed host context such as configuration and credential scope. - pub context: &'a Ctx, + /// Owned host context such as configuration and credential scope. + pub context: Ctx, } -impl std::fmt::Debug for EmbeddingRequest<'_, Ctx> { +impl std::fmt::Debug for EmbeddingRequest { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter .debug_struct("EmbeddingRequest") @@ -197,12 +197,12 @@ impl std::fmt::Debug for EmbeddingRequest<'_, Ctx> { } } +/// A provider embedding factory's async result. +pub type EmbeddingFactoryFuture = BoxFuture>>; + /// Builds a provider embedding model. -pub type EmbeddingFactory = Arc< - dyn for<'a> Fn(&EmbeddingRequest<'a, Ctx>) -> EmbeddingResult> - + Send - + Sync, ->; +pub type EmbeddingFactory = + Arc) -> EmbeddingFactoryFuture + Send + Sync>; /// Owned host context passed to a catalog listing implementation. pub struct CatalogRequest { @@ -223,8 +223,21 @@ impl std::fmt::Debug for CatalogRequest { pub type CatalogListing = Arc) -> BoxFuture> + Send + Sync>; +/// Async result of a provider limit-discovery request. +pub type CatalogLimitsFuture = BoxFuture>; + +/// Input to an async provider catalog limits resolver. +#[derive(Debug)] +pub struct CatalogLimitsRequest { + /// Model whose supported limits are being discovered. + pub model: String, + /// Owned host context such as configuration and credential sources. + pub context: Ctx, +} + /// Builds the provider limit discovery request for a model, if supported. -pub type CatalogLimits = Arc Option + Send + Sync>; +pub type CatalogLimits = + Arc) -> CatalogLimitsFuture + Send + Sync>; /// A provider catalog listing and optional model-limit resolver. pub struct CatalogEntry { @@ -338,18 +351,18 @@ impl ProviderOverrides { /// /// # Errors /// Returns the selected factory's construction error. - pub fn build_embedding( + pub async fn build_embedding( &self, provider_id: &str, - request: &EmbeddingRequest<'_, Ctx>, + request: EmbeddingRequest, builtin: B, ) -> EmbeddingResult> where - B: FnOnce(&EmbeddingRequest<'_, Ctx>) -> EmbeddingResult>, + B: FnOnce(EmbeddingRequest) -> EmbeddingFactoryFuture, { match self.embedding_factory(provider_id) { - Some(factory) => factory(request), - None => builtin(request), + Some(factory) => factory(request).await, + None => builtin(request).await, } } diff --git a/crates/tinyinference-llm/src/providers/tinyhumans/cache_identity_tests.rs b/crates/tinyinference-llm/src/providers/tinyhumans/cache_identity_tests.rs index d859da2b..6aa44cab 100644 --- a/crates/tinyinference-llm/src/providers/tinyhumans/cache_identity_tests.rs +++ b/crates/tinyinference-llm/src/providers/tinyhumans/cache_identity_tests.rs @@ -35,17 +35,17 @@ fn cache_namespace_is_trimmed_and_blank_input_keeps_the_current_value() { #[test] fn cache_identity_keeps_endpoint_and_default_model_dimensions() { let first = TinyHumansModel::new("https://one.example.test", "", "model-a", || { - Ok("credential-a".into()) + Box::pin(async { Ok("credential-a".into()) }) }) .unwrap() .with_cache_namespace("openhuman"); let other_endpoint = TinyHumansModel::new("https://two.example.test", "", "model-a", || { - Ok("credential-b".into()) + Box::pin(async { Ok("credential-b".into()) }) }) .unwrap() .with_cache_namespace("openhuman"); let other_model = TinyHumansModel::new("https://one.example.test", "", "model-b", || { - Ok("credential-c".into()) + Box::pin(async { Ok("credential-c".into()) }) }) .unwrap() .with_cache_namespace("openhuman"); @@ -63,8 +63,11 @@ fn reading_cache_identity_never_resolves_or_includes_credentials() { let reads = Arc::new(AtomicUsize::new(0)); let resolver_reads = Arc::clone(&reads); let model = TinyHumansModel::new("https://api.example.test", "", "default-model", move || { - resolver_reads.fetch_add(1, Ordering::SeqCst); - Ok("must-not-appear-in-cache-identity".into()) + let resolver_reads = Arc::clone(&resolver_reads); + Box::pin(async move { + resolver_reads.fetch_add(1, Ordering::SeqCst); + Ok("must-not-appear-in-cache-identity".into()) + }) }) .unwrap() .with_cache_namespace("openhuman"); diff --git a/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs b/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs index cf1f068b..f051d117 100644 --- a/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs +++ b/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs @@ -9,10 +9,13 @@ mod catalog; mod reasoning; mod usage; +use std::future::Future; +use std::pin::Pin; use std::sync::{Arc, OnceLock}; use std::time::Duration; use async_trait::async_trait; +use sha2::{Digest, Sha256}; use tinyinference_core::sanitize; use crate::model::{ @@ -28,7 +31,8 @@ pub use reasoning::{REASONING_OFF_METADATA_KEY, apply_reasoning, with_thread_id} pub use usage::project_tinyhumans_envelope; /// Resolves the current bearer credential for a managed inference call. -pub type BearerResolver = Arc Result + Send + Sync>; +pub type BearerResolver = + Arc Pin> + Send>> + Send + Sync>; /// Which kind of credential a managed request authenticated with. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -48,14 +52,96 @@ pub struct TinyHumansCredential { pub kind: TinyHumansCredentialKind, } -type KindedResolver = - Arc Result<(String, Option)> + Send + Sync>; +type KindedResolver = Arc< + dyn Fn() -> Pin< + Box)>> + Send>, + > + Send + + Sync, +>; + +/// Opaque, non-secret identity for a bearer credential used by one request. +/// +/// The fingerprint is a domain-separated SHA-256 digest. It supports comparing +/// the credential used by a failed request with a host's current bearer without +/// retaining or exposing the bearer itself. +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct TinyHumansCredentialFingerprint([u8; 32]); + +impl TinyHumansCredentialFingerprint { + fn from_bearer(bearer: &str) -> Self { + let mut digest = Sha256::new(); + digest.update(b"tinyinference/tinyhumans-credential-fingerprint/v1\0"); + digest.update(bearer.as_bytes()); + Self(digest.finalize().into()) + } + + /// Returns the canonical lower-case hexadecimal representation of this + /// fingerprint. The bearer is not recoverable from this representation. + #[must_use] + pub fn to_hex(self) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut encoded = String::with_capacity(self.0.len() * 2); + for byte in self.0 { + encoded.push(char::from(HEX[usize::from(byte >> 4)])); + encoded.push(char::from(HEX[usize::from(byte & 0x0f)])); + } + encoded + } + + /// Parses a canonical lower-case hexadecimal fingerprint representation. + /// + /// Returns `None` when `encoded` is not exactly 64 lower-case hex digits. + #[must_use] + pub fn from_hex(encoded: &str) -> Option { + if encoded.len() != 64 { + return None; + } + let mut bytes = [0_u8; 32]; + for (index, pair) in encoded.as_bytes().as_chunks::<2>().0.iter().enumerate() { + let high = decode_lower_hex(pair[0])?; + let low = decode_lower_hex(pair[1])?; + bytes[index] = (high << 4) | low; + } + Some(Self(bytes)) + } + + /// Returns whether this fingerprint belongs to `bearer`. + #[must_use] + pub fn matches_bearer(self, bearer: &str) -> bool { + let candidate = Self::from_bearer(bearer); + // Compare the fixed-width digest without early exit. + self.0 + .iter() + .zip(candidate.0.iter()) + .fold(0_u8, |different, (left, right)| different | (left ^ right)) + == 0 + } +} + +fn decode_lower_hex(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + _ => None, + } +} + +impl std::fmt::Debug for TinyHumansCredentialFingerprint { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("TinyHumansCredentialFingerprint()") + } +} + +#[derive(Clone, Copy)] +struct SentCredential { + kind: Option, + fingerprint: TinyHumansCredentialFingerprint, +} tokio::task_local! { - /// The credential kind the bearer source handed to the request being - /// dispatched in this task, so a failure is attributed to the credential - /// that request carried rather than to whatever was resolved last. - static SENT_CREDENTIAL: std::cell::Cell>; + /// Request-time credential identity so a failure is attributed to the + /// credential actually sent, rather than whichever one resolves later. + static SENT_CREDENTIAL: std::cell::Cell>; } /// Operation that received a provider authentication failure. @@ -102,6 +188,7 @@ pub struct TinyHumansAuthFailure { operation: TinyHumansOperation, kind: TinyHumansAuthFailureKind, credential: Option, + credential_fingerprint: Option, } impl TinyHumansAuthFailure { @@ -124,6 +211,13 @@ impl TinyHumansAuthFailure { pub fn credential(self) -> Option { self.credential } + + /// The opaque identity of the bearer sent with the failed request, when + /// the request used a managed credential resolver. + #[must_use] + pub fn credential_fingerprint(self) -> Option { + self.credential_fingerprint + } } /// Handles authentication failures from a managed endpoint. @@ -187,13 +281,19 @@ impl TinyHumansModel { base_url: impl AsRef, model: impl AsRef, default_model: impl AsRef, - resolver: impl Fn() -> Result + Send + Sync + 'static, + resolver: impl Fn() -> Pin> + Send>> + + Send + + Sync + + 'static, ) -> Result { Self::build( base_url, model, default_model, - Arc::new(move || resolver().map(|bearer| (bearer, None))), + Arc::new(move || { + let future = resolver(); + Box::pin(async move { future.await.map(|bearer| (bearer, None)) }) + }), ) } @@ -207,14 +307,22 @@ impl TinyHumansModel { base_url: impl AsRef, model: impl AsRef, default_model: impl AsRef, - resolver: impl Fn() -> Result + Send + Sync + 'static, + resolver: impl Fn() -> Pin> + Send>> + + Send + + Sync + + 'static, ) -> Result { Self::build( base_url, model, default_model, Arc::new(move || { - resolver().map(|credential| (credential.bearer, Some(credential.kind))) + let future = resolver(); + Box::pin(async move { + future + .await + .map(|credential| (credential.bearer, Some(credential.kind))) + }) }), ) } @@ -409,7 +517,7 @@ impl TinyHumansModel { &self, error: &Error, operation: TinyHumansOperation, - credential: Option, + credential: Option, ) { if let Error::Provider(provider_error) = error && is_auth_failure(provider_error) @@ -428,7 +536,7 @@ impl TinyHumansModel { &self, error: &ProviderError, operation: TinyHumansOperation, - credential: Option, + credential: Option, ) { if let Some(handler) = &self.auth_failure_handler { handler(error, operation); @@ -439,7 +547,8 @@ impl TinyHumansModel { TinyHumansAuthFailure { operation, kind: auth_failure_kind(error), - credential, + credential: credential.and_then(|sent| sent.kind), + credential_fingerprint: credential.map(|sent| sent.fingerprint), }, ); } @@ -515,7 +624,8 @@ impl ChatModel<()> for TinyHumansModel { TinyHumansAuthFailure { operation: TinyHumansOperation::Stream, kind: auth_failure_kind(&error), - credential, + credential: credential.and_then(|sent| sent.kind), + credential_fingerprint: credential.map(|sent| sent.fingerprint), }, ); } @@ -579,10 +689,14 @@ struct ResolvedBearer(KindedResolver); #[async_trait] impl BearerSource for ResolvedBearer { async fn current(&self) -> Result> { - let (bearer, kind) = (self.0)()?; + let (bearer, kind) = (self.0)().await?; // Outside a dispatch scope (e.g. the readiness probe) there is nothing // to attribute, so a missing task-local is fine. - let _ = SENT_CREDENTIAL.try_with(|sent| sent.set(kind)); + let sent = SentCredential { + kind, + fingerprint: TinyHumansCredentialFingerprint::from_bearer(&bearer), + }; + let _ = SENT_CREDENTIAL.try_with(|credential| credential.set(Some(sent))); Ok(Some(bearer)) } } diff --git a/crates/tinyinference-llm/src/providers/tinyhumans/mod_tests.rs b/crates/tinyinference-llm/src/providers/tinyhumans/mod_tests.rs index 8e3f9190..c3599440 100644 --- a/crates/tinyinference-llm/src/providers/tinyhumans/mod_tests.rs +++ b/crates/tinyinference-llm/src/providers/tinyhumans/mod_tests.rs @@ -14,6 +14,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; struct Fixture { requests: Arc>>, response: Arc>, + rotate_token_before_response: Arc>>>>, } #[derive(Clone)] @@ -49,6 +50,9 @@ fn completion_body() -> String { async fn chat(State(fixture): State, headers: HeaderMap, body: String) -> Response { let body: Value = serde_json::from_str(&body).expect("request JSON"); fixture.requests.lock().unwrap().push((headers, body)); + if let Some(token) = fixture.rotate_token_before_response.lock().unwrap().take() { + *token.lock().unwrap() = "new-session-token".to_owned(); + } let response = fixture.response.lock().unwrap().clone(); Response::builder() .status(response.status) @@ -89,8 +93,10 @@ async fn serve(fixture: Fixture) -> String { } fn fixture_model(base: &str) -> TinyHumansModel { - TinyHumansModel::new(base, "", "default-model", || Ok("test-token".into())) - .expect("model construction") + TinyHumansModel::new(base, "", "default-model", || { + Box::pin(async { Ok("test-token".into()) }) + }) + .expect("model construction") } #[tokio::test] @@ -100,10 +106,13 @@ async fn dynamic_bearer_is_resolved_again_for_invoke_and_stream() { let reads = Arc::new(AtomicUsize::new(0)); let source = reads.clone(); let model = TinyHumansModel::new(&base, "", "default-model", move || { - Ok(format!( - "token-{}", - source.fetch_add(1, Ordering::SeqCst) + 1 - )) + let source = Arc::clone(&source); + Box::pin(async move { + Ok(format!( + "token-{}", + source.fetch_add(1, Ordering::SeqCst) + 1 + )) + }) }) .unwrap(); model @@ -133,7 +142,7 @@ async fn constructor_model_is_default_and_nonblank_request_model_overrides_it() let fixture = Fixture::default(); let base = serve(fixture.clone()).await; let model = TinyHumansModel::new(&base, " constructor-model ", "fallback-model", || { - Ok("token".into()) + Box::pin(async { Ok("token".into()) }) }) .unwrap(); for requested_model in [None, Some("request-model"), Some(" ")] { @@ -148,8 +157,10 @@ async fn constructor_model_is_default_and_nonblank_request_model_overrides_it() assert_eq!(requests[2].1["model"], "constructor-model"); } - let fallback_model = - TinyHumansModel::new(&base, " ", "fallback-model", || Ok("token".into())).unwrap(); + let fallback_model = TinyHumansModel::new(&base, " ", "fallback-model", || { + Box::pin(async { Ok("token".into()) }) + }) + .unwrap(); fallback_model .invoke( &(), @@ -168,7 +179,7 @@ async fn resolver_failure_stops_before_transport() { let fixture = Fixture::default(); let base = serve(fixture.clone()).await; let model = TinyHumansModel::new(&base, "m", "default-model", || { - Err(Error::Model("credential source failed".into())) + Box::pin(async { Err(Error::Model("credential source failed".into())) }) }) .unwrap(); let error = model @@ -473,7 +484,7 @@ async fn readiness_blocks_only_confirmed_missing_provider_configuration() { #[tokio::test] async fn readiness_fails_open_when_bearer_resolution_fails() { let model = TinyHumansModel::new("http://127.0.0.1:1", "m", "default-model", || { - Err(Error::Model("credential lookup unavailable".into())) + Box::pin(async { Err(Error::Model("credential lookup unavailable".into())) }) }) .unwrap(); model @@ -732,14 +743,17 @@ async fn managed_auth_failure_reports_the_credential_its_request_carried() { let seen = Arc::new(std::sync::Mutex::new(Vec::new())); let seen_by_handler = seen.clone(); let model = TinyHumansModel::new_with_credential(base, "", "default-model", move || { - let api_key = resolver_use_api_key.load(Ordering::SeqCst); - Ok(TinyHumansCredential { - bearer: if api_key { "key".into() } else { "jwt".into() }, - kind: if api_key { - TinyHumansCredentialKind::ApiKey - } else { - TinyHumansCredentialKind::Session - }, + let resolver_use_api_key = Arc::clone(&resolver_use_api_key); + Box::pin(async move { + let api_key = resolver_use_api_key.load(Ordering::SeqCst); + Ok(TinyHumansCredential { + bearer: if api_key { "key".into() } else { "jwt".into() }, + kind: if api_key { + TinyHumansCredentialKind::ApiKey + } else { + TinyHumansCredentialKind::Session + }, + }) }) }) .expect("model construction") @@ -775,6 +789,71 @@ async fn managed_auth_failure_reports_the_credential_its_request_carried() { ); } +#[tokio::test] +async fn managed_auth_failure_carries_a_non_secret_request_credential_fingerprint() { + let fixture = Fixture::default(); + let base = serve(fixture.clone()).await; + *fixture.response.lock().unwrap() = FixtureResponse { + status: StatusCode::UNAUTHORIZED, + content_type: "application/json", + body: json!({ "error": { "message": "Unauthorized" } }).to_string(), + }; + let token = Arc::new(std::sync::Mutex::new("old-request-token".to_owned())); + *fixture.rotate_token_before_response.lock().unwrap() = Some(token.clone()); + let token_for_resolver = token.clone(); + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen_by_handler = seen.clone(); + let model = TinyHumansModel::new_with_credential(base, "", "default-model", move || { + let token_for_resolver = Arc::clone(&token_for_resolver); + Box::pin(async move { + Ok(TinyHumansCredential { + bearer: token_for_resolver.lock().unwrap().clone(), + kind: TinyHumansCredentialKind::Session, + }) + }) + }) + .unwrap() + .with_managed_auth_failure_handler(move |_, failure| { + seen_by_handler + .lock() + .unwrap() + .push(failure.credential_fingerprint()); + }); + + let request = ModelRequest::new(vec![crate::message::Message::user("hi")]); + assert!(model.invoke(&(), request).await.is_err()); + assert_eq!(*token.lock().unwrap(), "new-session-token"); + + let fingerprints = seen.lock().unwrap(); + let fingerprint = fingerprints[0].expect("managed credential fingerprint"); + assert!(fingerprint.matches_bearer("old-request-token")); + assert!(!fingerprint.matches_bearer("new-session-token")); + assert!(!format!("{fingerprint:?}").contains("old-request-token")); +} + +#[test] +fn credential_fingerprint_has_a_stable_opaque_hex_round_trip() { + let fingerprint = TinyHumansCredentialFingerprint::from_bearer("session-token"); + let encoded = fingerprint.to_hex(); + + assert_eq!(encoded.len(), 64); + assert!( + encoded + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)) + ); + assert_eq!( + TinyHumansCredentialFingerprint::from_hex(&encoded), + Some(fingerprint) + ); + assert_eq!( + TinyHumansCredentialFingerprint::from_hex(&encoded.to_ascii_uppercase()), + None, + "the serialized form is canonical lower-case hex" + ); + assert_eq!(TinyHumansCredentialFingerprint::from_hex("xyz"), None); +} + #[tokio::test] async fn a_plain_bearer_resolver_reports_no_credential_kind() { let fixture = Fixture::default(); diff --git a/crates/tinyinference-llm/src/providers/tinyhumans/observer_tests.rs b/crates/tinyinference-llm/src/providers/tinyhumans/observer_tests.rs index 2a2cf3d6..abd7076f 100644 --- a/crates/tinyinference-llm/src/providers/tinyhumans/observer_tests.rs +++ b/crates/tinyinference-llm/src/providers/tinyhumans/observer_tests.rs @@ -215,8 +215,11 @@ async fn observer_classifies_readiness_and_non_provider_failures() { let errors = Arc::new(std::sync::Mutex::new(Vec::new())); let errors_seen = Arc::clone(&errors); let model = TinyHumansModel::new("http://127.0.0.1:9", "model", "default-model", move || { - calls.fetch_add(1, Ordering::SeqCst); - Err(Error::Model("resolver marker".into())) + let calls = Arc::clone(&calls); + Box::pin(async move { + calls.fetch_add(1, Ordering::SeqCst); + Err(Error::Model("resolver marker".into())) + }) }) .unwrap() .with_error_observer(move |error, point| { diff --git a/crates/tinyinference-llm/tests/provider_overrides.rs b/crates/tinyinference-llm/tests/provider_overrides.rs index 9767a467..ac160a9c 100644 --- a/crates/tinyinference-llm/tests/provider_overrides.rs +++ b/crates/tinyinference-llm/tests/provider_overrides.rs @@ -7,8 +7,9 @@ use tinyinference_embeddings::{EmbeddingModel, MockEmbeddingModel}; use tinyinference_llm::model::ChatModel; use tinyinference_llm::providers::MockModel; use tinyinference_llm::providers::overrides::{ - CatalogEntry, CatalogRequest, ChatModelEntry, ChatModelRequest, EmbeddingRequest, - ProviderModelMetadata, ProviderOverrides, ReadinessRequest, + CatalogEntry, CatalogLimitsFuture, CatalogLimitsRequest, CatalogRequest, ChatModelEntry, + ChatModelRequest, EmbeddingFactoryFuture, EmbeddingRequest, ProviderModelMetadata, + ProviderOverrides, ReadinessRequest, }; struct HostContext<'a> { @@ -17,11 +18,10 @@ struct HostContext<'a> { _ownership_marker: std::marker::PhantomData<&'a mut ()>, } -#[test] -fn synchronous_factories_accept_borrowed_non_clone_host_contexts() { - let config = String::from("borrowed-host-config"); +#[tokio::test] +async fn asynchronous_factories_accept_owned_non_clone_host_contexts() { let context = HostContext { - borrowed_config: &config, + borrowed_config: "borrowed-host-config", _ownership_marker: std::marker::PhantomData, }; let overrides = ProviderOverrides::new( @@ -40,11 +40,18 @@ fn synchronous_factories_accept_borrowed_non_clone_host_contexts() { }, |id| { (id == "hosted").then(|| { - Arc::new(|request: &EmbeddingRequest<'_, HostContext<'_>>| { - assert_eq!(request.context.borrowed_config, "borrowed-host-config"); - Ok(Box::new(MockEmbeddingModel::new(request.dimensions)) - as Box) - }) as _ + Arc::new( + |request: EmbeddingRequest>| -> EmbeddingFactoryFuture { + Box::pin(async move { + assert_eq!(request.context.borrowed_config, "borrowed-host-config"); + Ok(Box::new(MockEmbeddingModel::new(request.dimensions)) + as Box) + }) + }, + ) + as tinyinference_llm::providers::overrides::EmbeddingFactory< + HostContext<'static>, + > }) }, |_| None, @@ -64,15 +71,17 @@ fn synchronous_factories_accept_borrowed_non_clone_host_contexts() { ); assert_eq!(chat.unwrap().1, "borrowed-model"); - let embedding = overrides.build_embedding( - "hosted", - &EmbeddingRequest { - model: "embed", - dimensions: 3, - context: &context, - }, - |_| panic!("registered embedding factory should be used"), - ); + let embedding = overrides + .build_embedding( + "hosted", + EmbeddingRequest { + model: "embed".into(), + dimensions: 3, + context, + }, + |_| panic!("registered embedding factory should be used"), + ) + .await; assert_eq!(embedding.unwrap().dimensions(), 3); } @@ -118,11 +127,18 @@ async fn independent_live_family_lookups_override_only_their_builtin() { }, |id| { (id == "shared").then(|| { - Arc::new(|request: &EmbeddingRequest<'_, HostContext<'_>>| { - assert_eq!(request.context.borrowed_config, "scoped-config"); - Ok(Box::new(MockEmbeddingModel::new(request.dimensions)) - as Box) - }) as _ + Arc::new( + |request: EmbeddingRequest>| -> EmbeddingFactoryFuture { + Box::pin(async move { + assert_eq!(request.context.borrowed_config, "scoped-config"); + Ok(Box::new(MockEmbeddingModel::new(request.dimensions)) + as Box) + }) + }, + ) + as tinyinference_llm::providers::overrides::EmbeddingFactory< + HostContext<'static>, + > }) }, |id| { @@ -133,16 +149,20 @@ async fn independent_live_family_lookups_override_only_their_builtin() { Ok(serde_json::json!({"models": []})) }) }), - limits: Some(Arc::new(|model, context: &HostContext<'static>| { - (model == "catalog-model" && context.borrowed_config == "catalog-config").then( - || { - tinyinference_llm::model::discover::DiscoveryRequest::new( - "https://provider.invalid/v1", - model, - ) - }, - ) - })), + limits: Some(Arc::new( + |request: CatalogLimitsRequest>| -> CatalogLimitsFuture { + Box::pin(async move { + (request.model == "catalog-model" + && request.context.borrowed_config == "catalog-config") + .then(|| { + tinyinference_llm::model::discover::DiscoveryRequest::new( + "https://provider.invalid/v1", + request.model, + ) + }) + }) + }, + )), }) }, ); @@ -176,30 +196,38 @@ async fn independent_live_family_lookups_override_only_their_builtin() { ); assert!(override_error.is_err()); - let embedding = overrides.build_embedding( - "shared", - &EmbeddingRequest { - model: "embed-model", - dimensions: 8, - context: &scoped_context, - }, - |_| panic!("embedding override must run before builtin"), - ); + let embedding = overrides + .build_embedding( + "shared", + EmbeddingRequest { + model: "embed-model".into(), + dimensions: 8, + context: scoped_context, + }, + |_| panic!("embedding override must run before builtin"), + ) + .await; assert_eq!(embedding.unwrap().dimensions(), 8); let fallback_context = HostContext { borrowed_config: "fallback-config", _ownership_marker: std::marker::PhantomData, }; - let fallback = overrides.build_embedding( - "unregistered", - &EmbeddingRequest { - model: "embed-model", - dimensions: 4, - context: &fallback_context, - }, - |_| Ok(Box::new(MockEmbeddingModel::new(4)) as Box), - ); + let fallback = overrides + .build_embedding( + "unregistered", + EmbeddingRequest { + model: "embed-model".into(), + dimensions: 4, + context: fallback_context, + }, + |_| { + Box::pin(async { + Ok(Box::new(MockEmbeddingModel::new(4)) as Box) + }) + }, + ) + .await; assert_eq!(fallback.unwrap().dimensions(), 4); chat_enabled.store(false, Ordering::SeqCst); @@ -211,7 +239,10 @@ async fn independent_live_family_lookups_override_only_their_builtin() { default_model: Some("builtin-model"), native_tool_calling: false, thread_id: None, - context: &fallback_context, + context: &HostContext { + borrowed_config: "fallback-config", + _ownership_marker: std::marker::PhantomData, + }, }, |_| { Ok(( @@ -241,12 +272,22 @@ async fn independent_live_family_lookups_override_only_their_builtin() { _ownership_marker: std::marker::PhantomData, }; let limits = overrides.catalog_entry("shared").unwrap().limits.unwrap(); - assert!(limits("catalog-model", &catalog_context).is_some()); + assert!( + limits(CatalogLimitsRequest { + model: "catalog-model".into(), + context: catalog_context, + }) + .await + .is_some() + ); let listing = overrides .list_catalog( "shared", CatalogRequest { - context: catalog_context, + context: HostContext { + borrowed_config: "catalog-config", + _ownership_marker: std::marker::PhantomData, + }, }, |_| panic!("catalog override must run before builtin"), )