1use std::collections::BTreeMap;
23use std::pin::Pin;
24use std::sync::Arc;
25use std::time::Duration;
26
27use async_trait::async_trait;
28use futures::Stream;
29use serde::{Deserialize, Serialize};
30
31use super::account::{
32 seat_fingerprint, Account, AccountPolicy, AccountSelector, AccountSummary, CredentialRef,
33 SeatIdentity,
34};
35use super::cloud::{CloudProviderId, InvokeRequest};
36use super::provider::OAuthProviderId;
37use super::static_providers::{
38 allowed_proxy_paths, static_provider, StaticProviderSpec, LOCAL_DEFAULT_BASE_URL,
39 LOCAL_PROVIDER_KEY, STATIC_PROVIDERS,
40};
41use super::storage;
42
43pub const MAX_PROXY_REQUEST_BYTES: usize = 2 * 1024 * 1024;
52pub const MAX_PROXY_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
54pub const MAX_UPSTREAM_ERROR_BYTES: usize = 2 * 1024;
58pub const PROXY_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
60
61const ANTHROPIC_USAGE_URL: &str = "https://api.anthropic.com/api/oauth/usage";
64const ANTHROPIC_OAUTH_BETA: &str = "oauth-2025-04-20";
65
66pub const GOOGLE_GEMINI_CODE_ASSIST_BASE_URL: &str = "https://cloudcode-pa.googleapis.com";
70
71pub(crate) fn is_allowed_google_gemini_path(path: &str) -> bool {
75 matches!(
76 path,
77 "/v1internal:loadCodeAssist"
78 | "/v1internal:onboardUser"
79 | "/v1internal:streamGenerateContent"
80 | "/v1internal:countTokens"
81 ) || path.starts_with("/v1internal/operations/")
82}
83
84pub const OPENAI_CODEX_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
86
87pub(crate) const OPENAI_CODEX_MODELS_PATH: &str = "/codex/models";
96
97pub(crate) fn is_allowed_openai_codex_path(path: &str) -> bool {
98 let Some((base, query)) = path.split_once('?') else {
99 return false;
100 };
101 if base != OPENAI_CODEX_MODELS_PATH {
102 return false;
103 }
104 let mut client_version: Option<&str> = None;
105 for pair in query.split('&') {
106 let Some((k, v)) = pair.split_once('=') else {
107 return false;
108 };
109 if k != "client_version" || client_version.is_some() {
110 return false;
111 }
112 if v.trim().is_empty() {
113 return false;
114 }
115 client_version = Some(v);
116 }
117 client_version.is_some()
118}
119
120pub(crate) fn is_allowed_anthropic_path(path: &str) -> bool {
125 if path == "/v1/models" {
126 return true;
127 }
128 let Some((base, query)) = path.split_once('?') else {
129 return false;
130 };
131 if base != "/v1/models" {
132 return false;
133 }
134 if query.is_empty() {
135 return false;
136 }
137 let mut saw_limit = false;
138 let mut saw_after_id = false;
139 for pair in query.split('&') {
140 let Some((k, v)) = pair.split_once('=') else {
141 return false;
142 };
143 if v.is_empty() {
144 return false;
145 }
146 match k {
147 "limit" if !saw_limit => saw_limit = true,
148 "after_id" if !saw_after_id => saw_after_id = true,
149 _ => return false,
150 }
151 }
152 true
155}
156
157pub(crate) fn anthropic_oauth_catalog_request_headers() -> &'static [(&'static str, &'static str)] {
161 &[
162 ("anthropic-version", "2023-06-01"),
163 ("anthropic-beta", ANTHROPIC_OAUTH_BETA),
164 ("accept", "application/json"),
165 ]
166}
167
168#[derive(Debug, Clone, PartialEq, Eq)]
173#[non_exhaustive]
174pub enum BrokerError {
175 UnknownProvider(String),
177 RegistrationRequired {
179 provider: String,
180 remediation: String,
181 },
182 NotConfigured(String),
184 Unauthorized,
186 Denied(String),
188 UnsupportedCapability {
191 provider: String,
192 capability: String,
193 },
194 Transport(String),
196 Credential(String),
198 UnknownAccount { provider: String, label: String },
200 UnsupportedAccount { provider: String, label: String },
202 InvalidAccount(String),
204 NoAccountAvailable {
208 provider: String,
209 reason: String,
210 failure: Option<super::quota_policy::SelectionFailure>,
211 },
212 BrokerTooOld(String),
215 Refresh(super::token::RefreshFailure),
219}
220
221impl std::fmt::Display for BrokerError {
222 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
223 match self {
224 Self::UnknownProvider(p) => write!(f, "unknown provider: {p}"),
225 Self::RegistrationRequired {
226 provider,
227 remediation,
228 } => {
229 write!(f, "registration required for '{provider}': {remediation}")
230 }
231 Self::NotConfigured(p) => write!(
232 f,
233 "no credential configured for '{p}'. Run `synaps login` to add one."
234 ),
235 Self::Unauthorized => f.write_str(&super::broker_token::unauthorized_hint()),
236 Self::Denied(msg) => write!(f, "broker denied request: {msg}"),
237 Self::UnsupportedCapability {
238 provider,
239 capability,
240 } => write!(
241 f,
242 "cloud provider '{provider}' is text-only: {capability} are not supported yet \
243 (no credential was used and no network request was made)"
244 ),
245 Self::Transport(msg) => write!(f, "broker transport error: {msg}"),
246 Self::Credential(msg) => write!(f, "credential error: {msg}"),
247 Self::Refresh(failure) => {
248 use super::token::RefreshFailureKind as K;
249 let hint = match failure.kind {
250 K::Transient => " (network or provider error, not a login problem; retry shortly)",
251 K::LockBusy => " (another refresh of this account is still running; retry shortly)",
252 K::Other => " (retry shortly)",
253 K::Credential | K::Cooldown { .. } | K::ActionRequired => "",
254 };
255 write!(f, "credential refresh failed{hint}: {failure}")
256 }
257 Self::UnknownAccount { provider, label } => write!(
258 f,
259 "unknown account '{label}' for provider '{provider}' (no fallback; run \
260 `synaps auth list` or `synaps login --provider {provider} --account {label}`)"
261 ),
262 Self::UnsupportedAccount { provider, label } => write!(
263 f,
264 "this credential broker cannot address account '{label}' for provider '{provider}'"
265 ),
266 Self::InvalidAccount(msg) => write!(f, "invalid account: {msg}"),
267 Self::NoAccountAvailable {
268 provider,
269 reason,
270 failure: Some(failure),
271 } => write!(
272 f,
273 "no account with proven capacity for '{provider}' ({}): {reason}",
274 failure.kind
275 ),
276 Self::NoAccountAvailable {
277 provider, reason, ..
278 } => {
279 write!(
280 f,
281 "no account with proven capacity for '{provider}': {reason}"
282 )
283 }
284 Self::BrokerTooOld(what) => write!(
285 f,
286 "broker is older than this client; {what} unavailable (upgrade the broker host)"
287 ),
288 }
289 }
290}
291
292impl std::error::Error for BrokerError {}
293
294#[cfg(test)]
295#[path = "broker_remote_compat_tests.rs"]
296mod remote_compat_tests;
297
298#[derive(Debug, Clone, Serialize, Deserialize)]
305pub struct AccessToken {
306 pub token: String,
307 pub expires: u64,
309}
310
311#[derive(Debug, Clone)]
316pub struct PinnedToken {
317 pub credential: CredentialRef,
318 pub token: AccessToken,
319}
320
321impl PinnedToken {
322 pub fn seat_fingerprint(&self) -> Option<String> {
328 match self.credential.provider {
329 OAuthProviderId::OpenAiCodex => super::extract_codex_account_id(&self.token.token)
330 .and_then(|id| seat_fingerprint(self.credential.provider, &id)),
331 _ => None,
332 }
333 }
334}
335
336#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
338#[serde(rename_all = "lowercase")]
339pub enum ProxyMethod {
340 Get,
341 Post,
342}
343
344#[derive(Debug, Clone, Serialize, Deserialize)]
352pub struct ProxyRequest {
353 pub provider: String,
359 pub method: ProxyMethod,
360 pub path: String,
362 #[serde(default, skip_serializing_if = "Option::is_none")]
364 pub body: Option<serde_json::Value>,
365 #[serde(default)]
367 pub stream: bool,
368 #[serde(skip)]
377 pub body_bytes: Option<bytes::Bytes>,
378}
379
380impl ProxyRequest {
381 pub fn validate(&self) -> Result<(), BrokerError> {
384 let (provider, _account) = self.split_provider()?;
385 let provider = provider.as_str();
386 if provider != LOCAL_PROVIDER_KEY
387 && provider != "xai-auth"
388 && provider != "github-copilot"
389 && provider != "google-gemini"
390 && provider != "openai-codex"
391 && provider != "anthropic"
392 && provider != "kimi-code"
393 && provider != "muse"
394 && static_provider(provider).is_none()
395 {
396 return Err(BrokerError::UnknownProvider(self.provider.clone()));
397 }
398 if !self.path.starts_with('/') {
399 return Err(BrokerError::Denied("proxy path must be relative".into()));
400 }
401 if self.path.contains("://") || self.path.contains("..") {
402 return Err(BrokerError::Denied(
403 "proxy path is not a plain relative path".into(),
404 ));
405 }
406 let oauth_path_allowed = provider == "xai-auth" && self.path == "/responses"
410 || provider == "github-copilot"
411 && matches!(
412 self.path.as_str(),
413 "/models" | "/chat/completions" | "/responses"
414 )
415 || provider == "google-gemini" && is_allowed_google_gemini_path(&self.path)
416 || provider == "kimi-code"
419 && (self.path == "/chat/completions"
420 || self.method == ProxyMethod::Get
421 && matches!(self.path.as_str(), "/models" | "/me" | "/usages"))
422 || provider == "muse"
425 && (self.method == ProxyMethod::Post && self.path == "/v1/responses"
426 || self.method == ProxyMethod::Get && self.path == "/muse-code/models")
427 || provider == "openai-codex"
428 && self.method == ProxyMethod::Get
429 && is_allowed_openai_codex_path(&self.path)
430 || provider == "anthropic"
431 && self.method == ProxyMethod::Get
432 && is_allowed_anthropic_path(&self.path);
433 if !oauth_path_allowed && !allowed_proxy_paths(provider).contains(&self.path.as_str()) {
434 return Err(BrokerError::Denied(format!(
435 "proxy path '{}' is not in the provider's endpoint allowlist",
436 self.path
437 )));
438 }
439 let size = match (&self.body_bytes, &self.body) {
444 (Some(bytes), _) => Some(bytes.len()),
446 (None, Some(body)) => Some(
447 serde_json::to_vec(body)
448 .map(|v| v.len())
449 .unwrap_or(usize::MAX),
450 ),
451 (None, None) => None,
452 };
453 if let Some(size) = size {
454 if size > MAX_PROXY_REQUEST_BYTES {
455 return Err(BrokerError::Denied(format!(
456 "request body exceeds the {MAX_PROXY_REQUEST_BYTES}-byte broker limit"
457 )));
458 }
459 }
460 #[cfg(any(test, debug_assertions))]
466 if let (Some(bytes), Some(body)) = (&self.body_bytes, &self.body) {
467 let parsed = serde_json::from_slice::<serde_json::Value>(bytes).ok();
468 if parsed.as_ref() != Some(body) {
469 return Err(BrokerError::Denied(
470 "body_bytes does not match the JSON body (exact-byte handoff incoherent)"
471 .into(),
472 ));
473 }
474 }
475 Ok(())
476 }
477
478 pub fn post_json_exact(
487 provider: impl Into<String>,
488 path: impl Into<String>,
489 body: serde_json::Value,
490 stream: bool,
491 ) -> Result<(Self, bytes::Bytes), BrokerError> {
492 let bytes =
493 bytes::Bytes::from(serde_json::to_vec(&body).map_err(|e| {
494 BrokerError::Denied(format!("request body failed to serialize: {e}"))
495 })?);
496 let request = Self {
497 provider: provider.into(),
498 method: ProxyMethod::Post,
499 path: path.into(),
500 body: Some(body),
501 stream,
502 body_bytes: Some(bytes.clone()),
503 };
504 Ok((request, bytes))
505 }
506
507 pub fn with_account(mut self, account: &Account) -> Self {
510 let base = self
511 .provider
512 .split_once('@')
513 .map(|(p, _)| p.to_string())
514 .unwrap_or_else(|| self.provider.clone());
515 self.provider = format!("{base}@{}", account.label_str());
516 self
517 }
518
519 pub fn split_provider(&self) -> Result<(String, Option<Account>), BrokerError> {
522 match self.provider.split_once('@') {
523 None => Ok((self.provider.clone(), None)),
524 Some((base, account)) => {
525 if base.parse::<OAuthProviderId>().is_err() {
526 return Err(BrokerError::UnknownProvider(self.provider.clone()));
527 }
528 let account = Account::parse(account).map_err(BrokerError::InvalidAccount)?;
529 Ok((base.to_string(), Some(account)))
530 }
531 }
532 }
533
534 pub fn base_provider(&self) -> Result<String, BrokerError> {
536 Ok(self.split_provider()?.0)
537 }
538
539 pub fn explicit_account(&self) -> Result<Option<Account>, BrokerError> {
541 Ok(self.split_provider()?.1)
542 }
543}
544
545#[derive(Debug, Clone, Serialize, Deserialize)]
547pub struct ProxyResponse {
548 pub status: u16,
549 pub body: String,
550}
551
552pub type ProxyByteStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, BrokerError>> + Send>>;
554
555#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
557#[serde(rename_all = "snake_case")]
558pub enum CredentialKind {
559 OAuth,
560 StaticKey,
561 LocalEndpoint,
562}
563
564#[derive(Debug, Clone, Serialize, Deserialize)]
566pub struct ProviderStatus {
567 pub key: String,
568 pub name: String,
569 pub kind: CredentialKind,
570 pub configured: bool,
573 #[serde(default, skip_serializing_if = "Vec::is_empty")]
575 pub accounts: Vec<AccountSummary>,
576}
577
578#[derive(Debug, Clone, PartialEq, Eq)]
580pub enum StaticKeyStatus {
581 NotSet,
582 Configured {
585 masked: String,
586 },
587 FromEnv,
589}
590
591impl StaticKeyStatus {
592 pub fn is_configured(&self) -> bool {
593 !matches!(self, StaticKeyStatus::NotSet)
594 }
595}
596
597#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
601pub struct CloudCatalogEntry {
602 pub provider: CloudProviderId,
603 pub id: String,
604 pub display_name: String,
605 pub context_ref: String,
607 pub context_label: String,
608 pub stale: bool,
609 pub fetched_at: u64,
611}
612
613#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
615#[serde(tag = "type", rename_all = "snake_case")]
616pub enum CloudEvent {
617 TextDelta {
618 delta: String,
619 },
620 ToolArguments {
621 id: String,
622 name: Option<String>,
623 delta: String,
624 },
625 Usage {
626 input_tokens: u64,
627 output_tokens: u64,
628 },
629 Done,
630}
631pub type CloudEventStream = Pin<Box<dyn Stream<Item = Result<CloudEvent, BrokerError>> + Send>>;
632
633#[async_trait]
637pub trait CredentialBroker: Send + Sync {
638 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError>;
641
642 fn account_selector(&self, provider: OAuthProviderId) -> AccountSelector {
645 let _ = provider;
646 AccountSelector::Account(Account::Default)
647 }
648
649 async fn access_token_for(&self, cred: &CredentialRef) -> Result<AccessToken, BrokerError> {
655 match &cred.account {
656 Account::Default => self.access_token(cred.provider).await,
657 Account::Named(label) => Err(BrokerError::UnsupportedAccount {
658 provider: cred.provider.as_str().to_string(),
659 label: label.as_str().to_string(),
660 }),
661 }
662 }
663
664 async fn recover_rejected_token(
668 &self,
669 cred: &CredentialRef,
670 _rejected: &str,
671 ) -> Result<AccessToken, BrokerError> {
672 self.access_token_for(cred).await
673 }
674
675 async fn access_token_pinned(
678 &self,
679 provider: OAuthProviderId,
680 ) -> Result<PinnedToken, BrokerError> {
681 let token = self.access_token(provider).await?;
682 Ok(PinnedToken {
683 credential: CredentialRef::default_for(provider),
684 token,
685 })
686 }
687
688 async fn access_token_pinned_for(
694 &self,
695 provider: OAuthProviderId,
696 model: Option<&str>,
697 ) -> Result<PinnedToken, BrokerError> {
698 let _ = model;
699 self.access_token_pinned(provider).await
700 }
701
702 async fn accounts(&self, provider: OAuthProviderId) -> Result<Vec<AccountSummary>, BrokerError> {
704 Err(BrokerError::UnsupportedCapability {
705 provider: provider.as_str().to_string(),
706 capability: "accounts".into(),
707 })
708 }
709
710 async fn usage(&self, cred: &CredentialRef) -> Result<super::usage::UsageSnapshot, BrokerError> {
714 Err(BrokerError::UnsupportedCapability {
715 provider: cred.provider.as_str().to_string(),
716 capability: "usage".into(),
717 })
718 }
719
720 async fn report_cooldown(
724 &self,
725 cred: &CredentialRef,
726 until_ms: Option<u64>,
727 reason: &str,
728 ) -> Result<(), BrokerError> {
729 let _ = (cred, until_ms, reason);
730 Ok(())
731 }
732
733 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError>;
735
736 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError>;
740
741 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError>;
748
749 async fn cloud_catalog(
751 &self,
752 provider: CloudProviderId,
753 context_ref: &str,
754 allow_stale: bool,
755 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
756 let _ = (context_ref, allow_stale);
757 Err(BrokerError::NotConfigured(provider.to_string()))
758 }
759
760 async fn cloud_invoke(
763 &self,
764 provider: CloudProviderId,
765 context_ref: &str,
766 model_id: &str,
767 request: InvokeRequest,
768 ) -> Result<CloudEventStream, BrokerError> {
769 let _ = (context_ref, model_id, request);
770 Err(BrokerError::NotConfigured(provider.to_string()))
771 }
772
773 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError>;
775}
776
777#[async_trait]
780pub trait CloudBackend: Send + Sync {
781 async fn catalog(
782 &self,
783 provider: CloudProviderId,
784 context_ref: &str,
785 allow_stale: bool,
786 ) -> Result<Vec<CloudCatalogEntry>, BrokerError>;
787 async fn invoke(
788 &self,
789 provider: CloudProviderId,
790 context_ref: &str,
791 model_id: &str,
792 request: InvokeRequest,
793 ) -> Result<CloudEventStream, BrokerError>;
794}
795
796#[derive(Clone)]
797struct ProductionCloudBackend {
798 http: reqwest::Client,
801 refresh_lock: Arc<tokio::sync::Mutex<()>>,
802}
803impl ProductionCloudBackend {
804 fn new(_http: reqwest::Client) -> Self {
805 let http = reqwest::Client::builder()
806 .redirect(reqwest::redirect::Policy::none())
807 .connect_timeout(Duration::from_secs(10))
808 .timeout(Duration::from_secs(90))
809 .build()
810 .expect("cloud HTTPS client configuration is valid");
811 Self {
812 http,
813 refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
814 }
815 }
816
817 fn state(&self, provider: CloudProviderId) -> Result<serde_json::Value, BrokerError> {
818 storage::load_cloud_state(provider.as_str())
819 .map_err(BrokerError::Credential)?
820 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))
821 }
822
823 fn context_ref(&self, provider: CloudProviderId) -> Result<String, BrokerError> {
824 use sha2::{Digest, Sha256};
825 let state = self.state(provider)?;
826 let public = serde_json::to_vec(&state["config"])
827 .map_err(|_| BrokerError::Credential("invalid cloud context".into()))?;
828 let digest = Sha256::digest([provider.as_str().as_bytes(), &public].concat());
829 Ok(format!(
830 "ctx-{}",
831 digest[..16]
832 .iter()
833 .map(|b| format!("{b:02x}"))
834 .collect::<String>()
835 ))
836 }
837
838 fn validate_context(
839 &self,
840 provider: CloudProviderId,
841 supplied: &str,
842 ) -> Result<String, BrokerError> {
843 let opaque = self.context_ref(provider)?;
844 if supplied == provider.as_str() || supplied == opaque {
847 Ok(opaque)
848 } else {
849 Err(BrokerError::Denied(
850 "cloud context does not match stored provider".into(),
851 ))
852 }
853 }
854
855 async fn aws(
856 &self,
857 ) -> Result<super::aws_bedrock::AwsBedrockBroker<super::aws_bedrock::AwsHttpApi>, BrokerError>
858 {
859 #[derive(Deserialize)]
860 struct State {
861 config: super::cloud::AwsBedrockConfig,
862 access_key: String,
863 secret_key: String,
864 session_token: String,
865 expires_at: u64,
866 registered_client: serde_json::Value,
867 sso_access_token: String,
868 sso_refresh_token: Option<String>,
869 sso_expires_at: u64,
870 }
871 let decode = |value| {
872 serde_json::from_value::<State>(value)
873 .map_err(|_| BrokerError::Credential("invalid AWS broker state".into()))
874 };
875 let mut raw = self.state(CloudProviderId::AwsBedrock)?;
876 let now = std::time::SystemTime::now()
877 .duration_since(std::time::UNIX_EPOCH)
878 .unwrap_or_default()
879 .as_millis() as u64;
880 let mut s = decode(raw.clone())?;
881 if s.expires_at <= now + 60_000 {
882 let _guard = self.refresh_lock.lock().await;
885 raw = self.state(CloudProviderId::AwsBedrock)?;
886 s = decode(raw.clone())?;
887 if s.expires_at <= now + 60_000 {
888 let api =
889 super::aws_bedrock::AwsHttpApi::new(self.http.clone(), &s.config.sso_region);
890 if s.sso_expires_at <= now + 60_000 {
891 let refresh = s
892 .sso_refresh_token
893 .as_deref()
894 .filter(|v| !v.is_empty())
895 .ok_or_else(|| {
896 BrokerError::Credential("AWS SSO session expired; login again".into())
897 })?;
898 let client = super::aws_bedrock::RegisteredClient::new(
899 s.registered_client["id"].as_str().ok_or_else(|| {
900 BrokerError::Credential("invalid AWS client registration".into())
901 })?,
902 s.registered_client["secret"].as_str().ok_or_else(|| {
903 BrokerError::Credential("invalid AWS client registration".into())
904 })?,
905 s.registered_client["expires_at"].as_u64().ok_or_else(|| {
906 BrokerError::Credential("invalid AWS client registration".into())
907 })?,
908 );
909 if client.expires_at <= now / 1000 {
910 return Err(BrokerError::Credential(
911 "AWS client registration expired; login again".into(),
912 ));
913 }
914 use super::aws_bedrock::AwsApi;
915 let token = api
916 .create_token(
917 &client,
918 &s.config.sso_region,
919 super::aws_bedrock::TokenGrant::RefreshToken(refresh),
920 )
921 .await
922 .map_err(|_| {
923 BrokerError::Credential("AWS SSO refresh rejected; login again".into())
924 })?;
925 raw["sso_access_token"] = token.access().into();
926 if let Some(rotated) = token.refresh() {
927 raw["sso_refresh_token"] = rotated.into();
928 }
929 raw["sso_expires_at"] = (now + token.expires_in * 1000).into();
930 storage::save_cloud_state("aws-bedrock", &raw)
935 .map_err(BrokerError::Credential)?;
936 s = decode(raw.clone())?;
937 }
938 use super::aws_bedrock::AwsApi;
939 let credentials = api
940 .get_role_credentials(
941 &s.config.sso_region,
942 &s.sso_access_token,
943 &s.config.account_id,
944 &s.config.role_name,
945 )
946 .await
947 .map_err(|_| {
948 BrokerError::Credential("AWS role refresh rejected; login again".into())
949 })?;
950 raw["access_key"] = credentials.access_key().into();
951 raw["secret_key"] = credentials.secret_key().into();
952 raw["session_token"] = credentials.session_token().into();
953 raw["expires_at"] = credentials.expires_at.into();
954 storage::save_cloud_state("aws-bedrock", &raw).map_err(BrokerError::Credential)?;
955 s = decode(raw)?;
956 }
957 }
958 let api = super::aws_bedrock::AwsHttpApi::new(self.http.clone(), &s.config.sso_region);
959 Ok(super::aws_bedrock::AwsBedrockBroker::from_credentials(
960 api,
961 s.config,
962 super::aws_bedrock::RoleCredentials::new(
963 s.access_key,
964 s.secret_key,
965 s.session_token,
966 s.expires_at,
967 ),
968 ))
969 }
970 async fn azure_request(
971 &self,
972 state: &mut serde_json::Value,
973 audience: super::azure_openai::AzureAudience,
974 ) -> Result<String, BrokerError> {
975 let now = std::time::SystemTime::now()
976 .duration_since(std::time::UNIX_EPOCH)
977 .unwrap_or_default()
978 .as_millis() as u64;
979 let key = match audience {
980 super::azure_openai::AzureAudience::Arm => "arm",
981 super::azure_openai::AzureAudience::Inference => "inference",
982 };
983 if state[key]["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
984 return state[key]["access_token"]
985 .as_str()
986 .map(str::to_owned)
987 .ok_or_else(|| BrokerError::Credential("invalid Azure token state".into()));
988 }
989 let _refresh = self.refresh_lock.lock().await;
992 *state = self.state(CloudProviderId::AzureOpenAi)?;
993 if state[key]["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
994 return state[key]["access_token"]
995 .as_str()
996 .map(str::to_owned)
997 .ok_or_else(|| BrokerError::Credential("invalid Azure token state".into()));
998 }
999 let config: super::cloud::AzureOpenAiConfig =
1000 serde_json::from_value(state["config"].clone())
1001 .map_err(|_| BrokerError::Credential("invalid Azure broker state".into()))?;
1002 let client_id = state["client_id"]
1003 .as_str()
1004 .ok_or_else(|| BrokerError::Credential("invalid Azure client registration".into()))?;
1005 let refresh = state["refresh_token"]
1006 .as_str()
1007 .ok_or_else(|| BrokerError::Credential("invalid Azure refresh state".into()))?;
1008 let reg = super::azure_openai::AzureRegistration::production(Some(client_id.into()))
1009 .map_err(|e| BrokerError::Credential(e.to_string()))?;
1010 let r = super::azure_openai::refresh_request(&config, ®, audience, refresh)
1011 .map_err(|e| BrokerError::Credential(e.to_string()))?;
1012 let response = self
1013 .http
1014 .post(r.url)
1015 .form(&r.form)
1016 .send()
1017 .await
1018 .map_err(|_| BrokerError::Transport("Azure refresh failed".into()))?;
1019 if !response.status().is_success() {
1020 return Err(BrokerError::Credential(
1021 "Azure refresh rejected; login again".into(),
1022 ));
1023 }
1024 let wire: serde_json::Value = response
1025 .json()
1026 .await
1027 .map_err(|_| BrokerError::Transport("invalid Azure token response".into()))?;
1028 let access = wire["access_token"]
1029 .as_str()
1030 .filter(|v| !v.is_empty())
1031 .ok_or_else(|| BrokerError::Transport("invalid Azure token response".into()))?
1032 .to_owned();
1033 if let Some(r) = wire["refresh_token"].as_str().filter(|v| !v.is_empty()) {
1034 state["refresh_token"] = r.into();
1035 }
1036 state[key] = serde_json::json!({"access_token": access, "expires_at": now + wire["expires_in"].as_u64().unwrap_or(3600) * 1000});
1037 storage::save_cloud_state("azure-openai", state).map_err(BrokerError::Credential)?;
1038 Ok(access)
1039 }
1040
1041 async fn vertex_request(&self, state: &mut serde_json::Value) -> Result<String, BrokerError> {
1042 let now = std::time::SystemTime::now()
1043 .duration_since(std::time::UNIX_EPOCH)
1044 .unwrap_or_default()
1045 .as_millis() as u64;
1046 if state["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
1047 return state["access_token"]
1048 .as_str()
1049 .map(str::to_owned)
1050 .ok_or_else(|| BrokerError::Credential("invalid Vertex token state".into()));
1051 }
1052 let _refresh = self.refresh_lock.lock().await;
1053 *state = self.state(CloudProviderId::GoogleVertex)?;
1054 if state["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
1055 return state["access_token"]
1056 .as_str()
1057 .map(str::to_owned)
1058 .ok_or_else(|| BrokerError::Credential("invalid Vertex token state".into()));
1059 }
1060 let client_id = state["client_id"]
1061 .as_str()
1062 .ok_or_else(|| BrokerError::Credential("invalid Vertex registration".into()))?;
1063 let refresh = state["refresh_token"]
1064 .as_str()
1065 .ok_or_else(|| BrokerError::Credential("invalid Vertex refresh state".into()))?;
1066 let response = self
1067 .http
1068 .post(super::google_vertex::TOKEN_URL)
1069 .form(&[
1070 ("client_id", client_id),
1071 ("grant_type", "refresh_token"),
1072 ("refresh_token", refresh),
1073 ])
1074 .send()
1075 .await
1076 .map_err(|_| BrokerError::Transport("Vertex refresh failed".into()))?;
1077 if !response.status().is_success() {
1078 return Err(BrokerError::Credential(
1079 "Vertex refresh rejected; login again".into(),
1080 ));
1081 }
1082 let wire: serde_json::Value = response
1083 .json()
1084 .await
1085 .map_err(|_| BrokerError::Transport("invalid Vertex token response".into()))?;
1086 let access = wire["access_token"]
1087 .as_str()
1088 .filter(|v| !v.is_empty())
1089 .ok_or_else(|| BrokerError::Transport("invalid Vertex token response".into()))?
1090 .to_owned();
1091 if let Some(r) = wire["refresh_token"].as_str().filter(|v| !v.is_empty()) {
1092 state["refresh_token"] = r.into();
1093 }
1094 state["access_token"] = access.clone().into();
1095 state["expires_at"] = (now + wire["expires_in"].as_u64().unwrap_or(3600) * 1000).into();
1096 storage::save_cloud_state("google-vertex", state).map_err(BrokerError::Credential)?;
1097 Ok(access)
1098 }
1099}
1100
1101#[async_trait]
1102impl CloudBackend for ProductionCloudBackend {
1103 async fn catalog(
1104 &self,
1105 provider: CloudProviderId,
1106 context_ref: &str,
1107 _allow_stale: bool,
1108 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
1109 let opaque_context = self.validate_context(provider, context_ref)?;
1110 let fetched_at = std::time::SystemTime::now()
1111 .duration_since(std::time::UNIX_EPOCH)
1112 .unwrap_or_default()
1113 .as_millis() as u64;
1114 match provider {
1115 CloudProviderId::AwsBedrock => {
1116 let broker = self.aws().await?;
1117 let label = broker.public_context().bedrock_region;
1118 broker
1119 .catalog()
1120 .await
1121 .map_err(|_| {
1122 BrokerError::Transport("AWS catalog failed (details redacted)".into())
1123 })
1124 .map(|xs| {
1125 xs.into_iter()
1126 .map(|x| CloudCatalogEntry {
1127 provider,
1128 id: x.id,
1129 display_name: x.display_name,
1130 context_ref: opaque_context.clone(),
1131 context_label: label.clone(),
1132 stale: false,
1133 fetched_at,
1134 })
1135 .collect()
1136 })
1137 }
1138 CloudProviderId::AzureOpenAi => {
1139 let mut state = self.state(provider)?;
1140 let config: super::cloud::AzureOpenAiConfig =
1141 serde_json::from_value(state["config"].clone()).map_err(|_| {
1142 BrokerError::Credential("invalid Azure broker state".into())
1143 })?;
1144 let token = self
1145 .azure_request(&mut state, super::azure_openai::AzureAudience::Arm)
1146 .await?;
1147 let mut discovery =
1148 super::azure_openai::DeploymentDiscovery::new(config.clone(), 20, 1000);
1149 let mut url = Some(discovery.initial_url());
1150 while let Some(next) = url {
1151 let r = self
1152 .http
1153 .get(next)
1154 .bearer_auth(&token)
1155 .send()
1156 .await
1157 .map_err(|_| BrokerError::Transport("Azure catalog failed".into()))?;
1158 if !r.status().is_success() {
1159 return Err(BrokerError::Transport("Azure catalog rejected".into()));
1160 }
1161 let body = read_body_capped(r, MAX_CLOUD_CATALOG_BODY_BYTES)
1162 .await
1163 .map_err(|_| {
1164 BrokerError::Transport("Azure catalog exceeded body limit".into())
1165 })?;
1166 url = discovery
1167 .accept_page(&body)
1168 .map_err(|_| BrokerError::Transport("invalid Azure catalog".into()))?;
1169 }
1170 discovery
1171 .finish()
1172 .map_err(|_| BrokerError::Transport("Azure catalog unavailable".into()))
1173 .map(|xs| {
1174 xs.into_iter()
1175 .map(|x| CloudCatalogEntry {
1176 provider,
1177 id: x.id,
1178 display_name: x.display_name,
1179 context_ref: opaque_context.clone(),
1180 context_label: format!(
1181 "{}/{}",
1182 config.resource_name, config.resource_group
1183 ),
1184 stale: false,
1185 fetched_at,
1186 })
1187 .collect()
1188 })
1189 }
1190 CloudProviderId::GoogleVertex => {
1191 let mut state = self.state(provider)?;
1192 let token = self.vertex_request(&mut state).await?;
1193 let config: super::cloud::GoogleVertexConfig =
1194 serde_json::from_value(state["config"].clone()).map_err(|_| {
1195 BrokerError::Credential("invalid Vertex broker state".into())
1196 })?;
1197 let mut page: Option<String> = None;
1198 let mut out = Vec::new();
1199 let mut seen = std::collections::HashSet::new();
1200 for _ in 0..20 {
1201 let mut url=format!("https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models",config.location,config.project_id,config.location);
1202 if let Some(p) = &page {
1203 let mut parsed = url::Url::parse(&url).map_err(|_| {
1204 BrokerError::Transport("invalid Vertex catalog URL".into())
1205 })?;
1206 parsed.query_pairs_mut().append_pair("pageToken", p);
1207 url = parsed.into();
1208 }
1209 let r = self
1210 .http
1211 .get(url)
1212 .bearer_auth(&token)
1213 .send()
1214 .await
1215 .map_err(|_| BrokerError::Transport("Vertex catalog failed".into()))?;
1216 if !r.status().is_success() {
1217 return Err(BrokerError::Transport("Vertex catalog rejected".into()));
1218 }
1219 let v = cloud_catalog_json(r).await?;
1220 for m in v["publisherModels"].as_array().into_iter().flatten() {
1221 if let Some(name) = m["name"]
1222 .as_str()
1223 .filter(|n| n.starts_with("publishers/google/models/"))
1224 {
1225 if out.len() >= MAX_CLOUD_CATALOG_ENTRIES {
1226 return Err(BrokerError::Transport(
1227 "Vertex catalog exceeded entry limit".into(),
1228 ));
1229 }
1230 if !out.iter().any(|entry: &CloudCatalogEntry| {
1231 entry.id == format!("google-vertex/{name}")
1232 }) {
1233 out.push(CloudCatalogEntry {
1234 provider,
1235 id: format!("google-vertex/{name}"),
1236 display_name: m["displayName"].as_str().unwrap_or(name).into(),
1237 context_ref: opaque_context.clone(),
1238 context_label: format!(
1239 "{}/{}",
1240 config.project_id, config.location
1241 ),
1242 stale: false,
1243 fetched_at,
1244 });
1245 }
1246 }
1247 }
1248 page = v["nextPageToken"].as_str().map(str::to_owned);
1249 match &page {
1250 Some(p) if seen.insert(p.clone()) => {}
1251 Some(_) => {
1252 return Err(BrokerError::Transport("Vertex pagination loop".into()))
1253 }
1254 None => break,
1255 }
1256 }
1257 if out.is_empty() {
1258 Err(BrokerError::Transport("Vertex catalog is empty".into()))
1259 } else {
1260 Ok(out)
1261 }
1262 }
1263 }
1264 }
1265 async fn invoke(
1266 &self,
1267 provider: CloudProviderId,
1268 context_ref: &str,
1269 model_id: &str,
1270 request: InvokeRequest,
1271 ) -> Result<CloudEventStream, BrokerError> {
1272 self.validate_context(provider, context_ref)?;
1273 if !request.tools.is_empty() {
1274 return Err(BrokerError::Denied(
1275 "tools are not yet supported by cloud providers".into(),
1276 ));
1277 }
1278 let authorized = self.catalog(provider, context_ref, false).await?;
1279 if !authorized.iter().any(|entry| entry.id == model_id) {
1280 return Err(BrokerError::Denied(
1281 "model is not present in the current provider catalog".into(),
1282 ));
1283 }
1284 match provider {
1285 CloudProviderId::AwsBedrock => {
1286 let broker = self.aws().await?;
1287 if request.stream {
1288 use futures::StreamExt;
1289 let stream = broker
1290 .converse_stream(model_id, request)
1291 .await
1292 .map_err(|_| {
1293 BrokerError::Transport(
1294 "AWS invocation failed (details redacted)".into(),
1295 )
1296 })?
1297 .map(|event| {
1298 event
1299 .map_err(|_| {
1300 BrokerError::Transport(
1301 "AWS stream failed (details redacted)".into(),
1302 )
1303 })
1304 .map(|e| match e {
1305 super::aws_bedrock::ConverseEvent::TextDelta(delta) => {
1306 CloudEvent::TextDelta { delta }
1307 }
1308 super::aws_bedrock::ConverseEvent::ToolArguments {
1309 id,
1310 delta,
1311 } => CloudEvent::ToolArguments {
1312 id,
1313 name: None,
1314 delta,
1315 },
1316 super::aws_bedrock::ConverseEvent::Usage(u) => {
1317 CloudEvent::Usage {
1318 input_tokens: u.input_tokens,
1319 output_tokens: u.output_tokens,
1320 }
1321 }
1322 super::aws_bedrock::ConverseEvent::Done => CloudEvent::Done,
1323 })
1324 });
1325 Ok(Box::pin(stream))
1326 } else {
1327 let o = broker.converse(model_id, request).await.map_err(|_| {
1328 BrokerError::Transport("AWS invocation failed (details redacted)".into())
1329 })?;
1330 Ok(Box::pin(futures::stream::iter(vec![
1331 Ok(CloudEvent::TextDelta { delta: o.text }),
1332 Ok(CloudEvent::Usage {
1333 input_tokens: o.usage.input_tokens,
1334 output_tokens: o.usage.output_tokens,
1335 }),
1336 Ok(CloudEvent::Done),
1337 ])))
1338 }
1339 }
1340 CloudProviderId::AzureOpenAi => {
1341 let mut state = self.state(provider)?;
1342 let token = self
1343 .azure_request(&mut state, super::azure_openai::AzureAudience::Inference)
1344 .await?;
1345 let endpoint = super::azure_openai::AzureEndpoint::parse(
1346 state["endpoint"]
1347 .as_str()
1348 .ok_or_else(|| BrokerError::Credential("invalid Azure endpoint".into()))?,
1349 )
1350 .map_err(|_| BrokerError::Credential("invalid Azure endpoint".into()))?;
1351 let deployment = model_id
1352 .strip_prefix("azure-openai/")
1353 .ok_or_else(|| BrokerError::Denied("invalid Azure model".into()))?;
1354 let body = serde_json::json!({"input":request.messages.into_iter().map(|m|serde_json::json!({"role":match m.role{super::cloud::MessageRole::Assistant=>"assistant",super::cloud::MessageRole::System=>"system",_=>"user"},"content":m.content})).collect::<Vec<_>>(),"stream":request.stream});
1355 let rr = super::azure_openai::responses_request(
1356 &endpoint,
1357 deployment,
1358 &serde_json::to_vec(&body).unwrap(),
1359 )
1360 .map_err(|_| BrokerError::Denied("invalid Azure request".into()))?;
1361 let r = self
1362 .http
1363 .post(rr.url)
1364 .bearer_auth(token)
1365 .header("content-type", "application/json")
1366 .body(rr.body)
1367 .send()
1368 .await
1369 .map_err(|_| BrokerError::Transport("Azure invocation failed".into()))?;
1370 if !r.status().is_success() {
1371 return Err(BrokerError::Transport("Azure invocation rejected".into()));
1372 }
1373 if request.stream {
1374 use futures::StreamExt;
1375 use std::sync::{
1376 atomic::{AtomicBool, Ordering},
1377 Arc,
1378 };
1379 let terminal = Arc::new(AtomicBool::new(false));
1380 let seen = terminal.clone();
1381 let events = sse_json_stream(r).filter_map(move |item| {
1382 let seen = seen.clone();
1383 async move {
1384 match item {
1385 Ok(Some(v)) => v["delta"]
1386 .as_str()
1387 .or_else(|| v["text"].as_str())
1388 .map(|delta| {
1389 Ok(CloudEvent::TextDelta {
1390 delta: delta.into(),
1391 })
1392 }),
1393 Ok(None) => {
1394 seen.store(true, Ordering::SeqCst);
1395 Some(Ok(CloudEvent::Done))
1396 }
1397 Err(error) => Some(Err(error)),
1398 }
1399 }
1400 });
1401 let stream = events.chain(futures::stream::unfold(
1402 (terminal, false),
1403 |(terminal, emitted)| async move {
1404 if terminal.load(Ordering::SeqCst) || emitted {
1405 None
1406 } else {
1407 Some((
1408 Err(BrokerError::Transport(
1409 "Azure stream ended without [DONE]".into(),
1410 )),
1411 (terminal, true),
1412 ))
1413 }
1414 },
1415 ));
1416 Ok(Box::pin(stream))
1417 } else {
1418 let text = read_body_capped(r, MAX_CLOUD_STREAM_EVENT_BYTES).await?;
1419 let v: serde_json::Value = serde_json::from_str(&text)
1420 .map_err(|_| BrokerError::Transport("invalid Azure response".into()))?;
1421 let delta = v["output_text"]
1422 .as_str()
1423 .ok_or_else(|| {
1424 BrokerError::Transport("Azure response omitted output".into())
1425 })?
1426 .to_owned();
1427 Ok(Box::pin(futures::stream::iter(vec![
1428 Ok(CloudEvent::TextDelta { delta }),
1429 Ok(CloudEvent::Done),
1430 ])))
1431 }
1432 }
1433 CloudProviderId::GoogleVertex => {
1434 let mut state = self.state(provider)?;
1435 let token = self.vertex_request(&mut state).await?;
1436 let config: super::cloud::GoogleVertexConfig =
1437 serde_json::from_value(state["config"].clone())
1438 .map_err(|_| BrokerError::Credential("invalid Vertex state".into()))?;
1439 let model = model_id
1440 .strip_prefix("google-vertex/publishers/google/models/")
1441 .filter(|m| {
1442 m.bytes()
1443 .all(|b| b.is_ascii_alphanumeric() || b"-._".contains(&b))
1444 })
1445 .ok_or_else(|| BrokerError::Denied("invalid Vertex model".into()))?;
1446 let action = if request.stream {
1447 "streamGenerateContent?alt=sse"
1448 } else {
1449 "generateContent"
1450 };
1451 let url=format!("https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:{}",config.location,config.project_id,config.location,model,action);
1452 let body = serde_json::json!({"contents":request.messages.into_iter().map(|m|serde_json::json!({"role":if matches!(m.role,super::cloud::MessageRole::Assistant){"model"}else{"user"},"parts":[{"text":m.content}]})).collect::<Vec<_>>()});
1453 let r = self
1454 .http
1455 .post(url)
1456 .bearer_auth(token)
1457 .json(&body)
1458 .send()
1459 .await
1460 .map_err(|_| BrokerError::Transport("Vertex invocation failed".into()))?;
1461 if !r.status().is_success() {
1462 return Err(BrokerError::Transport("Vertex invocation rejected".into()));
1463 }
1464 if request.stream {
1465 use futures::StreamExt;
1466 let events = sse_json_stream(r).map(|item| {
1467 item.and_then(|value| {
1468 value.ok_or_else(|| {
1469 BrokerError::Transport(
1470 "Vertex stream terminated without a provider event".into(),
1471 )
1472 })
1473 })
1474 });
1475 use std::sync::{
1476 atomic::{AtomicBool, Ordering},
1477 Arc,
1478 };
1479 let terminal = Arc::new(AtomicBool::new(false));
1480 let terminal_in_events = terminal.clone();
1481 let stream = events
1482 .flat_map(move |item| {
1483 let mut out = Vec::new();
1484 match item {
1485 Ok(v) => {
1486 if terminal_in_events.load(Ordering::SeqCst) {
1487 out.push(Err(BrokerError::Transport(
1488 "Vertex emitted data after terminal metadata".into(),
1489 )));
1490 return futures::stream::iter(out);
1491 }
1492 for c in v["candidates"].as_array().into_iter().flatten() {
1493 for p in
1494 c["content"]["parts"].as_array().into_iter().flatten()
1495 {
1496 if let Some(delta) = p["text"].as_str() {
1497 out.push(Ok(CloudEvent::TextDelta {
1498 delta: delta.into(),
1499 }));
1500 }
1501 }
1502 if let Some(reason) = c["finishReason"].as_str() {
1503 const VALID: &[&str] = &[
1504 "STOP",
1505 "MAX_TOKENS",
1506 "SAFETY",
1507 "RECITATION",
1508 "OTHER",
1509 "BLOCKLIST",
1510 "PROHIBITED_CONTENT",
1511 "SPII",
1512 "MALFORMED_FUNCTION_CALL",
1513 ];
1514 if !VALID.contains(&reason) {
1515 out.push(Err(BrokerError::Transport(
1516 "Vertex returned an invalid finish reason"
1517 .into(),
1518 )));
1519 } else if terminal_in_events
1520 .swap(true, Ordering::SeqCst)
1521 {
1522 out.push(Err(BrokerError::Transport(
1523 "Vertex returned duplicate terminal metadata"
1524 .into(),
1525 )));
1526 } else {
1527 out.push(Ok(CloudEvent::Done));
1528 }
1529 }
1530 }
1531 if let Some(u) = v.get("usageMetadata") {
1532 out.push(Ok(CloudEvent::Usage {
1533 input_tokens: u["promptTokenCount"]
1534 .as_u64()
1535 .unwrap_or(0),
1536 output_tokens: u["candidatesTokenCount"]
1537 .as_u64()
1538 .unwrap_or(0),
1539 }));
1540 }
1541 }
1542 Err(error) => out.push(Err(error)),
1543 }
1544 futures::stream::iter(out)
1545 })
1546 .chain(futures::stream::unfold(
1547 (terminal, false),
1548 |(terminal, emitted)| async move {
1549 if terminal.load(Ordering::SeqCst) || emitted {
1550 None
1551 } else {
1552 Some((
1553 Err(BrokerError::Transport(
1554 "Vertex stream ended without finish metadata".into(),
1555 )),
1556 (terminal, true),
1557 ))
1558 }
1559 },
1560 ));
1561 Ok(Box::pin(stream))
1562 } else {
1563 let text = read_body_capped(r, MAX_CLOUD_STREAM_EVENT_BYTES).await?;
1564 let v: serde_json::Value = serde_json::from_str(&text)
1565 .map_err(|_| BrokerError::Transport("invalid Vertex response".into()))?;
1566 let mut out = Vec::new();
1567 for c in v["candidates"].as_array().into_iter().flatten() {
1568 for p in c["content"]["parts"].as_array().into_iter().flatten() {
1569 if let Some(delta) = p["text"].as_str() {
1570 out.push(Ok(CloudEvent::TextDelta {
1571 delta: delta.into(),
1572 }));
1573 }
1574 }
1575 }
1576 out.push(Ok(CloudEvent::Done));
1577 Ok(Box::pin(futures::stream::iter(out)))
1578 }
1579 }
1580 }
1581 }
1582}
1583
1584const MAX_CLOUD_CATALOG_BODY_BYTES: usize = 2 * 1024 * 1024;
1585const MAX_CLOUD_CATALOG_ENTRIES: usize = 10_000;
1586const MAX_CLOUD_STREAM_EVENT_BYTES: usize = 1024 * 1024;
1587
1588fn sse_json_stream(
1592 resp: reqwest::Response,
1593) -> Pin<Box<dyn Stream<Item = Result<Option<serde_json::Value>, BrokerError>> + Send>> {
1594 use futures::StreamExt;
1595 let chunks = Box::pin(resp.bytes_stream());
1596 Box::pin(futures::stream::unfold(
1597 (chunks, Vec::<u8>::new(), false),
1598 |(mut chunks, mut buf, done)| async move {
1599 if done {
1600 return None;
1601 }
1602 loop {
1603 if let Some(end) = buf.windows(2).position(|w| w == b"\n\n") {
1604 let frame: Vec<u8> = buf.drain(..end + 2).collect();
1605 let text = match std::str::from_utf8(&frame[..end]) {
1606 Ok(v) => v,
1607 Err(_) => {
1608 return Some((
1609 Err(BrokerError::Transport("invalid SSE encoding".into())),
1610 (chunks, buf, true),
1611 ))
1612 }
1613 };
1614 let data = text
1615 .lines()
1616 .filter_map(|line| line.strip_prefix("data:"))
1617 .map(str::trim_start)
1618 .collect::<Vec<_>>()
1619 .join("\n");
1620 if data.is_empty() {
1621 continue;
1622 }
1623 if data == "[DONE]" {
1624 return Some((Ok(None), (chunks, buf, true)));
1625 }
1626 let value = serde_json::from_str(&data)
1627 .map_err(|_| BrokerError::Transport("invalid SSE event".into()));
1628 let terminal = value.is_err();
1629 return Some((value.map(Some), (chunks, buf, terminal)));
1630 }
1631 match chunks.next().await {
1632 Some(Ok(chunk)) if buf.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES => {
1633 buf.extend_from_slice(&chunk)
1634 }
1635 Some(_) => {
1636 return Some((
1637 Err(BrokerError::Transport(
1638 "cloud stream event exceeded limit".into(),
1639 )),
1640 (chunks, buf, true),
1641 ))
1642 }
1643 None if buf.is_empty() => return None,
1644 None => {
1645 return Some((
1646 Err(BrokerError::Transport("truncated SSE event".into())),
1647 (chunks, buf, true),
1648 ))
1649 }
1650 }
1651 }
1652 },
1653 ))
1654}
1655
1656async fn cloud_catalog_json(resp: reqwest::Response) -> Result<serde_json::Value, BrokerError> {
1658 let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
1659 serde_json::from_str(&body)
1660 .map_err(|_| BrokerError::Transport("invalid cloud catalog response".into()))
1661}
1662
1663pub struct LocalBroker {
1667 http: reqwest::Client,
1668 local_base_url: Option<String>,
1670 anthropic_usage_url: Option<String>,
1672 google_gemini_base_url: Option<String>,
1674 openai_codex_base_url: Option<String>,
1676 muse_base_url: Option<String>,
1678 refresh_endpoints: super::provider::RefreshEndpoints,
1680 stream_idle_timeout: Duration,
1682 request_timeout: Duration,
1684 max_response_bytes: usize,
1686 cloud_backend: Option<Arc<dyn CloudBackend>>,
1687 account_policy: Option<AccountPolicy>,
1690 auto_policy: Option<super::AutoSelectionPolicy>,
1691 auto_policy_cache: std::sync::Mutex<
1692 Option<(
1693 std::path::PathBuf,
1694 Option<std::time::SystemTime>,
1695 super::AutoSelectionPolicy,
1696 )>,
1697 >,
1698 current_accounts: SharedMap<(CredentialRef, SeatIdentity)>,
1700 cooldowns: Arc<std::sync::Mutex<BTreeMap<String, Cooldown>>>,
1702 usage_endpoint_override: Option<String>,
1704 max_snapshot_age: Duration,
1706 snapshots: SharedMap<super::usage::UsageSnapshot>,
1713 usage_reads: Arc<super::usage_reads::UsageReads>,
1716 instance: u64,
1718 policy_cache: Arc<std::sync::Mutex<Option<CachedPolicy>>>,
1721 observations: SharedMap<SlotObservations>,
1725 selections: SharedMap<super::auth_state::SelectionState>,
1727}
1728
1729#[derive(Debug, Clone, Default)]
1731struct SlotObservations {
1732 last_usage: Option<super::auth_state::ObservedOutcome>,
1733 last_usage_ok_ms: Option<u64>,
1734}
1735
1736type SharedMap<T> = Arc<std::sync::Mutex<BTreeMap<String, T>>>;
1738type CachedPolicy = (Option<std::time::SystemTime>, AccountPolicy);
1740
1741#[derive(Debug, Clone, PartialEq, Eq)]
1743struct Cooldown {
1744 until_ms: u64,
1745 reason: String,
1746}
1747
1748const DEFAULT_COOLDOWN: Duration = Duration::from_secs(15 * 60);
1750const MAX_COOLDOWN: Duration = Duration::from_secs(8 * 24 * 60 * 60);
1753pub const DEFAULT_MAX_SNAPSHOT_AGE: Duration = Duration::from_secs(60);
1756
1757pub const BROKER_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
1761
1762pub const GEMINI_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
1766
1767fn next_broker_instance() -> u64 {
1769 static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
1770 NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
1771}
1772
1773impl LocalBroker {
1774 pub fn new(http: reqwest::Client) -> Self {
1775 let cloud_backend = Arc::new(ProductionCloudBackend::new(http.clone()));
1776 Self {
1777 http,
1778 local_base_url: None,
1779 anthropic_usage_url: None,
1780 google_gemini_base_url: None,
1781 openai_codex_base_url: None,
1782 muse_base_url: None,
1783 refresh_endpoints: Default::default(),
1784 stream_idle_timeout: BROKER_STREAM_IDLE_TIMEOUT,
1785 request_timeout: PROXY_REQUEST_TIMEOUT,
1786 max_response_bytes: MAX_PROXY_RESPONSE_BYTES,
1787 cloud_backend: Some(cloud_backend),
1788 account_policy: None,
1789 auto_policy: None,
1790 auto_policy_cache: std::sync::Mutex::new(None),
1791 current_accounts: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
1792 cooldowns: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
1793 usage_endpoint_override: None,
1794 max_snapshot_age: DEFAULT_MAX_SNAPSHOT_AGE,
1795 snapshots: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
1796 usage_reads: super::usage_reads::process_table(),
1797 instance: next_broker_instance(),
1798 policy_cache: Arc::new(std::sync::Mutex::new(None)),
1799 observations: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
1800 selections: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
1801 }
1802 }
1803
1804 pub fn with_cloud_backend(mut self, backend: Arc<dyn CloudBackend>) -> Self {
1805 self.cloud_backend = Some(backend);
1806 self
1807 }
1808
1809 pub fn with_account_policy(mut self, policy: AccountPolicy) -> Self {
1811 self.account_policy = Some(policy);
1812 self
1813 }
1814
1815 pub fn with_auto_policy(mut self, policy: super::AutoSelectionPolicy) -> Self {
1817 self.auto_policy = Some(policy);
1818 self
1819 }
1820
1821 fn auto_policy(&self) -> super::AutoSelectionPolicy {
1822 if let Some(policy) = &self.auto_policy {
1823 return policy.clone();
1824 }
1825 let path = crate::config::resolve_read_path("config");
1826 let mtime = std::fs::metadata(&path).and_then(|m| m.modified()).ok();
1827 let mut cache = self
1828 .auto_policy_cache
1829 .lock()
1830 .unwrap_or_else(|e| e.into_inner());
1831 if let Some((cached_path, cached_time, policy)) = cache.as_ref() {
1832 if cached_path == &path && *cached_time == mtime {
1833 return policy.clone();
1834 }
1835 }
1836 let (policy, warnings) =
1837 super::AutoSelectionPolicy::from_config_map(&crate::config::load_config().auth.auto);
1838 for warning in warnings {
1839 tracing::warn!(%warning, "automatic account policy");
1840 }
1841 *cache = Some((path, mtime, policy.clone()));
1842 policy
1843 }
1844
1845 fn current_key(provider: OAuthProviderId, model: Option<&str>) -> String {
1846 format!(
1847 "{}|{}|{}",
1848 storage::auth_file_path().display(),
1849 provider,
1850 model.unwrap_or("")
1851 )
1852 }
1853
1854 #[doc(hidden)]
1858 pub fn with_openai_codex_base_url_for_tests(mut self, base_url: impl Into<String>) -> Self {
1859 self.openai_codex_base_url = Some(base_url.into().trim_end_matches('/').to_string());
1860 self
1861 }
1862
1863 #[doc(hidden)]
1867 pub fn with_muse_base_url_for_tests(mut self, base_url: impl Into<String>) -> Self {
1868 self.muse_base_url = Some(base_url.into().trim_end_matches('/').to_string());
1869 self
1870 }
1871
1872 fn stream_idle_timeout_for(&self, provider_key: &str) -> Duration {
1874 if provider_key == "google-gemini" {
1875 GEMINI_STREAM_IDLE_TIMEOUT.max(self.stream_idle_timeout)
1876 } else {
1877 self.stream_idle_timeout
1878 }
1879 }
1880
1881 #[doc(hidden)]
1883 pub fn with_stream_idle_timeout_for_tests(mut self, idle: Duration) -> Self {
1884 self.stream_idle_timeout = idle;
1885 self
1886 }
1887
1888 #[doc(hidden)]
1893 pub fn with_muse_mint_url_for_tests(mut self, mint_url: impl Into<String>) -> Self {
1894 self.refresh_endpoints =
1895 std::mem::take(&mut self.refresh_endpoints).with_muse_mint_url_for_tests(mint_url);
1896 self
1897 }
1898
1899 #[doc(hidden)]
1903 pub fn with_anthropic_token_url_for_tests(mut self, url: impl Into<String>) -> Self {
1904 self.refresh_endpoints =
1905 std::mem::take(&mut self.refresh_endpoints).with_anthropic_token_url_for_tests(url);
1906 self
1907 }
1908
1909 pub async fn access_token_auto(
1915 &self,
1916 provider: OAuthProviderId,
1917 model: Option<&str>,
1918 ) -> Result<PinnedToken, BrokerError> {
1919 self.vend_auto(provider, model).await
1920 }
1921
1922 #[doc(hidden)]
1924 pub fn with_usage_endpoint_override(mut self, url: impl Into<String>) -> Self {
1925 self.usage_endpoint_override = Some(url.into());
1926 self
1927 }
1928
1929 pub fn with_max_snapshot_age(mut self, age: Duration) -> Self {
1931 self.max_snapshot_age = age;
1932 self
1933 }
1934
1935 fn policy(&self) -> AccountPolicy {
1939 if let Some(policy) = &self.account_policy {
1940 return policy.clone();
1941 }
1942 let mtime = std::fs::metadata(crate::config::resolve_read_path("config"))
1943 .and_then(|m| m.modified())
1944 .ok();
1945 let mut cache = self.policy_cache.lock().unwrap_or_else(|e| e.into_inner());
1946 let config_policy = match cache.as_ref() {
1947 Some((cached_mtime, policy)) if *cached_mtime == mtime => policy.clone(),
1948 _ => {
1949 let (policy, _warnings) = AccountPolicy::from_config_map(
1950 &crate::config::load_config().auth.accounts,
1951 );
1952 *cache = Some((mtime, policy.clone()));
1953 policy
1954 }
1955 };
1956 config_policy.with_env_overlay()
1957 }
1958
1959 fn cooldown_key(&self, cred: &CredentialRef) -> String {
1961 let seat = self
1962 .stored_seat(cred)
1963 .ok()
1964 .map(|s| s.as_str().to_string())
1965 .unwrap_or_else(|| cred.storage_key());
1966 format!(
1967 "{}|{}|{}",
1968 storage::auth_file_path().display(),
1969 cred.provider,
1970 seat
1971 )
1972 }
1973
1974 fn cooldown_until(&self, storage_key: &str, now_ms: u64) -> Option<u64> {
1976 let cred = CredentialRef::parse_storage_key(storage_key)?;
1977 let key = self.cooldown_key(&cred);
1978 self.cooldowns
1979 .lock()
1980 .unwrap_or_else(|e| e.into_inner())
1981 .get(&key)
1982 .filter(|c| c.until_ms > now_ms)
1983 .map(|c| c.until_ms)
1984 }
1985
1986 pub fn auth_state(
1992 &self,
1993 query: &super::auth_state::AuthStateQuery,
1994 ) -> Result<super::auth_state::AuthStateReport, BrokerError> {
1995 auth_state_report(Some(self), query)
1996 }
1997
1998 fn cooldown_state(
2000 &self,
2001 cred: &CredentialRef,
2002 now_ms: u64,
2003 ) -> Option<super::auth_state::CooldownState> {
2004 let key = self.cooldown_key(cred);
2005 self.cooldowns
2006 .lock()
2007 .unwrap_or_else(|e| e.into_inner())
2008 .get(&key)
2009 .filter(|c| c.until_ms > now_ms)
2010 .map(|c| super::auth_state::CooldownState {
2011 until_ms: c.until_ms,
2012 reason: super::auth_state::CooldownReason::from_reported(&c.reason),
2013 })
2014 }
2015
2016 fn clear_state(&self, creds: &[CredentialRef], whole_providers: &[OAuthProviderId]) {
2024 let store = storage::auth_file_path().display().to_string();
2025 {
2026 let mut cooldowns = self.cooldowns.lock().unwrap_or_else(|e| e.into_inner());
2027 for cred in creds {
2028 cooldowns.remove(&self.cooldown_key(cred));
2029 cooldowns.remove(&format!("{store}|{}|{}", cred.provider, cred.storage_key()));
2030 }
2031 for provider in whole_providers {
2032 let prefix = format!("{store}|{provider}|");
2033 cooldowns.retain(|k, _| !k.starts_with(&prefix));
2034 }
2035 }
2036 for cred in creds {
2037 self.forget_snapshots(cred);
2038 self.usage_reads.forget(&self.read_key(cred));
2039 }
2040 {
2041 let mut observations = self.observations.lock().unwrap_or_else(|e| e.into_inner());
2042 for cred in creds {
2043 observations.remove(&Self::state_key(cred));
2044 }
2045 }
2046 self.current_accounts
2047 .lock()
2048 .unwrap_or_else(|e| e.into_inner())
2049 .retain(|_, (current, _)| !creds.contains(current));
2050 let mut selections = self.selections.lock().unwrap_or_else(|e| e.into_inner());
2051 for provider in whole_providers {
2052 let prefix = format!("{store}|{provider}|");
2053 selections.retain(|k, _| !k.starts_with(&prefix));
2054 }
2055 }
2056
2057 fn check_slot_exists(&self, cred: &CredentialRef) -> Result<(), BrokerError> {
2062 if let Account::Named(label) = &cred.account {
2063 match storage::load_credential(cred) {
2064 Ok(Some(_)) => {}
2065 Ok(None) => {
2066 return Err(BrokerError::UnknownAccount {
2067 provider: cred.provider.as_str().to_string(),
2068 label: label.as_str().to_string(),
2069 })
2070 }
2071 Err(e) => return Err(BrokerError::Credential(e)),
2072 }
2073 }
2074 Ok(())
2075 }
2076
2077 async fn resolve_and_vend(
2081 &self,
2082 provider: OAuthProviderId,
2083 explicit: Option<Account>,
2084 model: Option<&str>,
2085 ) -> Result<PinnedToken, BrokerError> {
2086 let selector = match explicit {
2087 Some(account) => AccountSelector::Account(account),
2088 None => self.policy().selector(provider),
2089 };
2090 match selector {
2091 AccountSelector::Account(account) => {
2092 let credential = CredentialRef::new(provider, account);
2093 let token = self.access_token_for(&credential).await?;
2094 Ok(PinnedToken { credential, token })
2095 }
2096 AccountSelector::Auto => self.vend_auto(provider, model).await,
2097 AccountSelector::Invalid { source, reason } => {
2098 Err(BrokerError::InvalidAccount(format!("{source}: {reason}")))
2099 }
2100 }
2101 }
2102
2103 const MAX_AUTO_VEND_ROUNDS: usize = 2;
2106
2107 async fn vend_auto(
2113 &self,
2114 provider: OAuthProviderId,
2115 model: Option<&str>,
2116 ) -> Result<PinnedToken, BrokerError> {
2117 use super::token::{ensure_fresh_credential_for_seat, SeatVendError};
2118 for _ in 0..Self::MAX_AUTO_VEND_ROUNDS {
2119 let (credential, seat) = match self.select_auto(provider, model).await {
2120 Ok(choice) => choice,
2121 Err(e) => {
2122 let failure = match &e {
2123 BrokerError::NoAccountAvailable { failure, .. } => failure.clone(),
2124 _ => Some(super::quota_policy::SelectionFailure::new(
2127 super::quota_policy::SelectionFailureKind::CredentialStore,
2128 )),
2129 };
2130 self.record_selection(provider, model, None, failure, None);
2131 return Err(e);
2132 }
2133 };
2134 match ensure_fresh_credential_for_seat(
2135 &self.http,
2136 &credential,
2137 &seat,
2138 &self.refresh_endpoints,
2139 )
2140 .await
2141 {
2142 Ok(creds) => {
2143 self.current_accounts
2144 .lock()
2145 .unwrap_or_else(|e| e.into_inner())
2146 .insert(
2147 Self::current_key(provider, model),
2148 (credential.clone(), seat),
2149 );
2150 self.record_selection(provider, model, Some(&credential), None, None);
2151 return Ok(PinnedToken {
2152 credential,
2153 token: AccessToken {
2154 token: creds.access,
2155 expires: creds.expires,
2156 },
2157 })
2158 }
2159 Err(SeatVendError::SeatChanged) => {
2160 tracing::info!(credential = %credential, "slot changed after selection; re-selecting");
2161 self.forget_snapshots(&credential);
2162 continue;
2163 }
2164 Err(SeatVendError::Other(f)) => {
2167 let cause = refresh_failure(&f);
2168 tracing::warn!(
2169 credential = %credential,
2170 kind = %cause.kind,
2171 "automatic choice could not be vended: token refresh failed"
2172 );
2173 self.record_selection(provider, model, Some(&credential), None, Some(cause));
2174 return Err(BrokerError::Refresh(auto_vend_failure(&credential, f)));
2175 }
2176 }
2177 }
2178 let failure = super::quota_policy::SelectionFailure::new(
2179 super::quota_policy::SelectionFailureKind::SeatChurn,
2180 );
2181 self.record_selection(provider, model, None, Some(failure.clone()), None);
2182 Err(selection_failed(
2183 provider,
2184 model,
2185 "selected credential kept changing while its capacity was being proven".into(),
2186 failure,
2187 ))
2188 }
2189
2190 fn record_selection(
2192 &self,
2193 provider: OAuthProviderId,
2194 model: Option<&str>,
2195 chosen: Option<&CredentialRef>,
2196 failure: Option<super::quota_policy::SelectionFailure>,
2197 vend_failure: Option<super::quota_policy::ObservationFailure>,
2198 ) {
2199 let state = super::auth_state::SelectionState {
2200 provider: provider.as_str().to_string(),
2201 model: model.map(str::to_string),
2202 at_ms: crate::epoch_millis(),
2203 chosen: chosen.map(|c| c.account.label_str().to_string()),
2204 failure,
2205 vend_failure,
2206 };
2207 self.selections
2208 .lock()
2209 .unwrap_or_else(|e| e.into_inner())
2210 .insert(Self::current_key(provider, model), state);
2211 }
2212
2213 fn seat_key(cred: &CredentialRef, seat: &SeatIdentity) -> String {
2216 format!(
2217 "{}|{}|{}",
2218 storage::auth_file_path().display(),
2219 cred.storage_key(),
2220 seat.as_str()
2221 )
2222 }
2223
2224 fn stored_seat(&self, cred: &CredentialRef) -> Result<SeatIdentity, BrokerError> {
2226 match storage::load_credential(cred).map_err(BrokerError::Credential)? {
2227 Some(creds) => Ok(SeatIdentity::of(cred.provider, &creds)),
2228 None => Err(BrokerError::UnknownAccount {
2229 provider: cred.provider.as_str().to_string(),
2230 label: cred.account.label_str().to_string(),
2231 }),
2232 }
2233 }
2234
2235 fn cached_snapshot(&self, key: &str) -> Option<super::usage::UsageSnapshot> {
2236 let max_age_ms = self.max_snapshot_age.as_millis() as u64;
2237 let now_ms = crate::epoch_millis();
2238 self.snapshots
2239 .lock()
2240 .unwrap_or_else(|e| e.into_inner())
2241 .get(key)
2242 .filter(|s| !s.is_stale(now_ms, max_age_ms))
2243 .cloned()
2244 }
2245
2246 fn read_key(&self, cred: &CredentialRef) -> String {
2249 format!(
2250 "{}|{}|{}",
2251 storage::auth_file_path().display(),
2252 cred.storage_key(),
2253 self.usage_endpoint_override.as_deref().unwrap_or("")
2254 )
2255 }
2256
2257 fn forget_snapshots(&self, cred: &CredentialRef) {
2261 let prefix = format!("|{}|", cred.storage_key());
2262 self.snapshots
2263 .lock()
2264 .unwrap_or_else(|e| e.into_inner())
2265 .retain(|k, _| !k.contains(&prefix));
2266 }
2267
2268 async fn fresh_usage(
2277 &self,
2278 cred: &CredentialRef,
2279 ) -> Result<(super::usage::UsageSnapshot, SeatIdentity), super::usage_reads::ReadError> {
2280 let seat = self
2281 .stored_seat(cred)
2282 .map_err(|e| UsageReadError::slot(e).into_read_error())?;
2283 let key = Self::seat_key(cred, &seat);
2284 if let Some(cached) = self.cached_snapshot(&key) {
2285 return Ok((cached, seat));
2286 }
2287 let (snapshot, used) = self.shared_read(cred).await?;
2288 self.forget_snapshots(cred);
2291 let used_key = Self::seat_key(cred, &used);
2292 self.snapshots
2293 .lock()
2294 .unwrap_or_else(|e| e.into_inner())
2295 .insert(used_key, snapshot.clone());
2296 Ok((snapshot, used))
2297 }
2298
2299 async fn usage_with_seat(
2303 &self,
2304 cred: &CredentialRef,
2305 ) -> Result<(super::usage::UsageSnapshot, SeatIdentity), UsageReadError> {
2306 self.shared_read(cred)
2307 .await
2308 .map_err(|e| UsageReadError::from_read(&e))
2309 }
2310
2311 fn state_key(cred: &CredentialRef) -> String {
2314 format!(
2315 "{}|{}",
2316 storage::auth_file_path().display(),
2317 cred.storage_key()
2318 )
2319 }
2320
2321 fn record_usage(
2322 &self,
2323 cred: &CredentialRef,
2324 failure: Option<super::quota_policy::ObservationFailure>,
2325 ) {
2326 let at_ms = crate::epoch_millis();
2327 let mut observations = self.observations.lock().unwrap_or_else(|e| e.into_inner());
2328 let slot = observations.entry(Self::state_key(cred)).or_default();
2329 if failure.is_none() {
2330 slot.last_usage_ok_ms = Some(at_ms);
2331 }
2332 slot.last_usage = Some(super::auth_state::ObservedOutcome { at_ms, failure });
2333 }
2334
2335 async fn shared_read(&self, cred: &CredentialRef) -> super::usage_reads::ReadResult {
2341 use super::usage_reads::{ReadError, ReadPath, USAGE_READ_TARGET};
2342 let key = self.read_key(cred);
2343 let (result, path) = self
2344 .usage_reads
2345 .read(&key, || self.network_read(cred))
2346 .await;
2347 if path != ReadPath::NegativeCacheHit {
2348 self.record_usage(cred, result.as_ref().err().map(ReadError::cause));
2349 }
2350 let outcome = match (path, &result) {
2351 (ReadPath::Network, Ok(_)) => "ok",
2352 (ReadPath::Network, Err(_)) => "failed",
2353 (ReadPath::Joined, _) => "joined",
2354 (ReadPath::NegativeCacheHit, _) => "negative_cache_hit",
2355 };
2356 let now_ms = crate::epoch_millis();
2357 let (kind, status, retry_in_s) = match &result {
2358 Ok(_) => ("none", String::new(), 0),
2359 Err(ReadError::Transient {
2360 failure,
2361 retry_at_ms,
2362 }) => (
2363 failure.kind,
2364 failure.status.map(|s| s.to_string()).unwrap_or_default(),
2365 retry_at_ms.saturating_sub(now_ms) / 1000,
2366 ),
2367 Err(ReadError::Other { cause, .. }) => (
2368 cause.kind.as_str(),
2369 cause.http_status.map(|s| s.to_string()).unwrap_or_default(),
2370 0,
2371 ),
2372 };
2373 tracing::info!(
2374 target: USAGE_READ_TARGET,
2375 broker = self.instance,
2376 account = %cred.account.label_str(),
2377 outcome = %outcome,
2378 kind = %kind,
2379 status = %status,
2380 retry_in_s,
2381 "usage read"
2382 );
2383 result
2384 }
2385
2386 async fn network_read(&self, cred: &CredentialRef) -> super::usage_reads::Fetched {
2391 use super::usage::{fetch_usage, supports_usage, UsageClient, UsageError, UsageFetchOptions};
2392 use super::usage_reads::Fetched;
2393 if !supports_usage(cred.provider) {
2394 return Fetched::Usage(UsageError::UnsupportedProvider {
2395 provider: cred.provider.as_str().to_string(),
2396 });
2397 }
2398 if let Err(e) = self.check_slot_exists(cred) {
2399 return UsageReadError::slot(e).into_fetched();
2400 }
2401 let creds = match super::token::ensure_fresh_credential_via(
2402 &self.http,
2403 cred,
2404 &self.refresh_endpoints,
2405 )
2406 .await
2407 {
2408 Ok(creds) => creds,
2409 Err(f) => return UsageReadError::refresh(f).into_fetched(),
2410 };
2411 let seat = SeatIdentity::of(cred.provider, &creds);
2412 let client = match UsageClient::new() {
2413 Ok(client) => client,
2414 Err(e) => return Fetched::Usage(e),
2415 };
2416 let opts = UsageFetchOptions {
2417 endpoint_override: self.usage_endpoint_override.clone(),
2418 ..UsageFetchOptions::default()
2419 };
2420 match fetch_usage(
2421 &client,
2422 cred.provider,
2423 cred.account.label_str(),
2424 &creds.access,
2425 &opts,
2426 )
2427 .await
2428 {
2429 Ok(snapshot) => Fetched::Ok(snapshot, seat),
2430 Err(e) => Fetched::Usage(e),
2431 }
2432 }
2433
2434 async fn capacity_for(&self, cred: &CredentialRef) -> CandidateRead {
2442 use super::quota_policy::AccountCapacity;
2443 use super::usage_reads::ReadError;
2444 match self.fresh_usage(cred).await {
2445 Ok((snapshot, seat)) => CandidateRead {
2446 capacity: capacity_from_snapshot(cred, &snapshot),
2447 seat: Some(seat),
2448 transient: None,
2449 grace: None,
2450 },
2451 Err(read_err) => {
2452 let err = UsageReadError::from_read(&read_err);
2453 if err.failure.kind.needs_action() {
2454 tracing::warn!(
2458 credential = %cred,
2459 kind = %err.failure.kind,
2460 detail = %action_detail(cred, &err),
2461 "usage read needs action"
2462 );
2463 } else {
2464 tracing::debug!(credential = %cred, cause = %err.failure, "usage read failed");
2465 }
2466 let capacity = AccountCapacity::failed(cred.clone(), err.failure);
2467 let ReadError::Transient {
2468 failure,
2469 retry_at_ms,
2470 } = read_err
2471 else {
2472 return CandidateRead {
2473 capacity,
2474 seat: None,
2475 transient: None,
2476 grace: None,
2477 };
2478 };
2479 let graced = self.stored_seat(cred).ok().and_then(|seat| {
2482 let good = self
2483 .usage_reads
2484 .last_good_for(&self.read_key(cred), &seat)?;
2485 Some((good, seat))
2486 });
2487 let (grace, seat) = match graced {
2488 Some((good, seat)) => (Some(capacity_from_snapshot(cred, &good)), Some(seat)),
2489 None => (None, None),
2490 };
2491 CandidateRead {
2492 capacity,
2493 seat,
2494 transient: Some((failure, retry_at_ms)),
2495 grace,
2496 }
2497 }
2498 }
2499 }
2500
2501 pub async fn plan(
2505 &self,
2506 provider: OAuthProviderId,
2507 model: Option<&str>,
2508 ) -> Result<super::SelectionPlan, BrokerError> {
2509 self.selection_plan(provider, model)
2510 .await
2511 .map(|(plan, ..)| plan)
2512 }
2513
2514 #[allow(clippy::type_complexity)]
2515 async fn selection_plan(
2516 &self,
2517 provider: OAuthProviderId,
2518 model: Option<&str>,
2519 ) -> Result<
2520 (
2521 super::SelectionPlan,
2522 BTreeMap<String, SeatIdentity>,
2523 BTreeMap<String, (super::usage_reads::TransientReadFailure, u64)>,
2524 ),
2525 BrokerError,
2526 > {
2527 use super::quota_policy::SelectionRequest;
2528 let summaries = storage::list_accounts(provider).map_err(BrokerError::Credential)?;
2529 let mut candidates = Vec::with_capacity(summaries.len());
2530 let mut seats = BTreeMap::new();
2531 let mut transient = BTreeMap::new();
2532 let mut grace = Vec::new();
2533 let mut seen = std::collections::BTreeSet::new();
2534 let mut exclude = Vec::new();
2535 for summary in &summaries {
2536 let Some(cred) = summary.credential_ref() else {
2537 continue;
2538 };
2539 let read = self.capacity_for(&cred).await;
2540 let (mut capacity, seat) = (read.capacity, read.seat);
2541 if let Some(failed) = read.transient {
2542 transient.insert(cred.storage_key(), failed);
2543 }
2544 require_requested_model(provider, model, &mut capacity);
2547 grace.extend(read.grace.map(|mut good| {
2548 require_requested_model(provider, model, &mut good);
2549 good
2550 }));
2551 if let Some(seat) = seat {
2552 if !seen.insert(seat.as_str().to_string()) {
2553 exclude.push(cred.clone());
2554 }
2555 seats.insert(cred.storage_key(), seat);
2556 }
2557 candidates.push(capacity);
2558 }
2559 let now_ms = crate::epoch_millis();
2560 for capacity in &mut candidates {
2561 capacity.cooldown_until_ms =
2562 self.cooldown_until(&capacity.credential.storage_key(), now_ms);
2563 }
2564 let previous = self
2565 .current_accounts
2566 .lock()
2567 .unwrap_or_else(|e| e.into_inner())
2568 .get(&Self::current_key(provider, model))
2569 .cloned();
2570 let current = previous
2571 .filter(|(cred, seat)| seats.get(&cred.storage_key()) == Some(seat))
2572 .map(|(cred, _)| cred);
2573 let policy = self.auto_policy();
2574 let preference: Vec<_> = summaries.iter().map(|s| s.label.clone()).collect();
2575 let request = SelectionRequest {
2576 model,
2577 preference: &preference,
2578 exclude: &exclude,
2579 strategy: policy.strategy_for(provider),
2580 current: current.as_ref(),
2581 urgent_horizon_ms: policy.urgent_horizon_ms,
2582 sticky: policy.sticky,
2583 grace: &grace,
2584 ..SelectionRequest::new(provider, now_ms, self.max_snapshot_age.as_millis() as u64)
2585 };
2586 Ok((
2587 super::SelectionPlan::from_candidates(&request, &candidates),
2588 seats,
2589 transient,
2590 ))
2591 }
2592
2593 async fn select_auto(
2594 &self,
2595 provider: OAuthProviderId,
2596 model: Option<&str>,
2597 ) -> Result<(CredentialRef, SeatIdentity), BrokerError> {
2598 use super::quota_policy::{Selection, SelectionFailure, SelectionFailureKind};
2599 use super::usage_reads::USAGE_READ_TARGET;
2600 let (plan, mut seats, transient) = match self.selection_plan(provider, model).await {
2601 Ok(planned) => planned,
2602 Err(e) => {
2606 log_selection_failure(
2607 provider,
2608 model,
2609 &SelectionFailure::new(SelectionFailureKind::CredentialStore),
2610 );
2611 return Err(e);
2612 }
2613 };
2614 match plan.selection {
2615 Selection::Selected {
2616 credential,
2617 tier,
2618 budget_reset_ms,
2619 utilization,
2620 ..
2621 } => {
2622 let seat = seats.remove(&credential.storage_key());
2625 debug_assert!(
2626 seat.is_some(),
2627 "selected account has no paired seat identity"
2628 );
2629 let Some(seat) = seat else {
2630 return Err(selection_failed(
2631 provider,
2632 model,
2633 "selected account has no paired seat identity".into(),
2634 SelectionFailure::new(SelectionFailureKind::SeatChurn),
2635 ));
2636 };
2637 if let Some((failure, retry_at_ms)) = transient.get(&credential.storage_key()) {
2640 tracing::info!(
2641 target: USAGE_READ_TARGET,
2642 broker = self.instance,
2643 account = %credential.account.label_str(),
2644 outcome = %"grace",
2645 kind = %failure.kind,
2646 status = %failure.status.map(|s| s.to_string()).unwrap_or_default(),
2647 retry_in_s = retry_at_ms.saturating_sub(crate::epoch_millis()) / 1000,
2648 "usage read"
2649 );
2650 }
2651 tracing::info!(%credential, ?plan.strategy, tier, ?budget_reset_ms, ?utilization, "automatic subscription choice");
2652 Ok((credential, seat))
2653 }
2654 Selection::NoCapacity {
2655 rejections,
2656 earliest_reset_ms,
2657 } => {
2658 let detail: Vec<_> = rejections
2659 .iter()
2660 .map(|r| format!("{}: {}", r.credential.account, r.describe()))
2661 .collect();
2662 let reset = earliest_reset_ms
2663 .map(|t| format!("; earliest reset at {t}"))
2664 .unwrap_or_default();
2665 let mut failure = super::quota_policy::SelectionFailure::from_rejections(
2666 &rejections,
2667 plan.now_ms,
2668 );
2669 if failure.kind == SelectionFailureKind::UsageReadFailed {
2672 failure.usage_read_retry_ms = rejections
2673 .iter()
2674 .filter_map(|r| transient.get(&r.credential.storage_key()))
2675 .map(|(_, at)| *at)
2676 .min();
2677 }
2678 let reason = if detail.is_empty() {
2679 format!("no {provider} account is logged in")
2680 } else {
2681 format!("{}{reset}", detail.join(", "))
2682 };
2683 Err(selection_failed(provider, model, reason, failure))
2684 }
2685 }
2686 }
2687
2688 pub fn with_local_base_url(http: reqwest::Client, base_url: impl Into<String>) -> Self {
2690 Self {
2691 local_base_url: Some(base_url.into()),
2692 ..Self::new(http)
2693 }
2694 }
2695
2696 #[doc(hidden)]
2700 pub fn with_google_gemini_base_url_for_tests(
2701 http: reqwest::Client,
2702 base_url: impl Into<String>,
2703 ) -> Self {
2704 Self {
2705 google_gemini_base_url: Some(base_url.into().trim_end_matches('/').to_string()),
2706 ..Self::new(http)
2707 }
2708 }
2709
2710 #[doc(hidden)]
2712 pub fn with_anthropic_usage_url(mut self, url: impl Into<String>) -> Self {
2713 self.anthropic_usage_url = Some(url.into());
2714 self
2715 }
2716
2717 #[doc(hidden)]
2719 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
2720 self.request_timeout = timeout;
2721 self
2722 }
2723
2724 #[doc(hidden)]
2726 pub fn with_max_response_bytes(mut self, cap: usize) -> Self {
2727 self.max_response_bytes = cap;
2728 self
2729 }
2730
2731 fn resolve_static_key(&self, provider: &str) -> Result<String, BrokerError> {
2735 if provider == LOCAL_PROVIDER_KEY {
2736 return Ok(self.resolve_local_key());
2737 }
2738 let spec = static_provider(provider)
2739 .ok_or_else(|| BrokerError::UnknownProvider(provider.to_string()))?;
2740 if let Ok(Some(key)) = storage::load_static_key(provider) {
2741 return Ok(key);
2742 }
2743 if let Some(key) = discover_legacy_static_key(spec) {
2744 if let Err(e) = storage::save_static_key(provider, &key) {
2747 tracing::warn!(provider, "static key migration failed: {e}");
2748 }
2749 return Ok(key);
2750 }
2751 Err(BrokerError::NotConfigured(provider.to_string()))
2752 }
2753
2754 fn resolve_local_key(&self) -> String {
2755 storage::load_static_key(LOCAL_PROVIDER_KEY)
2756 .ok()
2757 .flatten()
2758 .or_else(|| {
2759 crate::config::get_provider_keys()
2760 .get(LOCAL_PROVIDER_KEY)
2761 .filter(|s| !s.is_empty())
2762 .cloned()
2763 })
2764 .or_else(|| {
2765 std::env::var("LOCAL_API_KEY")
2766 .ok()
2767 .filter(|s| !s.is_empty())
2768 })
2769 .unwrap_or_else(|| "local".to_string())
2770 }
2771
2772 fn base_url_for(&self, provider: &str) -> Result<String, BrokerError> {
2773 if provider == LOCAL_PROVIDER_KEY {
2774 if let Some(url) = &self.local_base_url {
2775 return Ok(url.trim_end_matches('/').to_string());
2776 }
2777 return Ok(local_endpoint_url());
2778 }
2779 static_provider(provider)
2780 .map(|s| s.base_url.trim_end_matches('/').to_string())
2781 .ok_or_else(|| BrokerError::UnknownProvider(provider.to_string()))
2782 }
2783
2784 async fn send(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
2785 request.validate()?;
2786 let (provider_key, explicit_account) = request.split_provider()?;
2787 let provider_key = provider_key.as_str();
2788 let oauth_provider: Option<OAuthProviderId> = provider_key.parse().ok();
2792 let mut pinned: Option<PinnedToken> = None;
2793 let muse_env_key = (provider_key == "muse" && explicit_account.is_none())
2797 .then(super::muse::api_key_from_env)
2798 .flatten();
2799 if let Some(provider) = oauth_provider.filter(|_| muse_env_key.is_none()) {
2800 if provider_key != LOCAL_PROVIDER_KEY && static_provider(provider_key).is_none() {
2801 pinned = Some(self.resolve_and_vend(provider, explicit_account, None).await?);
2802 }
2803 }
2804 let bearer = |pinned: &Option<PinnedToken>| -> Result<String, BrokerError> {
2805 pinned
2806 .as_ref()
2807 .map(|p| p.token.token.clone())
2808 .ok_or_else(|| BrokerError::UnknownProvider(request.provider.clone()))
2809 };
2810 let (key, base) = if provider_key == "xai-auth" {
2811 (bearer(&pinned)?, "https://api.x.ai/v1".to_string())
2812 } else if provider_key == "github-copilot" {
2813 (
2816 bearer(&pinned)?,
2817 super::github_copilot_models_base_url().to_string(),
2818 )
2819 } else if provider_key == "google-gemini" {
2820 let base = self
2823 .google_gemini_base_url
2824 .clone()
2825 .unwrap_or_else(|| GOOGLE_GEMINI_CODE_ASSIST_BASE_URL.to_string());
2826 (bearer(&pinned)?, base)
2827 } else if provider_key == "openai-codex" {
2828 let base = self
2831 .openai_codex_base_url
2832 .clone()
2833 .unwrap_or_else(|| OPENAI_CODEX_BACKEND_BASE_URL.to_string());
2834 (bearer(&pinned)?, base)
2835 } else if provider_key == "kimi-code" {
2836 (bearer(&pinned)?, super::kimi_code::API_BASE_URL.to_string())
2840 } else if provider_key == "muse" {
2841 let key = match &muse_env_key {
2846 Some(key) => key.clone(),
2847 None => bearer(&pinned)?,
2848 };
2849 let base = self
2850 .muse_base_url
2851 .clone()
2852 .unwrap_or_else(|| super::muse::API_ORIGIN.to_string());
2853 (key, base)
2854 } else if provider_key == "anthropic" {
2855 (bearer(&pinned)?, "https://api.anthropic.com".to_string())
2858 } else {
2859 (
2860 self.resolve_static_key(provider_key)?,
2861 self.base_url_for(provider_key)?,
2862 )
2863 };
2864 let response = self
2865 .dispatch(request, provider_key, &key, &base, &pinned)
2866 .await?;
2867 if provider_key == "muse" && response.status() == reqwest::StatusCode::UNAUTHORIZED {
2881 if let Some(rejected) = pinned.as_ref() {
2882 tracing::info!(credential = %rejected.credential, "muse key rejected (401); re-minting once");
2883 match self.recover_rejected_token(&rejected.credential, &key).await {
2884 Ok(fresh) => {
2885 let pinned = Some(PinnedToken {
2886 credential: rejected.credential.clone(),
2887 token: fresh,
2888 });
2889 let key = bearer(&pinned)?;
2890 return self
2891 .dispatch(request, provider_key, &key, &base, &pinned)
2892 .await;
2893 }
2894 Err(BrokerError::Refresh(f))
2900 if f.kind.is_retryable()
2901 || f.kind == super::token::RefreshFailureKind::ActionRequired =>
2902 {
2903 tracing::warn!(
2904 credential = %rejected.credential,
2905 reason = %f,
2906 "muse key re-mint failed; returning a retryable error"
2907 );
2908 return Err(BrokerError::Refresh(f));
2909 }
2910 Err(e) => {
2911 tracing::warn!(
2914 credential = %rejected.credential,
2915 reason = %e,
2916 "muse key re-mint refused or rejected; returning the upstream 401"
2917 );
2918 }
2919 }
2920 }
2921 }
2922 Ok(response)
2923 }
2924
2925 async fn dispatch(
2927 &self,
2928 request: &ProxyRequest,
2929 provider_key: &str,
2930 key: &str,
2931 base: &str,
2932 pinned: &Option<PinnedToken>,
2933 ) -> Result<reqwest::Response, BrokerError> {
2934 let url = format!("{base}{}", request.path);
2935 let mut builder = if provider_key == "google-gemini"
2938 || provider_key == "openai-codex"
2939 || provider_key == "anthropic"
2940 || provider_key == "kimi-code"
2941 || provider_key == "muse"
2942 {
2943 let no_redirect = reqwest::Client::builder()
2947 .redirect(reqwest::redirect::Policy::none())
2948 .connect_timeout(Duration::from_secs(10))
2949 .read_timeout(self.stream_idle_timeout_for(provider_key))
2950 .build()
2951 .map_err(|e| BrokerError::Transport(format!("client build failed: {e}")))?;
2952 match request.method {
2953 ProxyMethod::Get => no_redirect.get(&url),
2954 ProxyMethod::Post => no_redirect.post(&url),
2955 }
2956 } else {
2957 match request.method {
2958 ProxyMethod::Get => self.http.get(&url),
2959 ProxyMethod::Post => self.http.post(&url),
2960 }
2961 };
2962 builder = builder.bearer_auth(key);
2963 if provider_key == "openai-codex" {
2964 let cred = pinned
2968 .as_ref()
2969 .map(|p| p.credential.clone())
2970 .unwrap_or_else(|| CredentialRef::default_for(OAuthProviderId::OpenAiCodex));
2971 let account_id = super::extract_codex_account_id(key)
2972 .or_else(|| {
2973 storage::load_credential(&cred)
2974 .ok()
2975 .flatten()
2976 .and_then(|c| c.account_id)
2977 .filter(|s| !s.trim().is_empty())
2978 })
2979 .ok_or_else(|| {
2980 BrokerError::Credential(
2981 "openai-codex credential is missing chatgpt account id".into(),
2982 )
2983 })?;
2984 builder = builder
2985 .header("chatgpt-account-id", account_id)
2986 .header("originator", "synaps")
2987 .header("OpenAI-Beta", "responses=experimental")
2988 .header("accept", "application/json");
2989 }
2990 if provider_key == "anthropic" {
2991 for (name, value) in anthropic_oauth_catalog_request_headers() {
2994 builder = builder.header(*name, *value);
2995 }
2996 }
2997 if provider_key == "github-copilot" {
2998 for (name, value) in super::github_copilot_models_request_headers() {
2999 builder = builder.header(*name, *value);
3000 }
3001 if request.path != "/models" {
3002 builder = builder
3003 .header("Openai-Intent", "conversation-edits")
3004 .header("X-Initiator", "agent");
3005 }
3006 }
3007 if provider_key == "google-gemini" {
3008 if request.stream && request.path == "/v1internal:streamGenerateContent" {
3012 builder = builder.query(&[("alt", "sse")]);
3013 }
3014 builder = builder.header("user-agent", "SynapsCLI/0.6.0 (google-gemini)");
3015 }
3016 if provider_key == "muse" {
3017 builder = builder.header("user-agent", super::muse::USER_AGENT);
3020 }
3021 if provider_key == "kimi-code" {
3022 let device_id = super::kimi_code::load_or_create_device_id();
3026 for (name, value) in super::kimi_code::identity_request_headers(&device_id) {
3027 builder = builder.header(name, value);
3028 }
3029 }
3030 if let Some(bytes) = &request.body_bytes {
3031 builder = builder
3035 .header("content-type", "application/json")
3036 .body(bytes.clone());
3037 } else if let Some(body) = &request.body {
3038 builder = builder.json(body);
3039 }
3040 if request.stream {
3041 builder = builder.header("accept", "text/event-stream");
3042 } else {
3043 builder = builder.timeout(self.request_timeout);
3047 }
3048 builder.send().await.map_err(|e| {
3049 if e.is_connect() && provider_key == LOCAL_PROVIDER_KEY {
3050 BrokerError::Transport(format!(
3051 "can't reach local endpoint at {url} — is Ollama/LM Studio running?"
3052 ))
3053 } else {
3054 BrokerError::Transport(format!("request to {} failed: {e}", request.provider))
3056 }
3057 })
3058 }
3059}
3060
3061pub fn capacity_from_snapshot(
3067 cred: &CredentialRef,
3068 snapshot: &super::usage::UsageSnapshot,
3069) -> super::quota_policy::AccountCapacity {
3070 use super::quota_policy::{
3071 AccountCapacity, ModelAvailability, ModelState, QuotaObservation, WindowLimit,
3072 };
3073 use super::usage::{Availability, WindowScope};
3074 let mut windows: Vec<WindowLimit> = snapshot
3075 .windows
3076 .iter()
3077 .filter_map(|w| {
3078 let models = match &w.scope {
3079 WindowScope::Account => None,
3080 WindowScope::Model { model } => Some(vec![model.clone()]),
3081 WindowScope::Feature { .. } => return None,
3082 };
3083 Some(WindowLimit {
3084 id: w.id.clone(),
3085 duration_ms: w.duration_secs.map(|s| s.saturating_mul(1000)),
3086 used_percent: w.used_percent.valid(),
3087 limit_reached: w.limit_reached,
3088 resets_at_ms: w.reset_at,
3089 models,
3090 })
3091 })
3092 .collect();
3093 if snapshot.limit_reached == Some(true) {
3094 windows.push(WindowLimit {
3096 id: "limit_reached".into(),
3097 duration_ms: None,
3098 used_percent: None,
3099 limit_reached: Some(true),
3100 resets_at_ms: snapshot.earliest_reset_at(),
3101 models: None,
3102 });
3103 }
3104 let models = if snapshot.model_availability.is_empty() {
3105 None
3106 } else {
3107 Some(
3108 snapshot
3109 .model_availability
3110 .iter()
3111 .map(|m| ModelAvailability {
3112 model: m.model.clone(),
3113 state: match m.availability {
3114 Availability::Available => ModelState::Available,
3115 Availability::Exhausted => ModelState::Exhausted,
3116 Availability::Unknown => ModelState::Unknown,
3117 },
3118 })
3119 .collect(),
3120 )
3121 };
3122 AccountCapacity {
3123 credential: cred.clone(),
3124 observed_at_ms: Some(snapshot.observed_at),
3125 observation: QuotaObservation::Ok { windows, models },
3126 cooldown_until_ms: None,
3127 failure: None,
3128 }
3129}
3130
3131#[derive(Debug)]
3134struct UsageReadError {
3135 error: BrokerError,
3136 failure: super::quota_policy::ObservationFailure,
3137}
3138
3139struct CandidateRead {
3141 capacity: super::quota_policy::AccountCapacity,
3142 seat: Option<SeatIdentity>,
3144 transient: Option<(super::usage_reads::TransientReadFailure, u64)>,
3146 grace: Option<super::quota_policy::AccountCapacity>,
3148}
3149
3150impl UsageReadError {
3151 fn from_read(e: &super::usage_reads::ReadError) -> Self {
3154 Self {
3155 error: e.to_broker_error(),
3156 failure: e.cause(),
3157 }
3158 }
3159
3160 fn into_read_error(self) -> super::usage_reads::ReadError {
3163 super::usage_reads::ReadError::Other {
3164 error: self.error,
3165 cause: self.failure,
3166 }
3167 }
3168
3169 fn into_fetched(self) -> super::usage_reads::Fetched {
3171 super::usage_reads::Fetched::Failed {
3172 error: self.error,
3173 cause: self.failure,
3174 }
3175 }
3176
3177 fn refresh(f: super::token::RefreshFailure) -> Self {
3180 Self {
3181 failure: refresh_failure(&f),
3182 error: BrokerError::Credential(f.into()),
3183 }
3184 }
3185
3186 fn slot(error: BrokerError) -> Self {
3188 use super::quota_policy::{AccountFailureKind as K, ObservationFailure};
3189 let kind = match &error {
3190 BrokerError::UnknownAccount { .. } => K::CredentialMissing,
3191 _ => K::CredentialStore,
3192 };
3193 Self {
3194 failure: ObservationFailure::new(kind),
3195 error,
3196 }
3197 }
3198
3199 #[cfg(test)]
3202 fn usage(e: super::usage::UsageError) -> Self {
3203 let table = super::usage_reads::UsageReads::default();
3204 match table.record_at("slot", super::usage_reads::Fetched::Usage(e), 0) {
3205 Err(read) => Self::from_read(&read),
3206 Ok(_) => unreachable!("a usage error is never a reading"),
3207 }
3208 }
3209}
3210
3211fn require_requested_model(
3216 provider: OAuthProviderId,
3217 model: Option<&str>,
3218 capacity: &mut super::quota_policy::AccountCapacity,
3219) {
3220 use super::quota_policy::{model_matches, ModelAvailability, ModelState, QuotaObservation};
3221 if provider != OAuthProviderId::OpenAiCodex {
3222 return;
3223 }
3224 if let (Some(model), QuotaObservation::Ok { models, .. }) = (model, &mut capacity.observation) {
3225 let availability = models.get_or_insert_with(Vec::new);
3226 if !availability.iter().any(|m| model_matches(&m.model, model)) {
3227 availability.push(ModelAvailability {
3228 model: model.to_string(),
3229 state: ModelState::Unknown,
3230 });
3231 }
3232 }
3233}
3234
3235fn action_detail(cred: &CredentialRef, err: &UsageReadError) -> String {
3240 use super::quota_policy::AccountFailureKind as K;
3241 match err.failure.kind {
3242 K::CredentialStore => format!(
3243 "the stored login of '{}' could not be read from {}; repair or restore the \
3244 credential store",
3245 cred.account.label_str(),
3246 storage::auth_file_path().display()
3247 ),
3248 _ => err.error.to_string(),
3249 }
3250}
3251
3252fn refresh_failure(f: &super::token::RefreshFailure) -> super::quota_policy::ObservationFailure {
3255 let failure = refresh_kind_failure(f.kind);
3256 match f.http_status() {
3257 Some(status) => failure.with_status(status),
3258 None => failure,
3259 }
3260}
3261
3262fn refresh_kind_failure(
3264 kind: super::token::RefreshFailureKind,
3265) -> super::quota_policy::ObservationFailure {
3266 use super::quota_policy::{AccountFailureKind as K, ObservationFailure};
3267 use super::token::RefreshFailureKind as R;
3268 match kind {
3269 R::Transient => ObservationFailure::new(K::RefreshTransient),
3270 R::LockBusy => ObservationFailure::new(K::RefreshLockBusy),
3271 R::Other => ObservationFailure::new(K::RefreshOther),
3272 R::Cooldown { retry_in_secs } => {
3273 ObservationFailure::new(K::RefreshOther).with_retry_after(Some(retry_in_secs))
3274 }
3275 R::Credential => ObservationFailure::new(K::RefreshRejected),
3276 R::ActionRequired => ObservationFailure::new(K::RefreshActionRequired),
3277 }
3278}
3279
3280fn auto_vend_failure(
3285 credential: &CredentialRef,
3286 f: super::token::RefreshFailure,
3287) -> super::token::RefreshFailure {
3288 let label = credential.account.label_str();
3289 let mut message = format!("account '{label}': {}", f.message());
3290 if f.kind == super::token::RefreshFailureKind::Credential && !message.contains("synaps login") {
3291 let account = if credential.account.is_default() {
3292 String::new()
3293 } else {
3294 format!(" --account {label}")
3295 };
3296 message.push_str(&format!(
3297 "; log it in again with `synaps login --provider {}{account}`",
3298 credential.provider
3299 ));
3300 }
3301 super::token::RefreshFailure::new(f.kind, message).with_http_status(f.http_status())
3302}
3303
3304fn selection_failed(
3307 provider: OAuthProviderId,
3308 model: Option<&str>,
3309 reason: String,
3310 failure: super::quota_policy::SelectionFailure,
3311) -> BrokerError {
3312 log_selection_failure(provider, model, &failure);
3313 BrokerError::NoAccountAvailable {
3314 provider: provider.to_string(),
3315 reason,
3316 failure: Some(failure),
3317 }
3318}
3319
3320fn log_selection_failure(
3324 provider: OAuthProviderId,
3325 model: Option<&str>,
3326 failure: &super::quota_policy::SelectionFailure,
3327) {
3328 tracing::warn!(
3329 provider = %provider,
3330 model = model.unwrap_or("any"),
3331 kind = %failure.kind,
3332 accounts = %failure.accounts_field(),
3333 earliest_reset_ms = ?failure.earliest_reset_ms,
3334 "auth_selection_failed"
3335 );
3336}
3337
3338async fn read_body_capped(resp: reqwest::Response, cap: usize) -> Result<String, BrokerError> {
3343 use futures::StreamExt;
3344 let mut buf: Vec<u8> = Vec::new();
3345 let mut stream = resp.bytes_stream();
3346 while let Some(chunk) = stream.next().await {
3347 let chunk =
3348 chunk.map_err(|e| BrokerError::Transport(format!("failed to read response: {e}")))?;
3349 if buf.len() + chunk.len() > cap {
3350 return Err(BrokerError::Transport(format!(
3351 "response body exceeded the {cap}-byte broker buffering limit"
3352 )));
3353 }
3354 buf.extend_from_slice(&chunk);
3355 }
3356 String::from_utf8(buf)
3357 .map_err(|_| BrokerError::Transport("response body was not valid UTF-8".into()))
3358}
3359
3360const MAX_PROXY_ERROR_CLASSIFY_BYTES: usize = 8 * 1024;
3364
3365async fn read_error_body_for_classification(resp: reqwest::Response, cap: usize) -> Option<String> {
3370 use futures::StreamExt;
3371 let mut buf: Vec<u8> = Vec::with_capacity(cap.min(1024));
3372 let mut stream = resp.bytes_stream();
3373 while let Some(chunk) = stream.next().await {
3374 let chunk = chunk.ok()?;
3375 let room = cap.saturating_sub(buf.len());
3376 if chunk.len() >= room {
3377 buf.extend_from_slice(&chunk[..room]);
3378 return String::from_utf8(buf).ok();
3381 }
3382 buf.extend_from_slice(&chunk);
3383 }
3384 String::from_utf8(buf).ok()
3385}
3386
3387fn discover_legacy_static_key(spec: &StaticProviderSpec) -> Option<String> {
3390 if let Some(v) = crate::config::get_provider_keys().get(spec.key) {
3391 if !v.is_empty() {
3392 return Some(v.clone());
3393 }
3394 }
3395 spec.env_vars
3396 .iter()
3397 .find_map(|var| std::env::var(var).ok().filter(|v| !v.is_empty()))
3398}
3399
3400pub fn static_key_configured(provider: &str) -> bool {
3403 if provider == LOCAL_PROVIDER_KEY {
3404 return true; }
3406 let Some(spec) = static_provider(provider) else {
3407 return false;
3408 };
3409 storage::load_static_key(provider).ok().flatten().is_some()
3410 || discover_legacy_static_key(spec).is_some()
3411}
3412
3413pub fn static_key_status(provider: &str) -> StaticKeyStatus {
3417 let stored = storage::load_static_key(provider)
3418 .ok()
3419 .flatten()
3420 .or_else(|| {
3421 crate::config::get_provider_keys()
3422 .get(provider)
3423 .filter(|s| !s.is_empty())
3424 .cloned()
3425 });
3426 if let Some(key) = stored {
3427 return StaticKeyStatus::Configured {
3428 masked: mask_key(&key),
3429 };
3430 }
3431 if let Some(spec) = static_provider(provider) {
3432 if spec
3433 .env_vars
3434 .iter()
3435 .any(|var| std::env::var(var).is_ok_and(|v| !v.is_empty()))
3436 {
3437 return StaticKeyStatus::FromEnv;
3438 }
3439 }
3440 StaticKeyStatus::NotSet
3441}
3442
3443fn mask_key(key: &str) -> String {
3445 let chars: Vec<char> = key.chars().collect();
3446 if chars.len() <= 8 {
3447 return "…".to_string();
3448 }
3449 let head: String = chars[..4].iter().collect();
3450 let tail: String = chars[chars.len() - 4..].iter().collect();
3451 format!("{head}…{tail}")
3452}
3453
3454pub fn configured_static_provider_keys() -> std::collections::BTreeSet<String> {
3456 STATIC_PROVIDERS
3457 .iter()
3458 .filter(|s| static_key_configured(s.key))
3459 .map(|s| s.key.to_string())
3460 .collect()
3461}
3462
3463pub fn local_endpoint_url() -> String {
3466 local_endpoint_config()
3467 .unwrap_or_else(|| LOCAL_DEFAULT_BASE_URL.to_string())
3468 .trim_end_matches('/')
3469 .to_string()
3470}
3471
3472pub fn local_endpoint_config() -> Option<String> {
3474 crate::config::get_provider_keys()
3475 .get("local.url")
3476 .filter(|s| !s.is_empty())
3477 .cloned()
3478 .or_else(|| {
3479 std::env::var("LOCAL_ENDPOINT")
3480 .ok()
3481 .filter(|s| !s.is_empty())
3482 })
3483}
3484
3485pub fn local_model_ids() -> Vec<String> {
3487 crate::config::get_provider_keys()
3488 .get("local.models")
3489 .map(|value| {
3490 value
3491 .split(',')
3492 .map(|s| s.trim().to_string())
3493 .filter(|s| !s.is_empty())
3494 .collect()
3495 })
3496 .unwrap_or_default()
3497}
3498
3499pub fn anthropic_credential_available() -> bool {
3503 oauth_provider_logged_in(OAuthProviderId::Anthropic)
3504 || std::env::var("ANTHROPIC_API_KEY").is_ok_and(|v| !v.is_empty())
3505}
3506
3507pub fn oauth_provider_logged_in(provider: OAuthProviderId) -> bool {
3518 oauth_provider_logged_in_at(
3519 &storage::auth_file_path(),
3520 &AccountPolicy::from_environment(),
3521 provider,
3522 )
3523}
3524
3525pub(crate) fn oauth_provider_logged_in_at(
3527 path: &std::path::Path,
3528 policy: &AccountPolicy,
3529 provider: OAuthProviderId,
3530) -> bool {
3531 let Ok(inventory) = storage::list_accounts_detailed_at(path, Some(provider)) else {
3532 return false;
3533 };
3534 match policy.selector(provider) {
3535 AccountSelector::Account(Account::Named(label)) => inventory
3536 .accounts
3537 .iter()
3538 .any(|row| row.label == label.as_str()),
3539 AccountSelector::Account(Account::Default)
3540 | AccountSelector::Auto
3541 | AccountSelector::Invalid { .. } => !inventory.accounts.is_empty(),
3542 }
3543}
3544
3545#[async_trait]
3546impl CredentialBroker for LocalBroker {
3547 async fn recover_rejected_token(
3548 &self,
3549 cred: &CredentialRef,
3550 rejected: &str,
3551 ) -> Result<AccessToken, BrokerError> {
3552 self.check_slot_exists(cred)?;
3553 let creds = super::token::recover_rejected_credential(
3554 &self.http,
3555 cred,
3556 rejected,
3557 &self.refresh_endpoints,
3558 )
3559 .await
3560 .map_err(BrokerError::Refresh)?;
3561 Ok(AccessToken {
3562 token: creds.access,
3563 expires: creds.expires,
3564 })
3565 }
3566
3567 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
3568 Ok(self.access_token_pinned(provider).await?.token)
3569 }
3570
3571 fn account_selector(&self, provider: OAuthProviderId) -> AccountSelector {
3572 self.policy().selector(provider)
3573 }
3574
3575 async fn access_token_for(&self, cred: &CredentialRef) -> Result<AccessToken, BrokerError> {
3576 self.check_slot_exists(cred)?;
3577 let creds =
3578 super::token::ensure_fresh_credential_via(&self.http, cred, &self.refresh_endpoints)
3579 .await
3580 .map_err(|f| match f.kind {
3581 super::token::RefreshFailureKind::ActionRequired => BrokerError::Refresh(f),
3584 _ => BrokerError::Credential(f.into()),
3585 })?;
3586 Ok(AccessToken {
3588 token: creds.access,
3589 expires: creds.expires,
3590 })
3591 }
3592
3593 async fn access_token_pinned(
3594 &self,
3595 provider: OAuthProviderId,
3596 ) -> Result<PinnedToken, BrokerError> {
3597 self.access_token_pinned_for(provider, None).await
3598 }
3599
3600 async fn access_token_pinned_for(
3601 &self,
3602 provider: OAuthProviderId,
3603 model: Option<&str>,
3604 ) -> Result<PinnedToken, BrokerError> {
3605 self.resolve_and_vend(provider, None, model).await
3606 }
3607
3608 async fn accounts(&self, provider: OAuthProviderId) -> Result<Vec<AccountSummary>, BrokerError> {
3609 let mut rows = storage::list_accounts(provider).map_err(BrokerError::Credential)?;
3610 let selector = self.policy().selector(provider);
3611 let now_ms = crate::epoch_millis();
3612 for row in &mut rows {
3613 row.selected = matches!(&selector, AccountSelector::Account(a) if a.label_str() == row.label);
3614 if let Some(cred) = row.credential_ref() {
3615 row.cooldown_until = self.cooldown_until(&cred.storage_key(), now_ms);
3616 }
3617 }
3618 Ok(rows)
3619 }
3620
3621 async fn usage(&self, cred: &CredentialRef) -> Result<super::usage::UsageSnapshot, BrokerError> {
3622 self.usage_with_seat(cred)
3623 .await
3624 .map(|(snapshot, _seat)| snapshot)
3625 .map_err(|e| e.error)
3626 }
3627
3628 async fn report_cooldown(
3629 &self,
3630 cred: &CredentialRef,
3631 until_ms: Option<u64>,
3632 reason: &str,
3633 ) -> Result<(), BrokerError> {
3634 let now_ms = crate::epoch_millis();
3635 let max_until = now_ms.saturating_add(MAX_COOLDOWN.as_millis() as u64);
3636 let until_ms = until_ms
3637 .filter(|t| *t > now_ms)
3638 .unwrap_or(now_ms + DEFAULT_COOLDOWN.as_millis() as u64)
3639 .min(max_until);
3640 let reason = crate::truncate_str(reason, 64).to_string();
3641 tracing::info!(credential = %cred, until_ms, reason = %reason, "account cooldown reported");
3642 let key = self.cooldown_key(cred);
3643 self.cooldowns
3644 .lock()
3645 .unwrap_or_else(|e| e.into_inner())
3646 .insert(key.clone(), Cooldown { until_ms, reason });
3647 self.forget_snapshots(cred);
3651 self.usage_reads.forget_last_good(&self.read_key(cred));
3652 Ok(())
3653 }
3654
3655 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
3656 let resp = self.send(&request).await?;
3657 let status = resp.status().as_u16();
3658 let body = read_body_capped(resp, self.max_response_bytes).await?;
3659 Ok(ProxyResponse { status, body })
3660 }
3661
3662 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
3663 let mut request = request;
3664 request.stream = true;
3665 let resp = self.send(&request).await?;
3666 let status = resp.status();
3667 if !status.is_success() {
3668 let label = read_error_body_for_classification(resp, MAX_PROXY_ERROR_CLASSIFY_BYTES)
3678 .await
3679 .as_deref()
3680 .and_then(crate::error::vetted_proxy_error_label);
3681 tracing::warn!(
3682 provider = %request.provider,
3683 status = status.as_u16(),
3684 label = label.unwrap_or("-"),
3685 "provider request failed"
3686 );
3687 let label_suffix = label.map(|l| format!(" [{l}]")).unwrap_or_default();
3688 return Err(BrokerError::Transport(format!(
3689 "provider request failed: {status}{label_suffix}"
3690 )));
3691 }
3692 use futures::StreamExt;
3693 let stream = resp
3694 .bytes_stream()
3695 .map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
3696 Ok(Box::pin(stream))
3697 }
3698
3699 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
3700 let token = self.access_token(OAuthProviderId::Anthropic).await?;
3703 let url = self
3704 .anthropic_usage_url
3705 .as_deref()
3706 .unwrap_or(ANTHROPIC_USAGE_URL);
3707 let resp = self
3708 .http
3709 .get(url)
3710 .timeout(self.request_timeout)
3711 .bearer_auth(&token.token)
3712 .header("anthropic-beta", ANTHROPIC_OAUTH_BETA)
3713 .send()
3714 .await
3715 .map_err(|e| BrokerError::Transport(format!("usage request failed: {e}")))?;
3716 let status = resp.status();
3717 if !status.is_success() {
3718 drop(resp);
3721 return Err(BrokerError::Transport(format!(
3722 "usage request failed: {status}"
3723 )));
3724 }
3725 let body = read_body_capped(resp, self.max_response_bytes).await?;
3726 serde_json::from_str(&body)
3727 .map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
3728 }
3729
3730 async fn cloud_catalog(
3731 &self,
3732 provider: CloudProviderId,
3733 context_ref: &str,
3734 allow_stale: bool,
3735 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
3736 self.cloud_backend
3737 .as_ref()
3738 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
3739 .catalog(provider, context_ref, allow_stale)
3740 .await
3741 }
3742
3743 async fn cloud_invoke(
3744 &self,
3745 provider: CloudProviderId,
3746 context_ref: &str,
3747 model_id: &str,
3748 request: InvokeRequest,
3749 ) -> Result<CloudEventStream, BrokerError> {
3750 self.cloud_backend
3751 .as_ref()
3752 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
3753 .invoke(provider, context_ref, model_id, request)
3754 .await
3755 }
3756
3757 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
3758 let mut out = Vec::new();
3759 let policy = self.policy();
3760 for descriptor in super::provider::registry().iter() {
3761 let accounts = self.accounts(descriptor.id).await.unwrap_or_default();
3762 let configured = match policy.selector(descriptor.id) {
3763 AccountSelector::Account(account) => accounts
3764 .iter()
3765 .any(|a| a.label == account.label_str()),
3766 AccountSelector::Auto => !accounts.is_empty(),
3767 AccountSelector::Invalid { .. } => false,
3768 };
3769 out.push(ProviderStatus {
3770 key: descriptor.id.as_str().to_string(),
3771 name: descriptor.display_name.to_string(),
3772 kind: CredentialKind::OAuth,
3773 configured,
3774 accounts,
3775 });
3776 }
3777 for spec in STATIC_PROVIDERS {
3778 out.push(ProviderStatus {
3779 key: spec.key.to_string(),
3780 name: spec.name.to_string(),
3781 kind: CredentialKind::StaticKey,
3782 configured: static_key_configured(spec.key),
3783 accounts: Vec::new(),
3784 });
3785 }
3786 out.push(ProviderStatus {
3787 key: LOCAL_PROVIDER_KEY.to_string(),
3788 name: "Local endpoint".to_string(),
3789 kind: CredentialKind::LocalEndpoint,
3790 configured: true,
3791 accounts: Vec::new(),
3792 });
3793 Ok(out)
3794 }
3795}
3796
3797pub struct RemoteBroker {
3803 http: reqwest::Client,
3804 endpoint: String,
3805 machine_token: super::MachineToken,
3808 cache: super::TokenCache,
3809 account_policy: Option<AccountPolicy>,
3811}
3812
3813async fn route_missing(resp: reqwest::Response) -> bool {
3817 let json = resp
3818 .headers()
3819 .get(reqwest::header::CONTENT_TYPE)
3820 .and_then(|v| v.to_str().ok())
3821 .is_some_and(|v| v.contains("json"));
3822 !json
3823 && read_body_capped(resp, 1024)
3824 .await
3825 .is_ok_and(|body| body.trim().is_empty())
3826}
3827
3828impl RemoteBroker {
3829 async fn predates_usage_snapshot(&self) -> Result<bool, BrokerError> {
3833 let resp = self
3834 .machine_token
3835 .send(|t| {
3836 self.http
3837 .get(format!("{}/usage/snapshot", self.endpoint))
3838 .query(&[("provider", "anthropic"), ("account", "")])
3839 .bearer_auth(t)
3840 })
3841 .await
3842 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
3843 match resp.status().as_u16() {
3844 401 => Err(BrokerError::Unauthorized),
3845 404 => Ok(route_missing(resp).await),
3846 _ => Ok(false),
3847 }
3848 }
3849
3850 pub fn new(
3851 endpoint: impl Into<String>,
3852 machine_token: impl Into<String>,
3853 http: reqwest::Client,
3854 cache: super::TokenCache,
3855 ) -> Self {
3856 let endpoint = endpoint.into().trim_end_matches('/').to_string();
3857 let http = reqwest::Client::builder()
3860 .redirect(reqwest::redirect::Policy::none())
3861 .connect_timeout(Duration::from_secs(10))
3862 .timeout(Duration::from_secs(90))
3863 .build()
3864 .unwrap_or(http);
3865 let machine_token = super::MachineToken::for_endpoint(&endpoint, machine_token);
3866 Self {
3867 http,
3868 endpoint,
3869 machine_token,
3870 cache,
3871 account_policy: None,
3872 }
3873 }
3874
3875 pub fn with_account_policy(mut self, policy: AccountPolicy) -> Self {
3877 self.account_policy = Some(policy);
3878 self
3879 }
3880
3881 fn policy(&self) -> AccountPolicy {
3882 self.account_policy
3883 .clone()
3884 .unwrap_or_else(AccountPolicy::from_environment)
3885 }
3886
3887 fn fetcher(&self) -> super::BrokerClient {
3888 super::BrokerClient::with_machine_token(
3889 self.endpoint.clone(),
3890 self.machine_token.clone(),
3891 self.http.clone(),
3892 )
3893 }
3894
3895 fn scope(&self) -> String {
3897 super::credential_source::source_scope(&self.endpoint, &self.machine_token.get())
3898 }
3899
3900 async fn resolve_pinned(
3905 &self,
3906 provider: OAuthProviderId,
3907 model: Option<&str>,
3908 ) -> Result<PinnedToken, BrokerError> {
3909 match self.policy().selector(provider) {
3910 AccountSelector::Account(account) => {
3911 let credential = CredentialRef::new(provider, account);
3912 let token = self.access_token_for(&credential).await?;
3913 Ok(PinnedToken { credential, token })
3914 }
3915 AccountSelector::Auto => self.fetch_auto(provider, model).await,
3916 AccountSelector::Invalid { source, reason } => {
3917 Err(BrokerError::InvalidAccount(format!("{source}: {reason}")))
3918 }
3919 }
3920 }
3921
3922 async fn fetch_auto(
3926 &self,
3927 provider: OAuthProviderId,
3928 model: Option<&str>,
3929 ) -> Result<PinnedToken, BrokerError> {
3930 let mut query: Vec<(&str, &str)> = vec![
3931 ("provider", provider.as_str()),
3932 ("account", super::account::AUTO_ACCOUNT_NAME),
3933 ];
3934 if let Some(model) = model {
3935 query.push(("model", model));
3936 }
3937 let resp = self
3938 .machine_token
3939 .send(|t| {
3940 self.http
3941 .get(format!("{}/token", self.endpoint))
3942 .query(&query)
3943 .bearer_auth(t)
3944 })
3945 .await
3946 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
3947 match resp.status().as_u16() {
3948 401 => return Err(BrokerError::Unauthorized),
3949 503 => {
3950 drop(resp);
3951 return Err(BrokerError::NoAccountAvailable {
3952 provider: provider.as_str().to_string(),
3953 reason: "broker reported no account with proven capacity".into(),
3954 failure: Some(super::quota_policy::SelectionFailure::new(
3956 super::quota_policy::SelectionFailureKind::RemoteNoCapacity,
3957 )),
3958 });
3959 }
3960 s if !(200..300).contains(&s) => {
3961 drop(resp);
3962 return Err(BrokerError::Transport(format!("broker returned HTTP {s}")));
3963 }
3964 _ => {}
3965 }
3966 #[derive(Deserialize)]
3967 struct AutoToken {
3968 access_token: String,
3969 expires: u64,
3970 #[serde(default)]
3971 ttl_ms: Option<u64>,
3972 account: Option<String>,
3973 }
3974 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
3975 let tok: AutoToken = serde_json::from_str(&body)
3976 .map_err(|e| BrokerError::Transport(format!("invalid broker token response: {e}")))?;
3977 let account = tok
3978 .account
3979 .as_deref()
3980 .ok_or_else(|| BrokerError::Transport("broker did not report the selected account".into()))
3981 .and_then(|a| Account::parse(a).map_err(BrokerError::InvalidAccount))?;
3982 if tok.access_token.is_empty() {
3983 return Err(BrokerError::Transport("broker returned an empty access_token".into()));
3984 }
3985 let expires = match tok.ttl_ms {
3986 Some(ttl) => crate::epoch_millis().saturating_add(ttl),
3987 None => tok.expires,
3988 };
3989 Ok(PinnedToken {
3990 credential: CredentialRef::new(provider, account),
3991 token: AccessToken {
3992 token: tok.access_token,
3993 expires,
3994 },
3995 })
3996 }
3997
3998 async fn proxy_failure(resp: reqwest::Response, fallback: String) -> BrokerError {
4005 if resp.status() != reqwest::StatusCode::BAD_GATEWAY {
4006 return BrokerError::Transport(fallback);
4007 }
4008 let Ok(body) = read_body_capped(resp, 4096).await else {
4009 return BrokerError::Transport(fallback);
4010 };
4011 let text = serde_json::from_str::<serde_json::Value>(&body)
4012 .ok()
4013 .and_then(|v| v.get("error")?.as_str().map(str::to_string))
4014 .unwrap_or_default();
4015 use super::token::{RefreshFailure, RefreshFailureKind as K};
4016 for kind in [K::Transient, K::LockBusy, K::Other] {
4017 let prefix = BrokerError::Refresh(RefreshFailure::new(kind, "")).to_string();
4018 if text.starts_with(&prefix) {
4019 return BrokerError::Refresh(RefreshFailure::new(
4020 kind,
4021 "the remote credential broker could not refresh this credential",
4022 ));
4023 }
4024 }
4025 BrokerError::Transport(fallback)
4026 }
4027
4028 async fn post_proxy(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
4029 request.validate()?;
4030 let resp = self
4031 .machine_token
4032 .send(|t| {
4033 self.http
4034 .post(format!("{}/proxy", self.endpoint))
4035 .bearer_auth(t)
4036 .json(request)
4037 })
4038 .await
4039 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
4040 match resp.status().as_u16() {
4041 401 => Err(BrokerError::Unauthorized),
4042 s @ (400 | 403) => {
4043 drop(resp);
4047 Err(BrokerError::Denied(format!(
4048 "broker rejected the request (HTTP {s})"
4049 )))
4050 }
4051 _ => Ok(resp),
4052 }
4053 }
4054}
4055
4056#[async_trait]
4057impl CredentialBroker for RemoteBroker {
4058 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
4059 Ok(self.access_token_pinned(provider).await?.token)
4060 }
4061
4062 fn account_selector(&self, provider: OAuthProviderId) -> AccountSelector {
4063 self.policy().selector(provider)
4064 }
4065
4066 async fn access_token_for(&self, cred: &CredentialRef) -> Result<AccessToken, BrokerError> {
4067 let fetcher = self.fetcher();
4068 let tok = super::credential_source::resolve_remote_credential(
4069 &fetcher,
4070 &self.cache,
4071 &self.scope(),
4072 cred,
4073 super::DEFAULT_MARGIN_MS,
4074 )
4075 .await
4076 .map_err(|e| e.into_broker_error(cred))?;
4077 Ok(AccessToken {
4078 token: tok.access_token,
4079 expires: tok.expires,
4080 })
4081 }
4082
4083 async fn access_token_pinned(
4084 &self,
4085 provider: OAuthProviderId,
4086 ) -> Result<PinnedToken, BrokerError> {
4087 self.resolve_pinned(provider, None).await
4088 }
4089
4090 async fn access_token_pinned_for(
4091 &self,
4092 provider: OAuthProviderId,
4093 model: Option<&str>,
4094 ) -> Result<PinnedToken, BrokerError> {
4095 self.resolve_pinned(provider, model).await
4096 }
4097
4098 async fn accounts(&self, provider: OAuthProviderId) -> Result<Vec<AccountSummary>, BrokerError> {
4099 let caps = self.capabilities().await?;
4100 let accounts = caps
4101 .into_iter()
4102 .find(|c| c.key == provider.as_str())
4103 .map(|c| c.accounts)
4104 .unwrap_or_default();
4105 if accounts.is_empty() && self.predates_usage_snapshot().await? {
4108 return Err(BrokerError::BrokerTooOld("account listing".into()));
4109 }
4110 Ok(accounts)
4111 }
4112
4113 async fn usage(&self, cred: &CredentialRef) -> Result<super::usage::UsageSnapshot, BrokerError> {
4114 let resp = self
4115 .machine_token
4116 .send(|t| {
4117 self.http
4118 .get(format!("{}/usage/snapshot", self.endpoint))
4119 .query(&[
4120 ("provider", cred.provider.as_str()),
4121 ("account", cred.account.label_str()),
4122 ])
4123 .bearer_auth(t)
4124 })
4125 .await
4126 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
4127 match resp.status().as_u16() {
4128 401 => Err(BrokerError::Unauthorized),
4129 404 => {
4130 if !route_missing(resp).await {
4131 return Err(BrokerError::UnknownAccount {
4132 provider: cred.provider.as_str().to_string(),
4133 label: cred.account.label_str().to_string(),
4134 });
4135 }
4136 if cred.provider != OAuthProviderId::Anthropic || !cred.account.is_default() {
4137 return Err(BrokerError::BrokerTooOld("per-account usage".into()));
4138 }
4139 let body = self.anthropic_usage().await?.to_string();
4141 super::usage::parse_anthropic_usage(
4142 &body,
4143 cred.account.label_str(),
4144 crate::epoch_millis(),
4145 )
4146 .map_err(super::usage::UsageError::into_broker_error)
4147 }
4148 s if !(200..300).contains(&s) => {
4149 drop(resp);
4151 Err(BrokerError::Transport(format!(
4152 "broker usage snapshot returned HTTP {s}"
4153 )))
4154 }
4155 _ => {
4156 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
4157 let snapshot: super::usage::UsageSnapshot = serde_json::from_str(&body)
4158 .map_err(|e| BrokerError::Transport(format!("invalid usage snapshot: {e}")))?;
4159 if snapshot.provider != cred.provider.as_str()
4160 || snapshot.account != cred.account.label_str()
4161 {
4162 return Err(BrokerError::Transport(
4163 "broker returned a usage snapshot for a different account".into(),
4164 ));
4165 }
4166 Ok(snapshot)
4167 }
4168 }
4169 }
4170
4171 async fn report_cooldown(
4172 &self,
4173 cred: &CredentialRef,
4174 until_ms: Option<u64>,
4175 reason: &str,
4176 ) -> Result<(), BrokerError> {
4177 let resp = self
4178 .machine_token
4179 .send(|t| {
4180 self.http
4181 .post(format!("{}/accounts/cooldown", self.endpoint))
4182 .bearer_auth(t)
4183 .json(&serde_json::json!({
4184 "provider": cred.provider.as_str(),
4185 "account": cred.account.label_str(),
4186 "until_ms": until_ms,
4187 "reason": crate::truncate_str(reason, 64),
4188 }))
4189 })
4190 .await
4191 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
4192 match resp.status().as_u16() {
4193 401 => Err(BrokerError::Unauthorized),
4194 s if !(200..300).contains(&s) => {
4195 Err(BrokerError::Transport(format!("broker returned HTTP {s}")))
4196 }
4197 _ => Ok(()),
4198 }
4199 }
4200
4201 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
4202 let mut request = request;
4203 request.stream = false;
4204 let resp = self.post_proxy(&request).await?;
4205 if !resp.status().is_success() {
4206 let fallback = format!("broker proxy returned HTTP {}", resp.status());
4207 return Err(Self::proxy_failure(resp, fallback).await);
4208 }
4209 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
4210 serde_json::from_str::<ProxyResponse>(&body)
4211 .map_err(|e| BrokerError::Transport(format!("invalid broker proxy response: {e}")))
4212 }
4213
4214 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
4215 let mut request = request;
4216 request.stream = true;
4217 let resp = self.post_proxy(&request).await?;
4218 let status = resp.status();
4219 if !status.is_success() {
4220 let fallback = format!("broker proxy stream failed: {status}");
4223 return Err(Self::proxy_failure(resp, fallback).await);
4224 }
4225 use futures::StreamExt;
4226 let stream = resp
4227 .bytes_stream()
4228 .map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
4229 Ok(Box::pin(stream))
4230 }
4231
4232 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
4233 let resp = self
4236 .machine_token
4237 .send(|t| {
4238 self.http
4239 .get(format!("{}/usage", self.endpoint))
4240 .bearer_auth(t)
4241 })
4242 .await
4243 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
4244 match resp.status().as_u16() {
4245 401 => Err(BrokerError::Unauthorized),
4246 s if !(200..300).contains(&s) => {
4247 drop(resp);
4249 Err(BrokerError::Transport(format!(
4250 "broker usage returned HTTP {s}"
4251 )))
4252 }
4253 _ => {
4254 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
4255 serde_json::from_str(&body)
4256 .map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
4257 }
4258 }
4259 }
4260
4261 async fn cloud_catalog(
4262 &self,
4263 provider: CloudProviderId,
4264 context_ref: &str,
4265 allow_stale: bool,
4266 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
4267 let resp = self.machine_token.send(|t| self.http.post(format!("{}/cloud/catalog", self.endpoint)).bearer_auth(t).json(&serde_json::json!({"provider":provider,"context_ref":context_ref,"allow_stale":allow_stale}))).await.map_err(|e| BrokerError::Transport(e.to_string()))?;
4268 if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
4269 return Err(BrokerError::Unauthorized);
4270 }
4271 if !resp.status().is_success() {
4272 return Err(BrokerError::Transport(format!(
4273 "cloud catalog returned HTTP {}",
4274 resp.status()
4275 )));
4276 }
4277 let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
4278 let entries: Vec<CloudCatalogEntry> = serde_json::from_str(&body)
4279 .map_err(|e| BrokerError::Transport(format!("invalid cloud catalog: {e}")))?;
4280 if entries.len() > MAX_CLOUD_CATALOG_ENTRIES
4281 || entries.iter().any(|e| {
4282 e.provider != provider || !e.context_ref.starts_with("ctx-") || e.id.len() > 512
4283 })
4284 {
4285 return Err(BrokerError::Transport(
4286 "broker returned an invalid cloud catalog".into(),
4287 ));
4288 }
4289 Ok(entries)
4290 }
4291
4292 async fn cloud_invoke(
4293 &self,
4294 provider: CloudProviderId,
4295 context_ref: &str,
4296 model_id: &str,
4297 request: InvokeRequest,
4298 ) -> Result<CloudEventStream, BrokerError> {
4299 let resp = self.machine_token.send(|t| self.http.post(format!("{}/cloud/invoke", self.endpoint)).bearer_auth(t).json(&serde_json::json!({"provider":provider,"context_ref":context_ref,"model_id":model_id,"request":request}))).await.map_err(|e| BrokerError::Transport(e.to_string()))?;
4300 if !resp.status().is_success() {
4301 return Err(if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
4302 BrokerError::Unauthorized
4303 } else {
4304 BrokerError::Transport(format!("cloud invoke returned HTTP {}", resp.status()))
4305 });
4306 }
4307 use futures::StreamExt;
4308 let chunks = Box::pin(resp.bytes_stream());
4309 let stream = futures::stream::unfold(
4310 (chunks, Vec::<u8>::new(), false, false),
4311 |(mut chunks, mut buffer, done, failed)| async move {
4312 if failed {
4313 return None;
4314 }
4315 loop {
4316 if let Some(end) = buffer.iter().position(|b| *b == b'\n') {
4317 let line: Vec<u8> = buffer.drain(..=end).collect();
4318 let parsed = serde_json::from_slice::<CloudEvent>(&line[..end])
4319 .map_err(|_| BrokerError::Transport("invalid cloud event".into()));
4320 let event = match parsed {
4321 Ok(CloudEvent::Done) if done => Err(BrokerError::Transport(
4322 "duplicate cloud terminal event".into(),
4323 )),
4324 Ok(CloudEvent::Done) => Ok(CloudEvent::Done),
4325 Ok(_) if done => Err(BrokerError::Transport(
4326 "cloud data followed terminal event".into(),
4327 )),
4328 other => other,
4329 };
4330 let next_done = done || matches!(event, Ok(CloudEvent::Done));
4331 let next_failed = event.is_err();
4332 return Some((event, (chunks, buffer, next_done, next_failed)));
4333 }
4334 match chunks.next().await {
4335 Some(Ok(chunk))
4336 if buffer.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES =>
4337 {
4338 buffer.extend_from_slice(&chunk)
4339 }
4340 Some(Ok(_)) => {
4341 return Some((
4342 Err(BrokerError::Transport("cloud event exceeded limit".into())),
4343 (chunks, Vec::new(), done, true),
4344 ))
4345 }
4346 Some(Err(_)) => {
4347 return Some((
4348 Err(BrokerError::Transport("broker stream failed".into())),
4349 (chunks, Vec::new(), done, true),
4350 ))
4351 }
4352 None if buffer.is_empty() && done => return None,
4353 None if buffer.is_empty() => {
4354 return Some((
4355 Err(BrokerError::Transport(
4356 "cloud stream ended without terminal event".into(),
4357 )),
4358 (chunks, Vec::new(), done, true),
4359 ))
4360 }
4361 None => {
4362 return Some((
4363 Err(BrokerError::Transport("truncated cloud event".into())),
4364 (chunks, Vec::new(), done, true),
4365 ))
4366 }
4367 }
4368 }
4369 },
4370 );
4371 Ok(Box::pin(stream))
4372 }
4373
4374 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
4375 let resp = self
4376 .machine_token
4377 .send(|t| {
4378 self.http
4379 .get(format!("{}/capabilities", self.endpoint))
4380 .bearer_auth(t)
4381 })
4382 .await
4383 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
4384 match resp.status().as_u16() {
4385 401 => Err(BrokerError::Unauthorized),
4386 s if !(200..300).contains(&s) => {
4387 Err(BrokerError::Transport(format!("broker returned HTTP {s}")))
4388 }
4389 _ => {
4390 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
4391 serde_json::from_str::<Vec<ProviderStatus>>(&body).map_err(|e| {
4392 BrokerError::Transport(format!("invalid capabilities response: {e}"))
4393 })
4394 }
4395 }
4396 }
4397}
4398
4399pub fn preflight_cloud_capability(
4408 provider: CloudProviderId,
4409 needs_tools: bool,
4410) -> Result<(), BrokerError> {
4411 if needs_tools && !provider.supports_tools() {
4412 return Err(BrokerError::UnsupportedCapability {
4413 provider: provider.to_string(),
4414 capability: "tools".into(),
4415 });
4416 }
4417 Ok(())
4418}
4419
4420pub fn local_auth_state(
4426 cache: &super::TokenCache,
4427 query: &super::auth_state::AuthStateQuery,
4428 live: bool,
4429) -> Result<super::auth_state::AuthStateReport, BrokerError> {
4430 let key = (
4431 storage::auth_file_path(),
4432 crate::config::resolve_read_path("config"),
4433 );
4434 let broker = cache
4435 .local_brokers
4436 .lock()
4437 .unwrap_or_else(|e| e.into_inner())
4438 .get(&key)
4439 .cloned();
4440 let mut report = auth_state_report(broker.as_deref(), query)?;
4441 report.live = live;
4442 Ok(report)
4443}
4444
4445fn auth_state_report(
4446 broker: Option<&LocalBroker>,
4447 query: &super::auth_state::AuthStateQuery,
4448) -> Result<super::auth_state::AuthStateReport, BrokerError> {
4449 use super::auth_state::{
4450 AccountState, AuthStateReport, IdentityKind, ObservedOutcome, AUTH_STATE_SCHEMA_VERSION,
4451 };
4452 let provider = match query.provider.as_deref() {
4453 Some(p) => Some(
4454 p.parse::<OAuthProviderId>()
4455 .map_err(|_| BrokerError::UnknownProvider(p.to_string()))?,
4456 ),
4457 None => None,
4458 };
4459 let mut summaries = match provider {
4460 Some(p) => storage::list_accounts(p),
4461 None => storage::list_all_accounts(),
4462 }
4463 .map_err(BrokerError::Credential)?;
4464 if let Some(label) = query.account.as_deref() {
4465 summaries.retain(|s| s.label == label);
4466 if summaries.is_empty() {
4467 return Err(BrokerError::UnknownAccount {
4468 provider: provider.map_or("any provider", |p| p.as_str()).to_string(),
4469 label: label.to_string(),
4470 });
4471 }
4472 }
4473 let creds: Vec<CredentialRef> = summaries
4474 .iter()
4475 .filter_map(AccountSummary::credential_ref)
4476 .collect();
4477 let mut cleared = Vec::new();
4478 if query.clear {
4479 if let Some(broker) = broker {
4480 let whole: Vec<OAuthProviderId> = match (&query.account, provider) {
4481 (Some(_), _) => Vec::new(),
4482 (None, Some(p)) => vec![p],
4483 (None, None) => super::provider::DESCRIPTORS.iter().map(|d| d.id).collect(),
4484 };
4485 broker.clear_state(&creds, &whole);
4486 }
4487 cleared = creds
4488 .iter()
4489 .map(|c| format!("{}:{}", c.provider, c.account.label_str()))
4490 .collect();
4491 tracing::info!(
4492 accounts = %cleared.join(","),
4493 broker_active = broker.is_some(),
4494 "auth_state_cleared"
4495 );
4496 }
4497 let store = storage::auth_file_path();
4498 let now_ms = crate::epoch_millis();
4499 let mut accounts = Vec::with_capacity(summaries.len());
4500 for summary in &summaries {
4501 let Some(cred) = summary.credential_ref() else {
4502 continue;
4503 };
4504 let observed = broker
4505 .and_then(|b| {
4506 b.observations
4507 .lock()
4508 .unwrap_or_else(|e| e.into_inner())
4509 .get(&LocalBroker::state_key(&cred))
4510 .cloned()
4511 })
4512 .unwrap_or_default();
4513 accounts.push(AccountState {
4514 provider: summary.provider.clone(),
4515 label: summary.label.clone(),
4516 identity: if summary.seat_fingerprint.is_some() {
4517 IdentityKind::Stable
4518 } else {
4519 IdentityKind::Rotating
4520 },
4521 access_expires_ms: (summary.expires != 0).then_some(summary.expires),
4522 cooldown: broker.and_then(|b| b.cooldown_state(&cred, now_ms)),
4523 last_usage: observed.last_usage,
4524 last_usage_ok_ms: observed.last_usage_ok_ms,
4525 last_refresh: super::token::last_refresh_exchange(&store, &cred.storage_key()).map(
4526 |x| ObservedOutcome {
4527 at_ms: x.at_ms,
4528 failure: x.failure.map(|kind| {
4529 let failure = refresh_kind_failure(kind);
4530 match x.http_status {
4531 Some(status) => failure.with_status(status),
4532 None => failure,
4533 }
4534 }),
4535 },
4536 ),
4537 });
4538 }
4539 let store_prefix = format!("{}|", store.display());
4540 let mut selections: Vec<_> = broker
4541 .map(|b| {
4542 b.selections
4543 .lock()
4544 .unwrap_or_else(|e| e.into_inner())
4545 .iter()
4546 .filter(|(k, s)| {
4547 k.starts_with(&store_prefix)
4548 && provider.is_none_or(|p| s.provider == p.as_str())
4549 && query
4550 .account
4551 .as_deref()
4552 .is_none_or(|label| selection_involves(s, label))
4553 })
4554 .map(|(_, s)| s.clone())
4555 .collect()
4556 })
4557 .unwrap_or_default();
4558 selections.sort_by_key(|s| std::cmp::Reverse(s.at_ms));
4559 Ok(AuthStateReport {
4560 schema_version: AUTH_STATE_SCHEMA_VERSION,
4561 live: false,
4562 broker_active: broker.is_some(),
4563 generated_at_ms: now_ms,
4564 store: store.display().to_string(),
4565 accounts,
4566 selections,
4567 cleared,
4568 })
4569}
4570
4571fn selection_involves(s: &super::auth_state::SelectionState, label: &str) -> bool {
4574 s.chosen.as_deref() == Some(label)
4575 || s.failure
4576 .as_ref()
4577 .is_some_and(|f| f.accounts.iter().any(|a| a.account == label))
4578}
4579
4580pub fn broker_from_source(
4582 source: &super::CredentialSource,
4583 cache: &super::TokenCache,
4584 http: reqwest::Client,
4585) -> Arc<dyn CredentialBroker> {
4586 match source {
4587 super::CredentialSource::Local => {
4588 let key = (
4592 storage::auth_file_path(),
4593 crate::config::resolve_read_path("config"),
4594 );
4595 let mut brokers = cache
4596 .local_brokers
4597 .lock()
4598 .unwrap_or_else(|e| e.into_inner());
4599 brokers
4600 .entry(key)
4601 .or_insert_with(|| Arc::new(LocalBroker::new(http)))
4602 .clone()
4603 }
4604 super::CredentialSource::Remote {
4605 endpoint,
4606 machine_token,
4607 } => Arc::new(RemoteBroker::new(
4608 endpoint.clone(),
4609 machine_token.clone(),
4610 http,
4611 cache.clone(),
4612 )),
4613 }
4614}
4615
4616static GLOBAL_BROKER: std::sync::RwLock<Option<Arc<dyn CredentialBroker>>> =
4617 std::sync::RwLock::new(None);
4618
4619static GLOBAL_BROKER_INSTALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
4620
4621pub fn set_global_broker(broker: Arc<dyn CredentialBroker>) {
4626 *GLOBAL_BROKER.write().expect("broker registry poisoned") = Some(broker);
4627 let n = GLOBAL_BROKER_INSTALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
4628 if crate::core::memstat::mem_trace_enabled() {
4629 tracing::info!(target: "agent_core::memstat", installs = n, "global broker installed");
4630 }
4631}
4632
4633pub fn global_broker_install_count() -> u64 {
4635 GLOBAL_BROKER_INSTALLS.load(std::sync::atomic::Ordering::Relaxed)
4636}
4637
4638pub fn global_broker() -> Arc<dyn CredentialBroker> {
4641 if let Some(b) = GLOBAL_BROKER
4642 .read()
4643 .expect("broker registry poisoned")
4644 .clone()
4645 {
4646 return b;
4647 }
4648 let default: Arc<dyn CredentialBroker> = Arc::new(LocalBroker::new(
4649 reqwest::Client::builder()
4650 .connect_timeout(std::time::Duration::from_secs(10))
4651 .timeout(std::time::Duration::from_secs(300))
4652 .build()
4653 .unwrap_or_default(),
4654 ));
4655 let mut guard = GLOBAL_BROKER.write().expect("broker registry poisoned");
4656 if let Some(b) = guard.clone() {
4657 return b;
4658 }
4659 *guard = Some(default.clone());
4660 default
4661}
4662
4663pub async fn broker_access_token(
4666 provider: OAuthProviderId,
4667 source: &super::CredentialSource,
4668 cache: &super::TokenCache,
4669 http: &reqwest::Client,
4670) -> Result<AccessToken, BrokerError> {
4671 broker_from_source(source, cache, http.clone())
4672 .access_token(provider)
4673 .await
4674}
4675
4676pub fn static_key_status_map() -> BTreeMap<String, StaticKeyStatus> {
4679 STATIC_PROVIDERS
4680 .iter()
4681 .map(|s| (s.key.to_string(), static_key_status(s.key)))
4682 .collect()
4683}
4684
4685#[cfg(test)]
4686mod tests {
4687 use super::*;
4688
4689 #[test]
4690 fn login_availability_handles_named_only_and_explicit_policy() {
4691 let dir = tempfile::tempdir().unwrap();
4692 let path = dir.path().join("auth.json");
4693 let provider = OAuthProviderId::Anthropic;
4694 let empty = AccountPolicy::new();
4695 assert!(!oauth_provider_logged_in_at(&path, &empty, provider));
4696 std::fs::write(&path, serde_json::json!({
4697 "anthropic@work": {"type":"oauth", "access":"synthetic-access",
4698 "refresh":"synthetic-refresh", "expires":1}
4699 }).to_string()).unwrap();
4700 assert!(oauth_provider_logged_in_at(&path, &empty, provider));
4702 for selector in [AccountSelector::Auto,
4703 AccountSelector::parse_or_invalid("../bad", "test"),
4704 AccountSelector::Account(Account::named("work").unwrap())] {
4705 let policy = AccountPolicy::new().with(provider, selector);
4706 assert!(oauth_provider_logged_in_at(&path, &policy, provider));
4707 }
4708 let missing = AccountPolicy::new().with(provider,
4709 AccountSelector::Account(Account::named("missing").unwrap()));
4710 assert!(!oauth_provider_logged_in_at(&path, &missing, provider));
4711 assert!(!oauth_provider_logged_in_at(&path, &empty, OAuthProviderId::OpenAiCodex));
4712 std::fs::write(&path, "{broken").unwrap();
4713 assert!(!oauth_provider_logged_in_at(&path, &empty, provider));
4714 }
4715
4716 #[test]
4721 fn preflight_rejects_tool_requiring_cloud_routes_with_typed_error() {
4722 for provider in [
4723 CloudProviderId::AzureOpenAi,
4724 CloudProviderId::AwsBedrock,
4725 CloudProviderId::GoogleVertex,
4726 ] {
4727 let err = preflight_cloud_capability(provider, true)
4728 .expect_err("tool-requiring cloud route must fail pre-flight");
4729 assert_eq!(
4730 err,
4731 BrokerError::UnsupportedCapability {
4732 provider: provider.to_string(),
4733 capability: "tools".into(),
4734 }
4735 );
4736 assert!(err.to_string().contains("text-only"));
4737 assert_eq!(preflight_cloud_capability(provider, false), Ok(()));
4738 }
4739 }
4740
4741 async fn spawn_ndjson(body: &'static str) -> String {
4742 use axum::{body::Body, routing::post, Router};
4743 use bytes::Bytes;
4744 let app = Router::new().route(
4745 "/cloud/invoke",
4746 post(move || async move {
4747 let chunks = body
4748 .as_bytes()
4749 .chunks(3)
4750 .map(|chunk| Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(chunk)));
4751 axum::response::Response::builder()
4752 .header("content-type", "application/x-ndjson")
4753 .body(Body::from_stream(futures::stream::iter(chunks)))
4754 .unwrap()
4755 }),
4756 );
4757 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
4758 let addr = listener.local_addr().unwrap();
4759 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
4760 format!("http://{addr}")
4761 }
4762
4763 async fn remote_events(body: &'static str) -> Vec<Result<CloudEvent, BrokerError>> {
4764 use futures::StreamExt;
4765 let broker = RemoteBroker::new(
4766 spawn_ndjson(body).await,
4767 "opaque-machine-token",
4768 reqwest::Client::new(),
4769 super::super::TokenCache::new(),
4770 );
4771 broker
4772 .cloud_invoke(
4773 CloudProviderId::GoogleVertex,
4774 "ctx-opaque",
4775 "google-vertex/publishers/google/models/test",
4776 InvokeRequest {
4777 messages: vec![],
4778 tools: vec![],
4779 stream: true,
4780 options: Default::default(),
4781 },
4782 )
4783 .await
4784 .unwrap()
4785 .collect()
4786 .await
4787 }
4788
4789 #[tokio::test]
4790 async fn remote_ndjson_requires_exactly_one_terminal_and_stops_after_error() {
4791 let valid =
4792 remote_events("{\"type\":\"text_delta\",\"delta\":\"hi\"}\n{\"type\":\"done\"}\n")
4793 .await;
4794 assert_eq!(valid.len(), 2);
4795 assert!(matches!(valid[1], Ok(CloudEvent::Done)));
4796
4797 for body in [
4798 "{\"type\":\"text_delta\",\"delta\":\"hi\"}\n",
4799 "{\"type\":\"done\"}\n{\"type\":\"done\"}\n",
4800 "{\"type\":\"done\"}\n{\"type\":\"text_delta\",\"delta\":\"late\"}\n",
4801 ] {
4802 let events = remote_events(body).await;
4803 assert_eq!(events.iter().filter(|event| event.is_err()).count(), 1);
4804 assert!(events.last().unwrap().is_err());
4805 }
4806 }
4807
4808 #[test]
4811 fn proxy_rejects_unlisted_same_host_paths() {
4812 for (provider, path) in [
4813 ("groq", "/v1/keys"),
4814 ("groq", "/admin"),
4815 ("openrouter", "/api/v1/auth/keys"),
4816 ("local", "/audio/speech"),
4817 ] {
4818 let req = ProxyRequest {
4819 provider: provider.into(),
4820 method: ProxyMethod::Get,
4821 path: path.into(),
4822 body: None,
4823 stream: false,
4824 body_bytes: None,
4825 };
4826 match req.validate() {
4827 Err(BrokerError::Denied(msg)) => {
4828 assert!(msg.contains("allowlist"), "got: {msg}")
4829 }
4830 other => panic!("{provider} {path} must be denied, got {other:?}"),
4831 }
4832 }
4833 for path in ["/models", "/chat/completions"] {
4835 let req = ProxyRequest {
4836 provider: "groq".into(),
4837 method: ProxyMethod::Get,
4838 path: path.into(),
4839 body: None,
4840 stream: false,
4841 body_bytes: None,
4842 };
4843 assert!(req.validate().is_ok(), "{path} must be allowed");
4844 }
4845 }
4846
4847 #[test]
4850 fn proxy_rejects_oversize_request_body() {
4851 let req = ProxyRequest {
4852 provider: "groq".into(),
4853 method: ProxyMethod::Post,
4854 path: "/chat/completions".into(),
4855 body: Some(serde_json::json!({
4856 "blob": "x".repeat(MAX_PROXY_REQUEST_BYTES + 1)
4857 })),
4858 stream: false,
4859 body_bytes: None,
4860 };
4861 match req.validate() {
4862 Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
4863 other => panic!("oversize body must be denied, got {other:?}"),
4864 }
4865 }
4866
4867 #[test]
4871 fn proxy_rejects_oversize_body_bytes_without_json_body() {
4872 let req = ProxyRequest {
4873 provider: "groq".into(),
4874 method: ProxyMethod::Post,
4875 path: "/chat/completions".into(),
4876 body: None,
4877 stream: false,
4878 body_bytes: Some(bytes::Bytes::from(vec![b'x'; MAX_PROXY_REQUEST_BYTES + 1])),
4879 };
4880 match req.validate() {
4881 Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
4882 other => panic!("oversize body_bytes must be denied, got {other:?}"),
4883 }
4884 }
4885
4886 #[test]
4890 fn post_json_exact_sets_coherent_body_and_bytes() {
4891 let value = serde_json::json!({"model": "m", "stream": true});
4892 let (req, digest_bytes) =
4893 ProxyRequest::post_json_exact("groq", "/chat/completions", value.clone(), true)
4894 .expect("serializable body must construct");
4895 assert_eq!(req.method, ProxyMethod::Post);
4896 assert!(req.stream);
4897 assert_eq!(req.body.as_ref(), Some(&value));
4898 let stored = req.body_bytes.as_ref().expect("bytes must be set");
4899 assert_eq!(stored, &digest_bytes, "caller digest bytes == wire bytes");
4900 assert_eq!(
4901 serde_json::from_slice::<serde_json::Value>(stored).unwrap(),
4902 value
4903 );
4904 req.validate().expect("coherent request must validate");
4905 }
4906
4907 #[test]
4911 fn mismatched_body_bytes_rejected_by_debug_validation() {
4912 let req = ProxyRequest {
4913 provider: "groq".into(),
4914 method: ProxyMethod::Post,
4915 path: "/chat/completions".into(),
4916 body: Some(serde_json::json!({"model": "claimed"})),
4917 stream: true,
4918 body_bytes: Some(bytes::Bytes::from_static(b"{\"model\":\"actually-sent\"}")),
4919 };
4920 match req.validate() {
4921 Err(BrokerError::Denied(msg)) => {
4922 assert!(msg.contains("body_bytes"), "got: {msg}");
4923 }
4924 other => panic!("incoherent handoff must be denied, got {other:?}"),
4925 }
4926 }
4927
4928 #[test]
4929 fn proxy_request_validation_fails_closed() {
4930 let ok = ProxyRequest {
4931 provider: "groq".into(),
4932 method: ProxyMethod::Post,
4933 path: "/chat/completions".into(),
4934 body: None,
4935 stream: false,
4936 body_bytes: None,
4937 };
4938 assert!(ok.validate().is_ok());
4939
4940 let unknown = ProxyRequest {
4941 provider: "evil".into(),
4942 ..ok.clone()
4943 };
4944 assert!(matches!(
4945 unknown.validate(),
4946 Err(BrokerError::UnknownProvider(_))
4947 ));
4948
4949 let absolute = ProxyRequest {
4950 path: "https://evil.example/x".into(),
4951 ..ok.clone()
4952 };
4953 assert!(matches!(absolute.validate(), Err(BrokerError::Denied(_))));
4954
4955 let traversal = ProxyRequest {
4956 path: "/../secrets".into(),
4957 ..ok.clone()
4958 };
4959 assert!(matches!(traversal.validate(), Err(BrokerError::Denied(_))));
4960
4961 let relative = ProxyRequest {
4962 path: "chat/completions".into(),
4963 ..ok
4964 };
4965 assert!(matches!(relative.validate(), Err(BrokerError::Denied(_))));
4966 }
4967
4968 #[test]
4971 fn proxy_allows_only_pinned_kimi_code_paths() {
4972 let chat = ProxyRequest {
4973 provider: "kimi-code".into(),
4974 method: ProxyMethod::Post,
4975 path: "/chat/completions".into(),
4976 body: None,
4977 stream: true,
4978 body_bytes: None,
4979 };
4980 assert!(chat.validate().is_ok(), "kimi-code chat must be allowed");
4981
4982 for path in ["/models", "/me", "/usages"] {
4983 let get = ProxyRequest {
4984 provider: "kimi-code".into(),
4985 method: ProxyMethod::Get,
4986 path: path.into(),
4987 body: None,
4988 stream: false,
4989 body_bytes: None,
4990 };
4991 assert!(get.validate().is_ok(), "kimi-code GET {path} must pass");
4992 let post = ProxyRequest {
4994 method: ProxyMethod::Post,
4995 ..get
4996 };
4997 assert!(
4998 matches!(post.validate(), Err(BrokerError::Denied(_))),
4999 "kimi-code POST {path} must be denied"
5000 );
5001 }
5002
5003 for path in ["/feedback", "/search", "/fetch", "/completions"] {
5005 let denied = ProxyRequest {
5006 provider: "kimi-code".into(),
5007 method: ProxyMethod::Post,
5008 path: path.into(),
5009 body: None,
5010 stream: false,
5011 body_bytes: None,
5012 };
5013 assert!(
5014 matches!(denied.validate(), Err(BrokerError::Denied(_))),
5015 "kimi-code {path} must be denied"
5016 );
5017 }
5018
5019 let static_kimi = ProxyRequest {
5022 provider: "kimi".into(),
5023 method: ProxyMethod::Post,
5024 path: "/chat/completions".into(),
5025 body: None,
5026 stream: false,
5027 body_bytes: None,
5028 };
5029 assert!(static_kimi.validate().is_ok());
5030 }
5031
5032 #[test]
5036 fn proxy_rejects_oauth_providers_except_codex_catalog() {
5037 let claude = ProxyRequest {
5039 provider: "claude".into(),
5040 method: ProxyMethod::Get,
5041 path: "/models".into(),
5042 body: None,
5043 stream: false,
5044 body_bytes: None,
5045 };
5046 assert!(
5047 matches!(claude.validate(), Err(BrokerError::UnknownProvider(_))),
5048 "claude alias must not be proxyable"
5049 );
5050 let anth = ProxyRequest {
5052 provider: "anthropic".into(),
5053 method: ProxyMethod::Get,
5054 path: "/models".into(),
5055 body: None,
5056 stream: false,
5057 body_bytes: None,
5058 };
5059 assert!(
5060 matches!(anth.validate(), Err(BrokerError::Denied(_))),
5061 "anthropic /models must be denied"
5062 );
5063 let bare = ProxyRequest {
5065 provider: "openai-codex".into(),
5066 method: ProxyMethod::Get,
5067 path: "/models".into(),
5068 body: None,
5069 stream: false,
5070 body_bytes: None,
5071 };
5072 assert!(
5073 matches!(bare.validate(), Err(BrokerError::Denied(_))),
5074 "openai-codex /models without client_version must be denied"
5075 );
5076 }
5077
5078 #[test]
5082 fn proxy_allows_only_pinned_openai_codex_models_path() {
5083 let allowed = ProxyRequest {
5084 provider: "openai-codex".into(),
5085 method: ProxyMethod::Get,
5086 path: "/codex/models?client_version=0.153.3".into(),
5087 body: None,
5088 stream: false,
5089 body_bytes: None,
5090 };
5091 assert!(
5092 allowed.validate().is_ok(),
5093 "exact codex models path must be allowed"
5094 );
5095
5096 for path in [
5097 "/codex/models",
5098 "/codex/models?client_version=",
5099 "/codex/models?foo=1",
5100 "/codex/models?client_version=0.153.3&extra=1",
5101 "/codex/responses",
5102 "/models",
5103 "/models?client_version=0.153.3",
5104 "/backend-api/codex/models?client_version=0.153.3",
5105 "/v1/models?client_version=0.153.3",
5106 "/chat/completions",
5107 ] {
5108 let req = ProxyRequest {
5109 provider: "openai-codex".into(),
5110 method: ProxyMethod::Get,
5111 path: path.into(),
5112 body: None,
5113 stream: false,
5114 body_bytes: None,
5115 };
5116 assert!(
5117 matches!(req.validate(), Err(BrokerError::Denied(_))),
5118 "openai-codex path {path} must be denied"
5119 );
5120 }
5121
5122 let post = ProxyRequest {
5124 provider: "openai-codex".into(),
5125 method: ProxyMethod::Post,
5126 path: "/codex/models?client_version=0.153.3".into(),
5127 body: None,
5128 stream: false,
5129 body_bytes: None,
5130 };
5131 assert!(
5132 matches!(post.validate(), Err(BrokerError::Denied(_))),
5133 "POST openai-codex models must be denied"
5134 );
5135 }
5136
5137 #[test]
5139 fn proxy_allows_only_pinned_github_copilot_paths() {
5140 let models = ProxyRequest {
5141 provider: "github-copilot".into(),
5142 method: ProxyMethod::Get,
5143 path: "/models".into(),
5144 body: None,
5145 stream: false,
5146 body_bytes: None,
5147 };
5148 assert!(models.validate().is_ok());
5149
5150 let chat = ProxyRequest {
5151 provider: "github-copilot".into(),
5152 method: ProxyMethod::Post,
5153 path: "/chat/completions".into(),
5154 body: None,
5155 stream: false,
5156 body_bytes: None,
5157 };
5158 assert!(chat.validate().is_ok());
5159 let responses = ProxyRequest {
5160 provider: "github-copilot".into(),
5161 method: ProxyMethod::Post,
5162 path: "/responses".into(),
5163 body: None,
5164 stream: false,
5165 body_bytes: None,
5166 };
5167 assert!(responses.validate().is_ok());
5168 for path in ["/v1/messages", "/models?x=1", "/embeddings"] {
5169 let request = ProxyRequest {
5170 provider: "github-copilot".into(),
5171 method: ProxyMethod::Post,
5172 path: path.into(),
5173 body: None,
5174 stream: false,
5175 body_bytes: None,
5176 };
5177 assert!(matches!(request.validate(), Err(BrokerError::Denied(_))));
5178 }
5179 }
5180
5181 #[test]
5183 fn proxy_allows_only_pinned_anthropic_models_path() {
5184 for path in [
5185 "/v1/models",
5186 "/v1/models?limit=100",
5187 "/v1/models?limit=100&after_id=claude-opus-4-7",
5188 ] {
5189 let req = ProxyRequest {
5190 provider: "anthropic".into(),
5191 method: ProxyMethod::Get,
5192 path: path.into(),
5193 body: None,
5194 stream: false,
5195 body_bytes: None,
5196 };
5197 assert!(req.validate().is_ok(), "{path} must be allowed");
5198 }
5199 for path in [
5200 "/v1/messages",
5201 "/models",
5202 "/v1/models?foo=1",
5203 "/v1/models?limit=",
5204 "/api/oauth/usage",
5205 "/v1/models?limit=100&evil=1",
5206 "/v1/models?limit=100&limit=50",
5207 "/v1/models?after_id=a&after_id=b",
5208 "/v1/models?limit=100&after_id=a&limit=50",
5209 ] {
5210 let req = ProxyRequest {
5211 provider: "anthropic".into(),
5212 method: ProxyMethod::Get,
5213 path: path.into(),
5214 body: None,
5215 stream: false,
5216 body_bytes: None,
5217 };
5218 assert!(
5219 matches!(req.validate(), Err(BrokerError::Denied(_))),
5220 "anthropic path {path} must be denied"
5221 );
5222 }
5223 }
5224
5225 #[test]
5226 fn anthropic_oauth_catalog_headers_are_bearer_version_beta_not_x_api_key() {
5227 let headers = anthropic_oauth_catalog_request_headers();
5228 let names: Vec<_> = headers.iter().map(|(n, _)| *n).collect();
5229 assert!(names.contains(&"anthropic-version"));
5230 assert!(names.contains(&"anthropic-beta"));
5231 assert!(names.contains(&"accept"));
5232 assert!(!names.iter().any(|n| n.eq_ignore_ascii_case("x-api-key")));
5233 assert!(!names
5234 .iter()
5235 .any(|n| n.eq_ignore_ascii_case("authorization")));
5236 let beta = headers
5237 .iter()
5238 .find(|(n, _)| *n == "anthropic-beta")
5239 .map(|(_, v)| *v)
5240 .expect("beta");
5241 assert_eq!(beta, "oauth-2025-04-20");
5242 }
5243
5244 #[test]
5245 fn anthropic_models_path_round_trip_and_rejects_duplicates() {
5246 assert!(is_allowed_anthropic_path("/v1/models"));
5247 assert!(is_allowed_anthropic_path("/v1/models?limit=100"));
5248 assert!(is_allowed_anthropic_path(
5249 "/v1/models?limit=100&after_id=claude-opus-4-7"
5250 ));
5251 assert!(!is_allowed_anthropic_path("/v1/models?limit=100&limit=50"));
5252 assert!(!is_allowed_anthropic_path(
5253 "/v1/models?after_id=a&after_id=b"
5254 ));
5255 let page0 = "/v1/models?limit=100";
5257 let page1 = "/v1/models?limit=100&after_id=model-1";
5258 assert!(is_allowed_anthropic_path(page0));
5259 assert!(is_allowed_anthropic_path(page1));
5260 }
5261
5262 #[test]
5264 fn proxy_allows_only_pinned_google_gemini_paths() {
5265 for path in [
5266 "/v1internal:loadCodeAssist",
5267 "/v1internal:onboardUser",
5268 "/v1internal:streamGenerateContent",
5269 "/v1internal:countTokens",
5270 "/v1internal/operations/op-12345",
5271 ] {
5272 let req = ProxyRequest {
5273 provider: "google-gemini".into(),
5274 method: ProxyMethod::Post,
5275 path: path.into(),
5276 body: None,
5277 stream: false,
5278 body_bytes: None,
5279 };
5280 assert!(req.validate().is_ok(), "{path} must be allowed");
5281 }
5282 for path in [
5283 "/v1internal:listExperiments",
5285 "/v1internal:fetchAdminControls",
5286 "/v1internal:setCodeAssistGlobalUserSetting",
5287 "/v1internal:generateContent",
5288 "/v2/models",
5290 "/v1beta/models",
5291 "/",
5293 "/v1internal:",
5294 ] {
5295 let req = ProxyRequest {
5296 provider: "google-gemini".into(),
5297 method: ProxyMethod::Post,
5298 path: path.into(),
5299 body: None,
5300 stream: false,
5301 body_bytes: None,
5302 };
5303 assert!(
5304 matches!(req.validate(), Err(BrokerError::Denied(_))),
5305 "{path} must be denied"
5306 );
5307 }
5308 }
5309
5310 #[test]
5311 fn access_token_type_has_no_refresh_field() {
5312 let t: AccessToken =
5315 serde_json::from_str(r#"{"token":"sk-x","expires":123,"refresh":"MUST-NOT-EXIST"}"#)
5316 .unwrap();
5317 let round = serde_json::to_value(&t).unwrap();
5318 assert_eq!(round.get("refresh"), None);
5319 assert_eq!(round["token"], "sk-x");
5320 }
5321
5322 #[test]
5323 fn broker_error_display_never_echoes_values() {
5324 let e = BrokerError::NotConfigured("groq".into());
5326 assert!(!format!("{e}").contains("sk-"));
5327 let e = BrokerError::Unauthorized;
5328 let shown = format!("{e}");
5329 assert!(shown.starts_with("broker rejected machine auth"), "{shown}");
5330 assert!(shown.contains("auth.machine_token") && shown.contains("SYNAPS_MACHINE_TOKEN"), "{shown}");
5332 assert!(shown.contains("auth-broker.token") && shown.contains("restart"), "{shown}");
5333 }
5334
5335 #[test]
5336 fn mask_key_is_short_and_lossy() {
5337 assert_eq!(mask_key("gsk-live-1234567890abcdef"), "gsk-…cdef");
5338 assert_eq!(mask_key("short"), "…");
5339 }
5340
5341 #[test]
5342 fn proxy_request_serde_roundtrip() {
5343 let req = ProxyRequest {
5344 provider: "openrouter".into(),
5345 method: ProxyMethod::Get,
5346 path: "/models".into(),
5347 body: Some(serde_json::json!({"a": 1})),
5348 stream: true,
5349 body_bytes: None,
5350 };
5351 let json = serde_json::to_string(&req).unwrap();
5352 let back: ProxyRequest = serde_json::from_str(&json).unwrap();
5353 assert_eq!(back.provider, "openrouter");
5354 assert!(back.stream);
5355 assert!(matches!(back.method, ProxyMethod::Get));
5356 }
5357
5358 async fn spawn_upstream(app: axum::Router) -> String {
5361 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
5362 let addr = listener.local_addr().unwrap();
5363 tokio::spawn(async move {
5364 axum::serve(listener, app).await.unwrap();
5365 });
5366 format!("http://{addr}")
5367 }
5368
5369 #[tokio::test]
5373 async fn local_broker_streams_local_endpoint_and_applies_key() {
5374 use axum::routing::post;
5375 let seen_auth = Arc::new(std::sync::Mutex::new(String::new()));
5376 let seen = seen_auth.clone();
5377 let app = axum::Router::new().route(
5378 "/chat/completions",
5379 post(move |headers: axum::http::HeaderMap| {
5380 let seen = seen.clone();
5381 async move {
5382 *seen.lock().unwrap() = headers
5383 .get("authorization")
5384 .and_then(|v| v.to_str().ok())
5385 .unwrap_or("")
5386 .to_string();
5387 (
5388 [("content-type", "text/event-stream")],
5389 "data: {\"x\":1}\n\ndata: [DONE]\n\n",
5390 )
5391 }
5392 }),
5393 );
5394 let url = spawn_upstream(app).await;
5395 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
5396 let mut stream = broker
5397 .proxy_stream(ProxyRequest {
5398 provider: LOCAL_PROVIDER_KEY.into(),
5399 method: ProxyMethod::Post,
5400 path: "/chat/completions".into(),
5401 body: Some(serde_json::json!({"model": "m"})),
5402 stream: true,
5403 body_bytes: None,
5404 })
5405 .await
5406 .expect("stream must open");
5407 use futures::StreamExt;
5408 let mut collected = Vec::new();
5409 while let Some(chunk) = stream.next().await {
5410 collected.extend_from_slice(&chunk.unwrap());
5411 }
5412 let text = String::from_utf8(collected).unwrap();
5413 assert!(text.contains("data: {\"x\":1}"));
5414 assert!(text.contains("[DONE]"));
5415 assert_eq!(&*seen_auth.lock().unwrap(), "Bearer local");
5417 }
5418
5419 #[tokio::test]
5422 async fn local_broker_stream_error_is_typed_and_keyless() {
5423 use axum::routing::post;
5424 let app = axum::Router::new().route(
5425 "/chat/completions",
5426 post(|| async {
5427 (
5428 axum::http::StatusCode::UNAUTHORIZED,
5429 "{\"error\":\"bad key\"}",
5430 )
5431 }),
5432 );
5433 let url = spawn_upstream(app).await;
5434 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
5435 let err = match broker
5436 .proxy_stream(ProxyRequest {
5437 provider: LOCAL_PROVIDER_KEY.into(),
5438 method: ProxyMethod::Post,
5439 path: "/chat/completions".into(),
5440 body: None,
5441 stream: true,
5442 body_bytes: None,
5443 })
5444 .await
5445 {
5446 Err(e) => e,
5447 Ok(_) => panic!("non-2xx upstream must not yield a stream"),
5448 };
5449 let msg = format!("{err}");
5450 assert!(msg.contains("401"), "got: {msg}");
5451 assert!(
5452 !msg.to_lowercase().contains("bearer"),
5453 "no auth material in errors"
5454 );
5455 }
5456
5457 #[tokio::test]
5459 async fn local_broker_proxy_returns_status_and_body() {
5460 use axum::routing::get;
5461 let app =
5462 axum::Router::new().route("/models", get(|| async { "{\"data\":[{\"id\":\"m1\"}]}" }));
5463 let url = spawn_upstream(app).await;
5464 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
5465 let resp = broker
5466 .proxy(ProxyRequest {
5467 provider: LOCAL_PROVIDER_KEY.into(),
5468 method: ProxyMethod::Get,
5469 path: "/models".into(),
5470 body: None,
5471 stream: false,
5472 body_bytes: None,
5473 })
5474 .await
5475 .unwrap();
5476 assert_eq!(resp.status, 200);
5477 assert!(resp.body.contains("m1"));
5478 }
5479
5480 #[tokio::test]
5483 async fn local_broker_proxy_rejects_oversize_response() {
5484 use axum::routing::get;
5485 let app = axum::Router::new().route("/models", get(|| async { "x".repeat(4096) }));
5486 let url = spawn_upstream(app).await;
5487 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
5488 .with_max_response_bytes(64);
5489 let err = broker
5490 .proxy(ProxyRequest {
5491 provider: LOCAL_PROVIDER_KEY.into(),
5492 method: ProxyMethod::Get,
5493 path: "/models".into(),
5494 body: None,
5495 stream: false,
5496 body_bytes: None,
5497 })
5498 .await
5499 .expect_err("oversize body must be rejected");
5500 let msg = format!("{err}");
5501 assert!(msg.contains("limit"), "got: {msg}");
5502 }
5503
5504 #[tokio::test]
5507 async fn local_broker_stream_error_is_status_only() {
5508 use axum::routing::post;
5509 let app = axum::Router::new().route(
5510 "/chat/completions",
5511 post(|| async {
5512 (
5513 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
5514 format!("\x1b[2J{}", "E".repeat(64 * 1024)),
5515 )
5516 }),
5517 );
5518 let url = spawn_upstream(app).await;
5519 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
5520 let err = broker
5521 .proxy_stream(ProxyRequest {
5522 provider: LOCAL_PROVIDER_KEY.into(),
5523 method: ProxyMethod::Post,
5524 path: "/chat/completions".into(),
5525 body: None,
5526 stream: true,
5527 body_bytes: None,
5528 })
5529 .await
5530 .err()
5531 .expect("non-2xx upstream must not yield a stream");
5532 let msg = format!("{err}");
5533 assert!(
5534 msg.len() < 700,
5535 "error must be bounded, got {} bytes",
5536 msg.len()
5537 );
5538 assert!(!msg.contains('\x1b'), "control chars must be stripped");
5539 assert!(msg.contains("500"));
5540 }
5541
5542 #[tokio::test]
5545 async fn local_broker_buffered_request_times_out() {
5546 use axum::routing::get;
5547 let app = axum::Router::new().route(
5548 "/models",
5549 get(|| async {
5550 tokio::time::sleep(Duration::from_secs(30)).await;
5551 "too late"
5552 }),
5553 );
5554 let url = spawn_upstream(app).await;
5555 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
5556 .with_request_timeout(Duration::from_millis(200));
5557 let started = std::time::Instant::now();
5558 let err = broker
5559 .proxy(ProxyRequest {
5560 provider: LOCAL_PROVIDER_KEY.into(),
5561 method: ProxyMethod::Get,
5562 path: "/models".into(),
5563 body: None,
5564 stream: false,
5565 body_bytes: None,
5566 })
5567 .await
5568 .expect_err("hung upstream must time out");
5569 assert!(matches!(err, BrokerError::Transport(_)));
5570 assert!(
5571 started.elapsed() < Duration::from_secs(5),
5572 "timeout must be enforced promptly"
5573 );
5574 }
5575
5576 #[tokio::test]
5579 async fn global_broker_defaults_to_in_process_local() {
5580 let broker = global_broker();
5581 let caps = broker
5582 .capabilities()
5583 .await
5584 .expect("in-process broker needs no daemon");
5585 assert!(caps.iter().any(|c| c.key == LOCAL_PROVIDER_KEY));
5586 let again = global_broker();
5588 assert!(again.capabilities().await.is_ok());
5589 }
5590
5591 const HOSTILE_SENTINEL: &str = "ZX9-HOSTILE-SENTINEL-7Q";
5601
5602 fn hostile_body() -> String {
5605 format!(
5606 "{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} \
5607 {{\\\"messages\\\":[{{\\\"content\\\":\\\"my secret prompt\\\"}}],\
5608 \\\"tools\\\":[{{\\\"input_schema\\\":{{}}}}]}}\"}}}}"
5609 )
5610 }
5611
5612 fn hostile_broker_json() -> String {
5615 format!(
5616 "{{\"error\":\"ECHOED {HOSTILE_SENTINEL} input_schema \
5617 my secret prompt\"}}"
5618 )
5619 }
5620
5621 fn assert_no_hostile_bytes(msg: &str) {
5622 assert!(
5623 !msg.contains(HOSTILE_SENTINEL),
5624 "sentinel leaked into error: {msg}"
5625 );
5626 assert!(!msg.contains("ECHOED"), "echoed body leaked: {msg}");
5627 assert!(
5628 !msg.contains("input_schema"),
5629 "request-shaped data leaked: {msg}"
5630 );
5631 assert!(
5632 !msg.contains("secret prompt"),
5633 "prompt content leaked: {msg}"
5634 );
5635 }
5636
5637 #[tokio::test]
5642 async fn local_broker_stream_error_omits_upstream_body() {
5643 use axum::routing::post;
5644 let app = axum::Router::new().route(
5645 "/chat/completions",
5646 post(|| async {
5647 (
5648 axum::http::StatusCode::TOO_MANY_REQUESTS,
5649 format!("\x1b[2J{}", hostile_body()),
5650 )
5651 }),
5652 );
5653 let url = spawn_upstream(app).await;
5654 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
5655 let err = broker
5656 .proxy_stream(ProxyRequest {
5657 provider: LOCAL_PROVIDER_KEY.into(),
5658 method: ProxyMethod::Post,
5659 path: "/chat/completions".into(),
5660 body: None,
5661 stream: true,
5662 body_bytes: None,
5663 })
5664 .await
5665 .err()
5666 .expect("non-2xx upstream must not yield a stream");
5667 assert!(matches!(err, BrokerError::Transport(_)));
5668 let msg = format!("{err}");
5669 assert!(
5670 msg.contains("provider request failed: 429 Too Many Requests"),
5671 "status prefix must survive for engine classification: {msg}"
5672 );
5673 assert!(
5674 !msg.contains('\x1b'),
5675 "escapes must not reach errors: {msg}"
5676 );
5677 assert_no_hostile_bytes(&msg);
5678 }
5679
5680 async fn stream_error_message(status: axum::http::StatusCode, body: String) -> String {
5683 use axum::routing::post;
5684 let app = axum::Router::new().route(
5685 "/chat/completions",
5686 post(move || async move { (status, body) }),
5687 );
5688 let url = spawn_upstream(app).await;
5689 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
5690 let err = broker
5691 .proxy_stream(ProxyRequest {
5692 provider: LOCAL_PROVIDER_KEY.into(),
5693 method: ProxyMethod::Post,
5694 path: "/chat/completions".into(),
5695 body: None,
5696 stream: true,
5697 body_bytes: None,
5698 })
5699 .await
5700 .err()
5701 .expect("non-2xx upstream must not yield a stream");
5702 assert!(matches!(err, BrokerError::Transport(_)));
5703 format!("{err}")
5704 }
5705
5706 fn after_marker(msg: &str) -> &str {
5708 let marker = "provider request failed: ";
5709 let idx = msg.find(marker).expect("marker present");
5710 &msg[idx + marker.len()..]
5711 }
5712
5713 #[tokio::test]
5719 async fn local_broker_stream_403_kimi_quota_surfaces_vetted_label_only() {
5720 let kimi_body = "{\"error\":{\"type\":\"access_terminated_error\",\"message\":\
5721 \"You've reached your weekly (7-day) usage limit. Your quota will \
5722 reset when the current 7-day window ends. \"}}";
5723 let msg = stream_error_message(axum::http::StatusCode::FORBIDDEN, kimi_body.into()).await;
5724 assert!(
5725 msg.ends_with("provider request failed: 403 Forbidden [access_terminated_error]"),
5726 "expected exact labelled prefix: {msg}"
5727 );
5728 assert!(!msg.contains("weekly"), "body text leaked: {msg}");
5729 assert!(!msg.contains("quota"), "body text leaked: {msg}");
5730 assert!(!msg.contains('{'), "raw JSON leaked: {msg}");
5731 assert!(
5732 !after_marker(&msg).contains(':'),
5733 "suffix must not contain ':' (engine truncates there): {msg}"
5734 );
5735 }
5736
5737 #[tokio::test]
5740 async fn local_broker_stream_429_unvetted_type_yields_no_label() {
5741 let body = format!(
5742 "{{\"error\":{{\"type\":\"<script>ECHO:secret\",\"message\":\"ECHOED {HOSTILE_SENTINEL}\"}}}}"
5743 );
5744 let msg = stream_error_message(axum::http::StatusCode::TOO_MANY_REQUESTS, body).await;
5745 assert!(
5746 msg.ends_with("provider request failed: 429 Too Many Requests"),
5747 "no label expected: {msg}"
5748 );
5749 assert!(!msg.contains('['), "unexpected label: {msg}");
5750 assert!(!msg.contains("script"), "hostile type leaked: {msg}");
5751 assert!(!msg.contains("secret"), "hostile type leaked: {msg}");
5752 assert!(!after_marker(&msg).contains(':'), "got: {msg}");
5753 assert_no_hostile_bytes(&msg);
5754 }
5755
5756 #[tokio::test]
5758 async fn local_broker_stream_html_body_yields_no_label() {
5759 let body =
5760 format!("<html><body>ECHOED {HOSTILE_SENTINEL} access_terminated_error</body></html>");
5761 let msg = stream_error_message(axum::http::StatusCode::BAD_GATEWAY, body).await;
5762 assert!(
5763 msg.ends_with("provider request failed: 502 Bad Gateway"),
5764 "no label expected: {msg}"
5765 );
5766 assert!(!msg.contains("html"), "body leaked: {msg}");
5767 assert!(!msg.contains('['), "unexpected label: {msg}");
5768 assert_no_hostile_bytes(&msg);
5769 }
5770
5771 #[tokio::test]
5774 async fn local_broker_stream_oversized_body_yields_no_label_without_panic() {
5775 let pad = "x".repeat(MAX_PROXY_ERROR_CLASSIFY_BYTES * 4);
5778 let body = format!(
5779 "{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} {pad}\",\"type\":\"access_terminated_error\"}}}}"
5780 );
5781 let msg = stream_error_message(axum::http::StatusCode::FORBIDDEN, body).await;
5782 assert!(
5783 msg.ends_with("provider request failed: 403 Forbidden"),
5784 "no label expected for oversized body: {msg}"
5785 );
5786 assert!(!msg.contains('['), "unexpected label: {msg}");
5787 assert!(
5788 msg.len() < 200,
5789 "oversized body leaked: {} bytes",
5790 msg.len()
5791 );
5792 assert_no_hostile_bytes(&msg);
5793 }
5794
5795 #[tokio::test]
5799 async fn read_error_body_for_classification_truncates_and_never_fails() {
5800 use axum::routing::get;
5801 let app = axum::Router::new()
5802 .route("/big", get(|| async { "a".repeat(100) }))
5803 .route("/utf8", get(|| async { format!("{}é", "a".repeat(9)) }));
5804 let url = spawn_upstream(app).await;
5805 let client = reqwest::Client::new();
5806
5807 let resp = client.get(format!("{url}/big")).send().await.unwrap();
5808 let got = read_error_body_for_classification(resp, 10).await;
5809 assert_eq!(got.as_deref(), Some("aaaaaaaaaa"));
5810
5811 let resp = client.get(format!("{url}/utf8")).send().await.unwrap();
5813 let got = read_error_body_for_classification(resp, 10).await;
5814 assert_eq!(got, None);
5815
5816 let resp = client.get(format!("{url}/utf8")).send().await.unwrap();
5818 let got = read_error_body_for_classification(resp, 1024).await;
5819 assert_eq!(got.as_deref(), Some("aaaaaaaaaé"));
5820 }
5821
5822 #[tokio::test]
5826 #[serial_test::serial]
5827 async fn local_broker_usage_error_omits_upstream_body() {
5828 use axum::routing::get;
5829 let app = axum::Router::new().route(
5830 "/usage",
5831 get(|| async {
5832 (
5833 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
5834 hostile_body(),
5835 )
5836 }),
5837 );
5838 let url = spawn_upstream(app).await;
5839
5840 let dir = tempfile::tempdir().unwrap();
5846 let base = dir.path().join(".synaps-cli");
5847 std::fs::create_dir_all(&base).unwrap();
5848 let creds = crate::auth::OAuthCredentials {
5849 auth_type: "oauth".into(),
5850 refresh: String::new(),
5851 access: "test-access-token".into(),
5852 expires: crate::epoch_millis() + 3_600_000,
5853 account_id: None,
5854 };
5855 super::super::storage::save_provider_auth_at_test_hook(
5856 &base.join("auth.json"),
5857 "anthropic",
5858 &creds,
5859 )
5860 .unwrap();
5861 let old_base = std::env::var("SYNAPS_BASE_DIR").ok();
5862 let old_home = std::env::var("HOME").ok();
5863 std::env::remove_var("SYNAPS_BASE_DIR");
5864 std::env::set_var("HOME", dir.path());
5865
5866 let broker = LocalBroker::new(reqwest::Client::new())
5867 .with_anthropic_usage_url(format!("{url}/usage"));
5868 let result = broker.anthropic_usage().await;
5869
5870 match old_home {
5871 Some(v) => std::env::set_var("HOME", v),
5872 None => std::env::remove_var("HOME"),
5873 }
5874 if let Some(v) = old_base {
5875 std::env::set_var("SYNAPS_BASE_DIR", v);
5876 }
5877
5878 let err = result.expect_err("non-2xx usage response must fail");
5879 assert!(matches!(err, BrokerError::Transport(_)));
5880 let msg = format!("{err}");
5881 assert!(
5882 msg.contains("usage request failed") && msg.contains("500"),
5883 "operation + status must survive: {msg}"
5884 );
5885 assert_no_hostile_bytes(&msg);
5886 }
5887
5888 async fn spawn_remote_broker_error(status: u16) -> RemoteBroker {
5889 use axum::routing::{get, post};
5890 let code = axum::http::StatusCode::from_u16(status).unwrap();
5891 let app = axum::Router::new()
5892 .route(
5893 "/proxy",
5894 post(move || async move { (code, hostile_broker_json()) }),
5895 )
5896 .route(
5897 "/usage",
5898 get(move || async move { (code, hostile_broker_json()) }),
5899 );
5900 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
5901 let addr = listener.local_addr().unwrap();
5902 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
5903 RemoteBroker::new(
5904 format!("http://{addr}"),
5905 "opaque-machine-token",
5906 reqwest::Client::new(),
5907 super::super::TokenCache::new(),
5908 )
5909 }
5910
5911 fn groq_request(stream: bool) -> ProxyRequest {
5912 ProxyRequest {
5913 provider: "groq".into(),
5914 method: ProxyMethod::Post,
5915 path: "/chat/completions".into(),
5916 body: None,
5917 stream,
5918 body_bytes: None,
5919 }
5920 }
5921
5922 #[tokio::test]
5925 async fn remote_broker_denied_omits_broker_error_body() {
5926 for status in [400u16, 403] {
5927 let broker = spawn_remote_broker_error(status).await;
5928 let err = broker
5929 .proxy(groq_request(false))
5930 .await
5931 .expect_err("4xx broker response must fail");
5932 assert!(
5933 matches!(err, BrokerError::Denied(_)),
5934 "typed category must survive, got {err:?}"
5935 );
5936 let msg = format!("{err}");
5937 assert!(msg.contains(&status.to_string()), "status lost: {msg}");
5938 assert_no_hostile_bytes(&msg);
5939 }
5940 }
5941
5942 #[test]
5945 fn gemini_gets_the_longer_stream_idle_timeout() {
5946 let broker = LocalBroker::new(reqwest::Client::new());
5947 assert_eq!(broker.stream_idle_timeout_for("google-gemini"), Duration::from_secs(300));
5948 for key in ["muse", "kimi-code", "openai-codex", "anthropic"] {
5949 assert_eq!(broker.stream_idle_timeout_for(key), Duration::from_secs(90), "{key}");
5950 }
5951 let short = LocalBroker::new(reqwest::Client::new())
5952 .with_stream_idle_timeout_for_tests(Duration::from_millis(10));
5953 assert_eq!(short.stream_idle_timeout_for("muse"), Duration::from_millis(10));
5954 assert!(short.stream_idle_timeout_for("google-gemini") >= Duration::from_secs(300));
5955 }
5956
5957 async fn spawn_remote_broker_body(status: u16, body: String) -> RemoteBroker {
5958 use axum::routing::post;
5959 let code = axum::http::StatusCode::from_u16(status).unwrap();
5960 let app = axum::Router::new().route(
5961 "/proxy",
5962 post(move || {
5963 let body = body.clone();
5964 async move { (code, body) }
5965 }),
5966 );
5967 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
5968 let addr = listener.local_addr().unwrap();
5969 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
5970 RemoteBroker::new(
5971 format!("http://{addr}"),
5972 "opaque-machine-token",
5973 reqwest::Client::new(),
5974 super::super::TokenCache::new(),
5975 )
5976 }
5977
5978 #[tokio::test]
5984 async fn remote_broker_502_keeps_a_retryable_refresh_failure() {
5985 use super::super::token::{RefreshFailure, RefreshFailureKind as K};
5986 for kind in [K::Transient, K::LockBusy, K::Other] {
5987 let daemon_text = BrokerError::Refresh(RefreshFailure::new(
5988 kind,
5989 format!("HTTP error 503 {HOSTILE_SENTINEL} ECHOED input_schema"),
5990 ))
5991 .to_string();
5992 let body = serde_json::json!({ "error": daemon_text }).to_string();
5993 let broker = spawn_remote_broker_body(502, body.clone()).await;
5994 let stream_err = broker.proxy_stream(groq_request(true)).await.err().unwrap();
5995 let buffered_err = broker.proxy(groq_request(false)).await.unwrap_err();
5996 for err in [stream_err, buffered_err] {
5997 let BrokerError::Refresh(f) = &err else {
5998 panic!("{kind:?}: expected Refresh, got {err:?}");
5999 };
6000 assert_eq!(f.kind, kind);
6001 let msg = err.to_string();
6002 assert!(msg.contains("retry shortly"), "{msg}");
6003 assert_no_hostile_bytes(&msg);
6004 }
6005 }
6006 let broker = spawn_remote_broker_body(502, hostile_broker_json()).await;
6008 let err = broker.proxy_stream(groq_request(true)).await.err().unwrap();
6009 assert!(matches!(err, BrokerError::Transport(_)), "{err:?}");
6010 assert_no_hostile_bytes(&err.to_string());
6011 }
6012
6013 #[tokio::test]
6016 async fn remote_broker_stream_error_omits_broker_body() {
6017 let broker = spawn_remote_broker_error(500).await;
6018 let err = broker
6019 .proxy_stream(groq_request(true))
6020 .await
6021 .err()
6022 .expect("non-2xx broker response must not yield a stream");
6023 assert!(matches!(err, BrokerError::Transport(_)));
6024 let msg = format!("{err}");
6025 assert!(
6026 msg.contains("broker proxy stream failed") && msg.contains("500"),
6027 "operation + status must survive: {msg}"
6028 );
6029 assert_no_hostile_bytes(&msg);
6030 }
6031
6032 #[tokio::test]
6035 async fn remote_broker_usage_error_omits_broker_body() {
6036 let broker = spawn_remote_broker_error(502).await;
6037 let err = broker
6038 .anthropic_usage()
6039 .await
6040 .expect_err("non-2xx broker response must fail");
6041 assert!(matches!(err, BrokerError::Transport(_)));
6042 let msg = format!("{err}");
6043 assert!(
6044 msg.contains("broker usage returned HTTP 502"),
6045 "operation + status must survive: {msg}"
6046 );
6047 assert_no_hostile_bytes(&msg);
6048 }
6049
6050 #[tokio::test]
6053 async fn capabilities_expose_no_secret_fields() {
6054 let broker = LocalBroker::new(reqwest::Client::new());
6055 let caps = broker.capabilities().await.unwrap();
6056 assert!(caps
6057 .iter()
6058 .any(|c| c.key == "anthropic" && c.kind == CredentialKind::OAuth));
6059 assert!(caps
6060 .iter()
6061 .any(|c| c.key == "groq" && c.kind == CredentialKind::StaticKey));
6062 let json = serde_json::to_string(&caps).unwrap();
6063 for field in ["key\":", "name\":", "kind\":", "configured\":"] {
6064 assert!(json.contains(field));
6065 }
6066 assert!(!json.contains("refresh"));
6067 assert!(!json.contains("access"));
6068 }
6069
6070 #[test]
6074 fn auto_vend_failure_names_the_chosen_account_and_its_login_command() {
6075 use super::super::token::{RefreshFailure, RefreshFailureKind as R};
6076 let named = CredentialRef::new(
6077 OAuthProviderId::Anthropic,
6078 Account::named("claude3").unwrap(),
6079 );
6080 let refused =
6081 auto_vend_failure(&named, RefreshFailure::new(R::Credential, "invalid_grant"));
6082 assert_eq!(refused.kind, R::Credential);
6083 assert_eq!(
6084 refused.message(),
6085 "account 'claude3': invalid_grant; log it in again with \
6086 `synaps login --provider anthropic --account claude3`"
6087 );
6088 let default = CredentialRef::new(OAuthProviderId::Anthropic, Account::Default);
6089 let refused = auto_vend_failure(&default, RefreshFailure::new(R::Credential, "revoked"));
6090 assert_eq!(
6091 refused.message(),
6092 "account 'default': revoked; log it in again with \
6093 `synaps login --provider anthropic`"
6094 );
6095 let advised = auto_vend_failure(
6097 &named,
6098 RefreshFailure::new(
6099 R::Credential,
6100 "slot removed; run `synaps login --account claude3`",
6101 ),
6102 );
6103 assert_eq!(advised.message().matches("synaps login").count(), 1);
6104 for kind in [R::Transient, R::LockBusy, R::Other, R::ActionRequired] {
6106 let f = auto_vend_failure(&named, RefreshFailure::new(kind, "HTTP 503"));
6107 assert_eq!(f.kind, kind);
6108 assert_eq!(f.message(), "account 'claude3': HTTP 503");
6109 }
6110 }
6111
6112 #[test]
6117 fn refresh_failure_kinds_map_to_observation_kinds() {
6118 use super::super::quota_policy::{AccountFailureKind as K, QuotaObservation};
6119 use super::super::token::{RefreshFailure, RefreshFailureKind as R};
6120 for (kind, expected, transient, login) in [
6121 (R::Transient, K::RefreshTransient, true, false),
6122 (R::LockBusy, K::RefreshLockBusy, true, false),
6123 (R::Other, K::RefreshOther, true, false),
6124 (
6125 R::Cooldown { retry_in_secs: 42 },
6126 K::RefreshOther,
6127 true,
6128 false,
6129 ),
6130 (R::Credential, K::RefreshRejected, false, true),
6131 (R::ActionRequired, K::RefreshActionRequired, false, false),
6132 ] {
6133 let read = UsageReadError::refresh(RefreshFailure::new(kind, "why"));
6134 assert_eq!(read.failure.kind, expected, "{kind:?}");
6135 assert_eq!(read.failure.kind.is_transient(), transient, "{kind:?}");
6136 assert_eq!(read.failure.kind.needs_login(), login, "{kind:?}");
6137 assert_eq!(
6139 read.error,
6140 BrokerError::Credential("why".into()),
6141 "{kind:?}"
6142 );
6143 let observation = read.failure.observation();
6144 if login {
6145 assert_eq!(observation, QuotaObservation::AuthError);
6146 } else {
6147 assert_eq!(observation, QuotaObservation::Unknown);
6148 }
6149 }
6150 let cooldown =
6151 refresh_failure(&RefreshFailure::new(R::Cooldown { retry_in_secs: 42 }, "x"));
6152 assert_eq!(cooldown.retry_after_s, Some(42));
6153 let gone = UsageReadError::slot(BrokerError::UnknownAccount {
6155 provider: "anthropic".into(),
6156 label: "x".into(),
6157 });
6158 assert_eq!(gone.failure.kind, K::CredentialMissing);
6159 assert!(gone.failure.kind.needs_login());
6160 let store = UsageReadError::slot(BrokerError::Credential("unreadable".into()));
6161 assert_eq!(store.failure.kind, K::CredentialStore);
6162 assert!(!store.failure.kind.needs_login() && !store.failure.kind.is_transient());
6163 }
6164
6165 #[test]
6168 fn usage_errors_map_to_observation_kinds_with_status() {
6169 use super::super::quota_policy::AccountFailureKind as K;
6170 use super::super::usage::UsageError as U;
6171 for (err, kind, status) in [
6172 (
6173 U::RateLimited {
6174 retry_after_s: Some(9),
6175 },
6176 K::UsageRateLimited,
6177 Some(429),
6178 ),
6179 (
6180 U::Unauthorized { status: 403 },
6181 K::UsageUnauthorized,
6182 Some(403),
6183 ),
6184 (
6185 U::UpstreamStatus {
6186 status: 503,
6187 retry_after_s: None,
6188 },
6189 K::UsageUpstream,
6190 Some(503),
6191 ),
6192 (U::Redirected { status: 302 }, K::UsageUpstream, Some(302)),
6193 (U::Timeout, K::UsageTimeout, None),
6194 (
6195 U::Transport {
6196 detail: "connect".into(),
6197 },
6198 K::UsageTransport,
6199 None,
6200 ),
6201 (U::Malformed { detail: "x".into() }, K::UsageMalformed, None),
6202 (U::BodyTooLarge { cap: 1 }, K::UsageMalformed, None),
6203 (U::AccountIdUnavailable, K::UsageUnauthorized, None),
6204 (
6205 U::UnsupportedProvider {
6206 provider: "p".into(),
6207 },
6208 K::Unsupported,
6209 None,
6210 ),
6211 ] {
6212 let read = UsageReadError::usage(err.clone());
6213 assert_eq!(read.failure.kind, kind, "{err:?}");
6214 assert_eq!(read.failure.http_status, status, "{err:?}");
6215 assert_eq!(read.error, err.clone().into_broker_error(), "{err:?}");
6216 }
6217 assert_eq!(
6218 UsageReadError::usage(U::RateLimited {
6219 retry_after_s: Some(9)
6220 })
6221 .failure
6222 .retry_after_s,
6223 Some(9)
6224 );
6225 }
6226}