//! 401 attribution callback hook for the sampling client. //! //! Every 401 response site can optionally emit an attribution event so //! a downstream observer can split production 401s into "client sent a //! stale snapshot bearer that the server rejected" vs. "client sent //! the live token from its auth source and the server still rejected //! it" buckets. //! //! `kigi-sampler` is intentionally decoupled from `kigi-shell` //! (no shell types, no logging crate, no auth-manager dependency). The //! caller wires an implementation of [`Auth401AttributionCallback`] //! into [`crate::SamplerConfig::attribution_callback`]; the sampler //! invokes the callback at each UNAUTHORIZED arm with the bearer that //! was actually sent on the wire. The implementation is free to join //! the bearer with whatever live credential source it owns and emit //! the attribution however it wants. //! //! When the callback is `None` (the default), the 401 sites are silent //! and return the same `SamplingError::Auth` they would otherwise. use std::sync::Arc; /// A logical 401-emitting site inside the sampling client. The string /// identifier ends up in the consumer field of the attribution event /// so downstream queries can break down 401s by API path. /// /// # Scope: sampler endpoints only /// /// This enum enumerates the six HTTP endpoints owned by /// `SamplingClient` (chat completions, responses, messages -- each in /// streaming and non-streaming form). It does *not* cover image /// generation, video generation, web search, or embedding -- those /// tools live in `kigi-tools` /// (`crates/codegen/kigi-tools/src/implementations/`), have their /// own HTTP clients that do not flow through `SamplingClient`, and /// hook into the `kigi_tools::ApiKeyProvider` trait rather than /// this enum. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SamplingConsumer { /// `chat_completion_stream`: OpenAI-compatible streaming OpenAI Chat Completions API. ChatCompletionsStream, /// `chat_completion`: OpenAI-compatible non-streaming OpenAI Chat Completions API. ChatCompletions, /// `create_response_stream`: Responses API streaming. ResponsesStream, /// `create_response`: Responses API non-streaming. Responses, /// `messages_stream`: Anthropic Messages API streaming. MessagesStream, /// `messages`: Anthropic Messages API non-streaming. Messages, } impl SamplingConsumer { /// Stable string identifier for this emit site. Callbacks /// typically combine this with a fixed prefix (e.g. the client /// type) when building the consumer field of the attribution /// event. pub fn as_endpoint(self) -> &'static str { match self { Self::ChatCompletionsStream => "chat_completions_stream", Self::ChatCompletions => "chat_completions", Self::ResponsesStream => "responses_stream", Self::Responses => "responses", Self::MessagesStream => "messages_stream", Self::Messages => "messages", } } } /// Maximum prefix length the sampler shares with attribution /// callbacks across the crate boundary. Mirrors /// `kigi_shell::auth::token_suffix` (which truncates to 12 chars /// before any sink) so the two crates stay in lock-step on the /// "bearers leaving the sampler are 12-char prefixes only" invariant. /// /// The cross-crate boundary is the only place this constant is /// load-bearing -- changing it requires updating `token_suffix` in /// `kigi-shell/src/auth/manager.rs` to match, otherwise the /// shell's local-log payload and the sampler's callback argument /// will disagree on prefix length. pub const SENT_BEARER_PREFIX_LEN: usize = 12; /// Hook invoked by [`crate::SamplingClient`] at every 401 response site. /// /// Implementations are responsible for joining `sent_bearer_prefix` /// with whatever live credential source they own (e.g. an auth /// manager holding the most-recently-refreshed token) and emitting /// whatever attribution event makes sense for their observability /// stack. /// /// Implementations must be cheap to invoke and must not block. They /// run inside the request's response-handling path and any latency /// they add is paid by the user-visible 401 error path. // // The `Debug` bound is a structural requirement: [`crate::SamplerConfig`] // derives `Debug` and carries an `Option>` // field, which only compiles when the trait is `Debug`. Do not remove // the bound when factoring this trait out -- it will break // `derive(Debug)` on `SamplerConfig`. pub trait Auth401AttributionCallback: Send + Sync + std::fmt::Debug { /// Record a 401 attribution event for one logical 401 response. /// /// `sent_bearer_prefix` is the **first /// [`SENT_BEARER_PREFIX_LEN`] characters** of the bearer that /// was actually sent on the wire. The sampler extracts the /// bearer from the `Authorization` header (or `x-api-key` for /// Anthropic Messages API backends) and truncates it to the prefix /// length **before crossing this trait boundary** -- the full /// bearer never leaves [`crate::SamplingClient`]. This is the /// scrub-at-the-boundary invariant: even a misbehaving callback /// implementation that logs `sent_bearer_prefix` directly leaks /// only the prefix, never the full credential. /// /// `None` indicates the request had no bearer header at all /// (distinct from "had a bearer that turned out to be stale"). fn record_401(&self, consumer: SamplingConsumer, sent_bearer_prefix: Option<&str>); } /// Shared, cheap-to-clone alias for the attribution callback. pub type SharedAttributionCallback = Arc;