diff --git a/CHANGELOG.md b/CHANGELOG.md index 539f4df9..920994bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,14 @@ ## Unreleased +- TinyHumans managed model: `TinyHumansModel::new_with_credential` takes a + resolver that reports which kind of credential (`Session` or `ApiKey`) each + bearer is, and `TinyHumansAuthFailure::credential()` returns the kind the + failing request itself was sent with. A host can now tell a rejected API key + from an expired session even when the credential changes while requests are + in flight; reading host-side "last resolved" state could attribute a late + API-key 401 to the session that replaced it. `new` is unchanged and reports + `None`. - `response_to_model_response` keeps the provider's finish reason (new `ChatResponse::finish_reason`, passed through `canonical_finish_reason`), so a length cut is no longer reported as `stop`; without one it still infers diff --git a/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs b/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs index 74ba7d95..cf1f068b 100644 --- a/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs +++ b/crates/tinyinference-llm/src/providers/tinyhumans/mod.rs @@ -30,6 +30,34 @@ pub use usage::project_tinyhumans_envelope; /// Resolves the current bearer credential for a managed inference call. pub type BearerResolver = Arc Result + Send + Sync>; +/// Which kind of credential a managed request authenticated with. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TinyHumansCredentialKind { + /// The host's signed-in app session token. + Session, + /// A long-lived API key. + ApiKey, +} + +/// A resolved bearer together with the kind of credential it is. +#[derive(Clone, Debug)] +pub struct TinyHumansCredential { + /// The bearer sent with the request. + pub bearer: String, + /// What the bearer is, reported back on an authentication failure. + pub kind: TinyHumansCredentialKind, +} + +type KindedResolver = + Arc Result<(String, Option)> + Send + Sync>; + +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>; +} + /// Operation that received a provider authentication failure. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum TinyHumansOperation { @@ -73,6 +101,7 @@ pub enum TinyHumansAuthFailureKind { pub struct TinyHumansAuthFailure { operation: TinyHumansOperation, kind: TinyHumansAuthFailureKind, + credential: Option, } impl TinyHumansAuthFailure { @@ -87,6 +116,14 @@ impl TinyHumansAuthFailure { pub fn kind(self) -> TinyHumansAuthFailureKind { self.kind } + + /// The credential kind the failing request was sent with, when the model + /// was built with [`TinyHumansModel::new_with_credential`]; `None` for a + /// plain bearer resolver. + #[must_use] + pub fn credential(self) -> Option { + self.credential + } } /// Handles authentication failures from a managed endpoint. @@ -151,6 +188,42 @@ impl TinyHumansModel { model: impl AsRef, default_model: impl AsRef, resolver: impl Fn() -> Result + Send + Sync + 'static, + ) -> Result { + Self::build( + base_url, + model, + default_model, + Arc::new(move || resolver().map(|bearer| (bearer, None))), + ) + } + + /// Like [`Self::new`], but the resolver also reports which kind of + /// credential each bearer is, and every [`TinyHumansAuthFailure`] carries + /// the kind its own request was sent with. + /// + /// # Errors + /// Returns [`Error::Validation`] when the endpoint or default model is empty. + pub fn new_with_credential( + base_url: impl AsRef, + model: impl AsRef, + default_model: impl AsRef, + resolver: impl Fn() -> Result + Send + Sync + 'static, + ) -> Result { + Self::build( + base_url, + model, + default_model, + Arc::new(move || { + resolver().map(|credential| (credential.bearer, Some(credential.kind))) + }), + ) + } + + fn build( + base_url: impl AsRef, + model: impl AsRef, + default_model: impl AsRef, + resolver: KindedResolver, ) -> Result { let base = base_url.as_ref().trim().trim_end_matches('/'); if base.is_empty() { @@ -166,7 +239,7 @@ impl TinyHumansModel { } let endpoint = catalog::endpoint(base); let default_model = resolve_model(model.as_ref(), fallback_model); - let source = Arc::new(ResolvedBearer(Arc::new(resolver))); + let source = Arc::new(ResolvedBearer(resolver)); let wire = OpenAiModel::compatible_provider("tinyhumans", "", &endpoint, &default_model) .with_auth_style(AuthStyle::Bearer) .with_bearer_source(source) @@ -311,7 +384,13 @@ impl TinyHumansModel { request: ModelRequest, operation: TinyHumansOperation, ) -> Result { - match self.model.invoke(&(), self.prepare_request(request)).await { + let (result, credential) = SENT_CREDENTIAL + .scope(std::cell::Cell::new(None), async { + let result = self.model.invoke(&(), self.prepare_request(request)).await; + (result, SENT_CREDENTIAL.with(std::cell::Cell::get)) + }) + .await; + match result { Ok(response) => Ok(project_tinyhumans_envelope(response)), Err(error) => { let point = match operation { @@ -320,17 +399,22 @@ impl TinyHumansModel { TinyHumansOperation::Stream => unreachable!("invoke operation is not stream"), }; self.observe_error(&error, point); - self.observe_auth_error(&error, operation); + self.observe_auth_error(&error, operation, credential); Err(error) } } } - fn observe_auth_error(&self, error: &Error, operation: TinyHumansOperation) { + fn observe_auth_error( + &self, + error: &Error, + operation: TinyHumansOperation, + credential: Option, + ) { if let Error::Provider(provider_error) = error && is_auth_failure(provider_error) { - self.notify_auth_failure(provider_error, operation); + self.notify_auth_failure(provider_error, operation, credential); } } @@ -340,7 +424,12 @@ impl TinyHumansModel { } } - fn notify_auth_failure(&self, error: &ProviderError, operation: TinyHumansOperation) { + fn notify_auth_failure( + &self, + error: &ProviderError, + operation: TinyHumansOperation, + credential: Option, + ) { if let Some(handler) = &self.auth_failure_handler { handler(error, operation); } @@ -350,6 +439,7 @@ impl TinyHumansModel { TinyHumansAuthFailure { operation, kind: auth_failure_kind(error), + credential, }, ); } @@ -387,11 +477,17 @@ impl ChatModel<()> for TinyHumansModel { } async fn stream(&self, _state: &(), request: ModelRequest) -> Result { - let stream = match self.model.stream(&(), self.prepare_request(request)).await { + let (opened, credential) = SENT_CREDENTIAL + .scope(std::cell::Cell::new(None), async { + let opened = self.model.stream(&(), self.prepare_request(request)).await; + (opened, SENT_CREDENTIAL.with(std::cell::Cell::get)) + }) + .await; + let stream = match opened { Ok(stream) => stream, Err(error) => { self.observe_error(&error, TinyHumansFailurePoint::StreamOpen); - self.observe_auth_error(&error, TinyHumansOperation::Stream); + self.observe_auth_error(&error, TinyHumansOperation::Stream, credential); return Err(error); } }; @@ -419,6 +515,7 @@ impl ChatModel<()> for TinyHumansModel { TinyHumansAuthFailure { operation: TinyHumansOperation::Stream, kind: auth_failure_kind(&error), + credential, }, ); } @@ -477,12 +574,16 @@ fn resolve_model(model: &str, default_model: &str) -> String { } } -struct ResolvedBearer(BearerResolver); +struct ResolvedBearer(KindedResolver); #[async_trait] impl BearerSource for ResolvedBearer { async fn current(&self) -> Result> { - (self.0)().map(Some) + let (bearer, kind) = (self.0)()?; + // 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)); + 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 a22afd0f..8e3f9190 100644 --- a/crates/tinyinference-llm/src/providers/tinyhumans/mod_tests.rs +++ b/crates/tinyinference-llm/src/providers/tinyhumans/mod_tests.rs @@ -715,3 +715,81 @@ mod cache_identity_tests; #[path = "observer_tests.rs"] mod observer_tests; + +/// Each failure carries the credential kind its own request was sent with, +/// even when the resolver switches kinds between requests. +#[tokio::test] +async fn managed_auth_failure_reports_the_credential_its_request_carried() { + 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 use_api_key = Arc::new(std::sync::atomic::AtomicBool::new(true)); + let resolver_use_api_key = use_api_key.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 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") + .with_managed_auth_failure_handler(move |_, failure| { + seen_by_handler + .lock() + .unwrap() + .push((failure.operation(), failure.credential())); + }); + + let request = || ModelRequest::new(vec![crate::message::Message::user("hi")]); + assert!(model.invoke(&(), request()).await.is_err()); + use_api_key.store(false, Ordering::SeqCst); + assert!(model.invoke(&(), request()).await.is_err()); + assert!(model.stream(&(), request()).await.is_err()); + + assert_eq!( + *seen.lock().unwrap(), + vec![ + ( + TinyHumansOperation::Invoke, + Some(TinyHumansCredentialKind::ApiKey) + ), + ( + TinyHumansOperation::Invoke, + Some(TinyHumansCredentialKind::Session) + ), + ( + TinyHumansOperation::Stream, + Some(TinyHumansCredentialKind::Session) + ), + ] + ); +} + +#[tokio::test] +async fn a_plain_bearer_resolver_reports_no_credential_kind() { + 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 seen = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen_by_handler = seen.clone(); + let model = fixture_model(&base).with_managed_auth_failure_handler(move |_, failure| { + seen_by_handler.lock().unwrap().push(failure.credential()); + }); + let request = ModelRequest::new(vec![crate::message::Message::user("hi")]); + assert!(model.invoke(&(), request).await.is_err()); + assert_eq!(*seen.lock().unwrap(), vec![None]); +}