//! 401 attribution: callback hook + shared helpers for tool HTTP clients. use std::sync::Arc; /// Bearer prefix length shared across crate boundaries. pub const SENT_BEARER_PREFIX_LEN: usize = 12; /// Which tool endpoint produced the 401. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ToolConsumer { WebSearch, } impl ToolConsumer { pub fn as_str(self) -> &'static str { match self { Self::WebSearch => "WebSearch", } } } /// 401 attribution callback. Shell wires this to emit telemetry. pub trait Auth401AttributionCallback: Send + Sync + std::fmt::Debug { /// `sent_bearer_prefix` is truncated to [`SENT_BEARER_PREFIX_LEN`] /// before crossing this boundary. `None` = no bearer was sent. fn record_401(&self, consumer: ToolConsumer, sent_bearer_prefix: Option<&str>); } /// Shared, cheap-to-clone alias for the attribution callback. pub type SharedAttributionCallback = Arc; /// Record a 401 attribution event if a callback is wired. Truncates /// the bearer to [`SENT_BEARER_PREFIX_LEN`] before crossing the /// trait boundary. pub(crate) fn emit_401( callback: Option<&SharedAttributionCallback>, consumer: ToolConsumer, sent_bearer: Option<&str>, ) { if let Some(cb) = callback { let prefix = sent_bearer.map(|s| truncate_to_prefix(s.to_string())); cb.record_401(consumer, prefix.as_deref()); } } /// Truncate a bearer string to the first [`SENT_BEARER_PREFIX_LEN`] /// characters. Used by tool clients before passing the bearer across /// the [`Auth401AttributionCallback`] boundary. /// /// Bearer tokens are ASCII (per the `Authorization` header grammar) /// so the byte index is always a char boundary; this function uses /// `String::truncate` which would otherwise panic on a non-boundary /// cut. pub(crate) fn truncate_to_prefix(mut bearer: String) -> String { bearer.truncate(SENT_BEARER_PREFIX_LEN.min(bearer.len())); bearer } #[cfg(test)] mod tests { use super::*; #[test] fn truncate_to_prefix_long_string_cuts_at_12() { assert_eq!( truncate_to_prefix("xai-key-aaaaaaaaaaaaaaaaaaa".to_string()), "xai-key-aaaa" ); } #[test] fn truncate_to_prefix_short_string_unchanged() { assert_eq!(truncate_to_prefix("abc".to_string()), "abc"); } #[test] fn truncate_to_prefix_exact_12_unchanged() { assert_eq!( truncate_to_prefix("123456789012".to_string()), "123456789012" ); assert_eq!(truncate_to_prefix("123456789012".to_string()).len(), 12); } #[test] fn truncate_to_prefix_empty_unchanged() { assert_eq!(truncate_to_prefix(String::new()), ""); } #[test] fn tool_consumer_as_str_stable_identifiers() { assert_eq!(ToolConsumer::WebSearch.as_str(), "WebSearch"); } }