use std::collections::BTreeMap;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use futures::Stream;
use serde::{Deserialize, Serialize};
use super::account::{
seat_fingerprint, Account, AccountPolicy, AccountSelector, AccountSummary, CredentialRef,
SeatIdentity,
};
use super::cloud::{CloudProviderId, InvokeRequest};
use super::provider::OAuthProviderId;
use super::static_providers::{
allowed_proxy_paths, static_provider, StaticProviderSpec, LOCAL_DEFAULT_BASE_URL,
LOCAL_PROVIDER_KEY, STATIC_PROVIDERS,
};
use super::storage;
pub const MAX_PROXY_REQUEST_BYTES: usize = 2 * 1024 * 1024;
pub const MAX_PROXY_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_UPSTREAM_ERROR_BYTES: usize = 2 * 1024;
pub const PROXY_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
const ANTHROPIC_USAGE_URL: &str = "https://api.anthropic.com/api/oauth/usage";
const ANTHROPIC_OAUTH_BETA: &str = "oauth-2025-04-20";
pub const GOOGLE_GEMINI_CODE_ASSIST_BASE_URL: &str = "https://cloudcode-pa.googleapis.com";
pub(crate) fn is_allowed_google_gemini_path(path: &str) -> bool {
matches!(
path,
"/v1internal:loadCodeAssist"
| "/v1internal:onboardUser"
| "/v1internal:streamGenerateContent"
| "/v1internal:countTokens"
) || path.starts_with("/v1internal/operations/")
}
pub const OPENAI_CODEX_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
pub(crate) const OPENAI_CODEX_MODELS_PATH: &str = "/codex/models";
pub(crate) fn is_allowed_openai_codex_path(path: &str) -> bool {
let Some((base, query)) = path.split_once('?') else {
return false;
};
if base != OPENAI_CODEX_MODELS_PATH {
return false;
}
let mut client_version: Option<&str> = None;
for pair in query.split('&') {
let Some((k, v)) = pair.split_once('=') else {
return false;
};
if k != "client_version" || client_version.is_some() {
return false;
}
if v.trim().is_empty() {
return false;
}
client_version = Some(v);
}
client_version.is_some()
}
pub(crate) fn is_allowed_anthropic_path(path: &str) -> bool {
if path == "/v1/models" {
return true;
}
let Some((base, query)) = path.split_once('?') else {
return false;
};
if base != "/v1/models" {
return false;
}
if query.is_empty() {
return false;
}
let mut saw_limit = false;
let mut saw_after_id = false;
for pair in query.split('&') {
let Some((k, v)) = pair.split_once('=') else {
return false;
};
if v.is_empty() {
return false;
}
match k {
"limit" if !saw_limit => saw_limit = true,
"after_id" if !saw_after_id => saw_after_id = true,
_ => return false,
}
}
true
}
pub(crate) fn anthropic_oauth_catalog_request_headers() -> &'static [(&'static str, &'static str)] {
&[
("anthropic-version", "2023-06-01"),
("anthropic-beta", ANTHROPIC_OAUTH_BETA),
("accept", "application/json"),
]
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum BrokerError {
UnknownProvider(String),
RegistrationRequired {
provider: String,
remediation: String,
},
NotConfigured(String),
Unauthorized,
Denied(String),
UnsupportedCapability {
provider: String,
capability: String,
},
Transport(String),
Credential(String),
UnknownAccount { provider: String, label: String },
UnsupportedAccount { provider: String, label: String },
InvalidAccount(String),
NoAccountAvailable {
provider: String,
reason: String,
failure: Option<super::quota_policy::SelectionFailure>,
},
BrokerTooOld(String),
Refresh(super::token::RefreshFailure),
}
impl std::fmt::Display for BrokerError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::UnknownProvider(p) => write!(f, "unknown provider: {p}"),
Self::RegistrationRequired {
provider,
remediation,
} => {
write!(f, "registration required for '{provider}': {remediation}")
}
Self::NotConfigured(p) => write!(
f,
"no credential configured for '{p}'. Run `synaps login` to add one."
),
Self::Unauthorized => f.write_str(&super::broker_token::unauthorized_hint()),
Self::Denied(msg) => write!(f, "broker denied request: {msg}"),
Self::UnsupportedCapability {
provider,
capability,
} => write!(
f,
"cloud provider '{provider}' is text-only: {capability} are not supported yet \
(no credential was used and no network request was made)"
),
Self::Transport(msg) => write!(f, "broker transport error: {msg}"),
Self::Credential(msg) => write!(f, "credential error: {msg}"),
Self::Refresh(failure) => {
use super::token::RefreshFailureKind as K;
let hint = match failure.kind {
K::Transient => " (network or provider error, not a login problem; retry shortly)",
K::LockBusy => " (another refresh of this account is still running; retry shortly)",
K::Other => " (retry shortly)",
K::Credential | K::Cooldown { .. } | K::ActionRequired => "",
};
write!(f, "credential refresh failed{hint}: {failure}")
}
Self::UnknownAccount { provider, label } => write!(
f,
"unknown account '{label}' for provider '{provider}' (no fallback; run \
`synaps auth list` or `synaps login --provider {provider} --account {label}`)"
),
Self::UnsupportedAccount { provider, label } => write!(
f,
"this credential broker cannot address account '{label}' for provider '{provider}'"
),
Self::InvalidAccount(msg) => write!(f, "invalid account: {msg}"),
Self::NoAccountAvailable {
provider,
reason,
failure: Some(failure),
} => write!(
f,
"no account with proven capacity for '{provider}' ({}): {reason}",
failure.kind
),
Self::NoAccountAvailable {
provider, reason, ..
} => {
write!(
f,
"no account with proven capacity for '{provider}': {reason}"
)
}
Self::BrokerTooOld(what) => write!(
f,
"broker is older than this client; {what} unavailable (upgrade the broker host)"
),
}
}
}
impl std::error::Error for BrokerError {}
#[cfg(test)]
#[path = "broker_remote_compat_tests.rs"]
mod remote_compat_tests;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccessToken {
pub token: String,
pub expires: u64,
}
#[derive(Debug, Clone)]
pub struct PinnedToken {
pub credential: CredentialRef,
pub token: AccessToken,
}
impl PinnedToken {
pub fn seat_fingerprint(&self) -> Option<String> {
match self.credential.provider {
OAuthProviderId::OpenAiCodex => super::extract_codex_account_id(&self.token.token)
.and_then(|id| seat_fingerprint(self.credential.provider, &id)),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ProxyMethod {
Get,
Post,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProxyRequest {
pub provider: String,
pub method: ProxyMethod,
pub path: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body: Option<serde_json::Value>,
#[serde(default)]
pub stream: bool,
#[serde(skip)]
pub body_bytes: Option<bytes::Bytes>,
}
impl ProxyRequest {
pub fn validate(&self) -> Result<(), BrokerError> {
let (provider, _account) = self.split_provider()?;
let provider = provider.as_str();
if provider != LOCAL_PROVIDER_KEY
&& provider != "xai-auth"
&& provider != "github-copilot"
&& provider != "google-gemini"
&& provider != "openai-codex"
&& provider != "anthropic"
&& provider != "kimi-code"
&& provider != "muse"
&& static_provider(provider).is_none()
{
return Err(BrokerError::UnknownProvider(self.provider.clone()));
}
if !self.path.starts_with('/') {
return Err(BrokerError::Denied("proxy path must be relative".into()));
}
if self.path.contains("://") || self.path.contains("..") {
return Err(BrokerError::Denied(
"proxy path is not a plain relative path".into(),
));
}
let oauth_path_allowed = provider == "xai-auth" && self.path == "/responses"
|| provider == "github-copilot"
&& matches!(
self.path.as_str(),
"/models" | "/chat/completions" | "/responses"
)
|| provider == "google-gemini" && is_allowed_google_gemini_path(&self.path)
|| provider == "kimi-code"
&& (self.path == "/chat/completions"
|| self.method == ProxyMethod::Get
&& matches!(self.path.as_str(), "/models" | "/me" | "/usages"))
|| provider == "muse"
&& (self.method == ProxyMethod::Post && self.path == "/v1/responses"
|| self.method == ProxyMethod::Get && self.path == "/muse-code/models")
|| provider == "openai-codex"
&& self.method == ProxyMethod::Get
&& is_allowed_openai_codex_path(&self.path)
|| provider == "anthropic"
&& self.method == ProxyMethod::Get
&& is_allowed_anthropic_path(&self.path);
if !oauth_path_allowed && !allowed_proxy_paths(provider).contains(&self.path.as_str()) {
return Err(BrokerError::Denied(format!(
"proxy path '{}' is not in the provider's endpoint allowlist",
self.path
)));
}
let size = match (&self.body_bytes, &self.body) {
(Some(bytes), _) => Some(bytes.len()),
(None, Some(body)) => Some(
serde_json::to_vec(body)
.map(|v| v.len())
.unwrap_or(usize::MAX),
),
(None, None) => None,
};
if let Some(size) = size {
if size > MAX_PROXY_REQUEST_BYTES {
return Err(BrokerError::Denied(format!(
"request body exceeds the {MAX_PROXY_REQUEST_BYTES}-byte broker limit"
)));
}
}
#[cfg(any(test, debug_assertions))]
if let (Some(bytes), Some(body)) = (&self.body_bytes, &self.body) {
let parsed = serde_json::from_slice::<serde_json::Value>(bytes).ok();
if parsed.as_ref() != Some(body) {
return Err(BrokerError::Denied(
"body_bytes does not match the JSON body (exact-byte handoff incoherent)"
.into(),
));
}
}
Ok(())
}
pub fn post_json_exact(
provider: impl Into<String>,
path: impl Into<String>,
body: serde_json::Value,
stream: bool,
) -> Result<(Self, bytes::Bytes), BrokerError> {
let bytes =
bytes::Bytes::from(serde_json::to_vec(&body).map_err(|e| {
BrokerError::Denied(format!("request body failed to serialize: {e}"))
})?);
let request = Self {
provider: provider.into(),
method: ProxyMethod::Post,
path: path.into(),
body: Some(body),
stream,
body_bytes: Some(bytes.clone()),
};
Ok((request, bytes))
}
pub fn with_account(mut self, account: &Account) -> Self {
let base = self
.provider
.split_once('@')
.map(|(p, _)| p.to_string())
.unwrap_or_else(|| self.provider.clone());
self.provider = format!("{base}@{}", account.label_str());
self
}
pub fn split_provider(&self) -> Result<(String, Option<Account>), BrokerError> {
match self.provider.split_once('@') {
None => Ok((self.provider.clone(), None)),
Some((base, account)) => {
if base.parse::<OAuthProviderId>().is_err() {
return Err(BrokerError::UnknownProvider(self.provider.clone()));
}
let account = Account::parse(account).map_err(BrokerError::InvalidAccount)?;
Ok((base.to_string(), Some(account)))
}
}
}
pub fn base_provider(&self) -> Result<String, BrokerError> {
Ok(self.split_provider()?.0)
}
pub fn explicit_account(&self) -> Result<Option<Account>, BrokerError> {
Ok(self.split_provider()?.1)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProxyResponse {
pub status: u16,
pub body: String,
}
pub type ProxyByteStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, BrokerError>> + Send>>;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CredentialKind {
OAuth,
StaticKey,
LocalEndpoint,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderStatus {
pub key: String,
pub name: String,
pub kind: CredentialKind,
pub configured: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub accounts: Vec<AccountSummary>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StaticKeyStatus {
NotSet,
Configured {
masked: String,
},
FromEnv,
}
impl StaticKeyStatus {
pub fn is_configured(&self) -> bool {
!matches!(self, StaticKeyStatus::NotSet)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CloudCatalogEntry {
pub provider: CloudProviderId,
pub id: String,
pub display_name: String,
pub context_ref: String,
pub context_label: String,
pub stale: bool,
pub fetched_at: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CloudEvent {
TextDelta {
delta: String,
},
ToolArguments {
id: String,
name: Option<String>,
delta: String,
},
Usage {
input_tokens: u64,
output_tokens: u64,
},
Done,
}
pub type CloudEventStream = Pin<Box<dyn Stream<Item = Result<CloudEvent, BrokerError>> + Send>>;
#[async_trait]
pub trait CredentialBroker: Send + Sync {
async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError>;
fn account_selector(&self, provider: OAuthProviderId) -> AccountSelector {
let _ = provider;
AccountSelector::Account(Account::Default)
}
async fn access_token_for(&self, cred: &CredentialRef) -> Result<AccessToken, BrokerError> {
match &cred.account {
Account::Default => self.access_token(cred.provider).await,
Account::Named(label) => Err(BrokerError::UnsupportedAccount {
provider: cred.provider.as_str().to_string(),
label: label.as_str().to_string(),
}),
}
}
async fn recover_rejected_token(
&self,
cred: &CredentialRef,
_rejected: &str,
) -> Result<AccessToken, BrokerError> {
self.access_token_for(cred).await
}
async fn access_token_pinned(
&self,
provider: OAuthProviderId,
) -> Result<PinnedToken, BrokerError> {
let token = self.access_token(provider).await?;
Ok(PinnedToken {
credential: CredentialRef::default_for(provider),
token,
})
}
async fn access_token_pinned_for(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
let _ = model;
self.access_token_pinned(provider).await
}
async fn accounts(&self, provider: OAuthProviderId) -> Result<Vec<AccountSummary>, BrokerError> {
Err(BrokerError::UnsupportedCapability {
provider: provider.as_str().to_string(),
capability: "accounts".into(),
})
}
async fn usage(&self, cred: &CredentialRef) -> Result<super::usage::UsageSnapshot, BrokerError> {
Err(BrokerError::UnsupportedCapability {
provider: cred.provider.as_str().to_string(),
capability: "usage".into(),
})
}
async fn report_cooldown(
&self,
cred: &CredentialRef,
until_ms: Option<u64>,
reason: &str,
) -> Result<(), BrokerError> {
let _ = (cred, until_ms, reason);
Ok(())
}
async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError>;
async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError>;
async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError>;
async fn cloud_catalog(
&self,
provider: CloudProviderId,
context_ref: &str,
allow_stale: bool,
) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
let _ = (context_ref, allow_stale);
Err(BrokerError::NotConfigured(provider.to_string()))
}
async fn cloud_invoke(
&self,
provider: CloudProviderId,
context_ref: &str,
model_id: &str,
request: InvokeRequest,
) -> Result<CloudEventStream, BrokerError> {
let _ = (context_ref, model_id, request);
Err(BrokerError::NotConfigured(provider.to_string()))
}
async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError>;
}
#[async_trait]
pub trait CloudBackend: Send + Sync {
async fn catalog(
&self,
provider: CloudProviderId,
context_ref: &str,
allow_stale: bool,
) -> Result<Vec<CloudCatalogEntry>, BrokerError>;
async fn invoke(
&self,
provider: CloudProviderId,
context_ref: &str,
model_id: &str,
request: InvokeRequest,
) -> Result<CloudEventStream, BrokerError>;
}
#[derive(Clone)]
struct ProductionCloudBackend {
http: reqwest::Client,
refresh_lock: Arc<tokio::sync::Mutex<()>>,
}
impl ProductionCloudBackend {
fn new(_http: reqwest::Client) -> Self {
let http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(Duration::from_secs(10))
.timeout(Duration::from_secs(90))
.build()
.expect("cloud HTTPS client configuration is valid");
Self {
http,
refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
fn state(&self, provider: CloudProviderId) -> Result<serde_json::Value, BrokerError> {
storage::load_cloud_state(provider.as_str())
.map_err(BrokerError::Credential)?
.ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))
}
fn context_ref(&self, provider: CloudProviderId) -> Result<String, BrokerError> {
use sha2::{Digest, Sha256};
let state = self.state(provider)?;
let public = serde_json::to_vec(&state["config"])
.map_err(|_| BrokerError::Credential("invalid cloud context".into()))?;
let digest = Sha256::digest([provider.as_str().as_bytes(), &public].concat());
Ok(format!(
"ctx-{}",
digest[..16]
.iter()
.map(|b| format!("{b:02x}"))
.collect::<String>()
))
}
fn validate_context(
&self,
provider: CloudProviderId,
supplied: &str,
) -> Result<String, BrokerError> {
let opaque = self.context_ref(provider)?;
if supplied == provider.as_str() || supplied == opaque {
Ok(opaque)
} else {
Err(BrokerError::Denied(
"cloud context does not match stored provider".into(),
))
}
}
async fn aws(
&self,
) -> Result<super::aws_bedrock::AwsBedrockBroker<super::aws_bedrock::AwsHttpApi>, BrokerError>
{
#[derive(Deserialize)]
struct State {
config: super::cloud::AwsBedrockConfig,
access_key: String,
secret_key: String,
session_token: String,
expires_at: u64,
registered_client: serde_json::Value,
sso_access_token: String,
sso_refresh_token: Option<String>,
sso_expires_at: u64,
}
let decode = |value| {
serde_json::from_value::<State>(value)
.map_err(|_| BrokerError::Credential("invalid AWS broker state".into()))
};
let mut raw = self.state(CloudProviderId::AwsBedrock)?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let mut s = decode(raw.clone())?;
if s.expires_at <= now + 60_000 {
let _guard = self.refresh_lock.lock().await;
raw = self.state(CloudProviderId::AwsBedrock)?;
s = decode(raw.clone())?;
if s.expires_at <= now + 60_000 {
let api =
super::aws_bedrock::AwsHttpApi::new(self.http.clone(), &s.config.sso_region);
if s.sso_expires_at <= now + 60_000 {
let refresh = s
.sso_refresh_token
.as_deref()
.filter(|v| !v.is_empty())
.ok_or_else(|| {
BrokerError::Credential("AWS SSO session expired; login again".into())
})?;
let client = super::aws_bedrock::RegisteredClient::new(
s.registered_client["id"].as_str().ok_or_else(|| {
BrokerError::Credential("invalid AWS client registration".into())
})?,
s.registered_client["secret"].as_str().ok_or_else(|| {
BrokerError::Credential("invalid AWS client registration".into())
})?,
s.registered_client["expires_at"].as_u64().ok_or_else(|| {
BrokerError::Credential("invalid AWS client registration".into())
})?,
);
if client.expires_at <= now / 1000 {
return Err(BrokerError::Credential(
"AWS client registration expired; login again".into(),
));
}
use super::aws_bedrock::AwsApi;
let token = api
.create_token(
&client,
&s.config.sso_region,
super::aws_bedrock::TokenGrant::RefreshToken(refresh),
)
.await
.map_err(|_| {
BrokerError::Credential("AWS SSO refresh rejected; login again".into())
})?;
raw["sso_access_token"] = token.access().into();
if let Some(rotated) = token.refresh() {
raw["sso_refresh_token"] = rotated.into();
}
raw["sso_expires_at"] = (now + token.expires_in * 1000).into();
storage::save_cloud_state("aws-bedrock", &raw)
.map_err(BrokerError::Credential)?;
s = decode(raw.clone())?;
}
use super::aws_bedrock::AwsApi;
let credentials = api
.get_role_credentials(
&s.config.sso_region,
&s.sso_access_token,
&s.config.account_id,
&s.config.role_name,
)
.await
.map_err(|_| {
BrokerError::Credential("AWS role refresh rejected; login again".into())
})?;
raw["access_key"] = credentials.access_key().into();
raw["secret_key"] = credentials.secret_key().into();
raw["session_token"] = credentials.session_token().into();
raw["expires_at"] = credentials.expires_at.into();
storage::save_cloud_state("aws-bedrock", &raw).map_err(BrokerError::Credential)?;
s = decode(raw)?;
}
}
let api = super::aws_bedrock::AwsHttpApi::new(self.http.clone(), &s.config.sso_region);
Ok(super::aws_bedrock::AwsBedrockBroker::from_credentials(
api,
s.config,
super::aws_bedrock::RoleCredentials::new(
s.access_key,
s.secret_key,
s.session_token,
s.expires_at,
),
))
}
async fn azure_request(
&self,
state: &mut serde_json::Value,
audience: super::azure_openai::AzureAudience,
) -> Result<String, BrokerError> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
let key = match audience {
super::azure_openai::AzureAudience::Arm => "arm",
super::azure_openai::AzureAudience::Inference => "inference",
};
if state[key]["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
return state[key]["access_token"]
.as_str()
.map(str::to_owned)
.ok_or_else(|| BrokerError::Credential("invalid Azure token state".into()));
}
let _refresh = self.refresh_lock.lock().await;
*state = self.state(CloudProviderId::AzureOpenAi)?;
if state[key]["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
return state[key]["access_token"]
.as_str()
.map(str::to_owned)
.ok_or_else(|| BrokerError::Credential("invalid Azure token state".into()));
}
let config: super::cloud::AzureOpenAiConfig =
serde_json::from_value(state["config"].clone())
.map_err(|_| BrokerError::Credential("invalid Azure broker state".into()))?;
let client_id = state["client_id"]
.as_str()
.ok_or_else(|| BrokerError::Credential("invalid Azure client registration".into()))?;
let refresh = state["refresh_token"]
.as_str()
.ok_or_else(|| BrokerError::Credential("invalid Azure refresh state".into()))?;
let reg = super::azure_openai::AzureRegistration::production(Some(client_id.into()))
.map_err(|e| BrokerError::Credential(e.to_string()))?;
let r = super::azure_openai::refresh_request(&config, ®, audience, refresh)
.map_err(|e| BrokerError::Credential(e.to_string()))?;
let response = self
.http
.post(r.url)
.form(&r.form)
.send()
.await
.map_err(|_| BrokerError::Transport("Azure refresh failed".into()))?;
if !response.status().is_success() {
return Err(BrokerError::Credential(
"Azure refresh rejected; login again".into(),
));
}
let wire: serde_json::Value = response
.json()
.await
.map_err(|_| BrokerError::Transport("invalid Azure token response".into()))?;
let access = wire["access_token"]
.as_str()
.filter(|v| !v.is_empty())
.ok_or_else(|| BrokerError::Transport("invalid Azure token response".into()))?
.to_owned();
if let Some(r) = wire["refresh_token"].as_str().filter(|v| !v.is_empty()) {
state["refresh_token"] = r.into();
}
state[key] = serde_json::json!({"access_token": access, "expires_at": now + wire["expires_in"].as_u64().unwrap_or(3600) * 1000});
storage::save_cloud_state("azure-openai", state).map_err(BrokerError::Credential)?;
Ok(access)
}
async fn vertex_request(&self, state: &mut serde_json::Value) -> Result<String, BrokerError> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
if state["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
return state["access_token"]
.as_str()
.map(str::to_owned)
.ok_or_else(|| BrokerError::Credential("invalid Vertex token state".into()));
}
let _refresh = self.refresh_lock.lock().await;
*state = self.state(CloudProviderId::GoogleVertex)?;
if state["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
return state["access_token"]
.as_str()
.map(str::to_owned)
.ok_or_else(|| BrokerError::Credential("invalid Vertex token state".into()));
}
let client_id = state["client_id"]
.as_str()
.ok_or_else(|| BrokerError::Credential("invalid Vertex registration".into()))?;
let refresh = state["refresh_token"]
.as_str()
.ok_or_else(|| BrokerError::Credential("invalid Vertex refresh state".into()))?;
let response = self
.http
.post(super::google_vertex::TOKEN_URL)
.form(&[
("client_id", client_id),
("grant_type", "refresh_token"),
("refresh_token", refresh),
])
.send()
.await
.map_err(|_| BrokerError::Transport("Vertex refresh failed".into()))?;
if !response.status().is_success() {
return Err(BrokerError::Credential(
"Vertex refresh rejected; login again".into(),
));
}
let wire: serde_json::Value = response
.json()
.await
.map_err(|_| BrokerError::Transport("invalid Vertex token response".into()))?;
let access = wire["access_token"]
.as_str()
.filter(|v| !v.is_empty())
.ok_or_else(|| BrokerError::Transport("invalid Vertex token response".into()))?
.to_owned();
if let Some(r) = wire["refresh_token"].as_str().filter(|v| !v.is_empty()) {
state["refresh_token"] = r.into();
}
state["access_token"] = access.clone().into();
state["expires_at"] = (now + wire["expires_in"].as_u64().unwrap_or(3600) * 1000).into();
storage::save_cloud_state("google-vertex", state).map_err(BrokerError::Credential)?;
Ok(access)
}
}
#[async_trait]
impl CloudBackend for ProductionCloudBackend {
async fn catalog(
&self,
provider: CloudProviderId,
context_ref: &str,
_allow_stale: bool,
) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
let opaque_context = self.validate_context(provider, context_ref)?;
let fetched_at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64;
match provider {
CloudProviderId::AwsBedrock => {
let broker = self.aws().await?;
let label = broker.public_context().bedrock_region;
broker
.catalog()
.await
.map_err(|_| {
BrokerError::Transport("AWS catalog failed (details redacted)".into())
})
.map(|xs| {
xs.into_iter()
.map(|x| CloudCatalogEntry {
provider,
id: x.id,
display_name: x.display_name,
context_ref: opaque_context.clone(),
context_label: label.clone(),
stale: false,
fetched_at,
})
.collect()
})
}
CloudProviderId::AzureOpenAi => {
let mut state = self.state(provider)?;
let config: super::cloud::AzureOpenAiConfig =
serde_json::from_value(state["config"].clone()).map_err(|_| {
BrokerError::Credential("invalid Azure broker state".into())
})?;
let token = self
.azure_request(&mut state, super::azure_openai::AzureAudience::Arm)
.await?;
let mut discovery =
super::azure_openai::DeploymentDiscovery::new(config.clone(), 20, 1000);
let mut url = Some(discovery.initial_url());
while let Some(next) = url {
let r = self
.http
.get(next)
.bearer_auth(&token)
.send()
.await
.map_err(|_| BrokerError::Transport("Azure catalog failed".into()))?;
if !r.status().is_success() {
return Err(BrokerError::Transport("Azure catalog rejected".into()));
}
let body = read_body_capped(r, MAX_CLOUD_CATALOG_BODY_BYTES)
.await
.map_err(|_| {
BrokerError::Transport("Azure catalog exceeded body limit".into())
})?;
url = discovery
.accept_page(&body)
.map_err(|_| BrokerError::Transport("invalid Azure catalog".into()))?;
}
discovery
.finish()
.map_err(|_| BrokerError::Transport("Azure catalog unavailable".into()))
.map(|xs| {
xs.into_iter()
.map(|x| CloudCatalogEntry {
provider,
id: x.id,
display_name: x.display_name,
context_ref: opaque_context.clone(),
context_label: format!(
"{}/{}",
config.resource_name, config.resource_group
),
stale: false,
fetched_at,
})
.collect()
})
}
CloudProviderId::GoogleVertex => {
let mut state = self.state(provider)?;
let token = self.vertex_request(&mut state).await?;
let config: super::cloud::GoogleVertexConfig =
serde_json::from_value(state["config"].clone()).map_err(|_| {
BrokerError::Credential("invalid Vertex broker state".into())
})?;
let mut page: Option<String> = None;
let mut out = Vec::new();
let mut seen = std::collections::HashSet::new();
for _ in 0..20 {
let mut url=format!("https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models",config.location,config.project_id,config.location);
if let Some(p) = &page {
let mut parsed = url::Url::parse(&url).map_err(|_| {
BrokerError::Transport("invalid Vertex catalog URL".into())
})?;
parsed.query_pairs_mut().append_pair("pageToken", p);
url = parsed.into();
}
let r = self
.http
.get(url)
.bearer_auth(&token)
.send()
.await
.map_err(|_| BrokerError::Transport("Vertex catalog failed".into()))?;
if !r.status().is_success() {
return Err(BrokerError::Transport("Vertex catalog rejected".into()));
}
let v = cloud_catalog_json(r).await?;
for m in v["publisherModels"].as_array().into_iter().flatten() {
if let Some(name) = m["name"]
.as_str()
.filter(|n| n.starts_with("publishers/google/models/"))
{
if out.len() >= MAX_CLOUD_CATALOG_ENTRIES {
return Err(BrokerError::Transport(
"Vertex catalog exceeded entry limit".into(),
));
}
if !out.iter().any(|entry: &CloudCatalogEntry| {
entry.id == format!("google-vertex/{name}")
}) {
out.push(CloudCatalogEntry {
provider,
id: format!("google-vertex/{name}"),
display_name: m["displayName"].as_str().unwrap_or(name).into(),
context_ref: opaque_context.clone(),
context_label: format!(
"{}/{}",
config.project_id, config.location
),
stale: false,
fetched_at,
});
}
}
}
page = v["nextPageToken"].as_str().map(str::to_owned);
match &page {
Some(p) if seen.insert(p.clone()) => {}
Some(_) => {
return Err(BrokerError::Transport("Vertex pagination loop".into()))
}
None => break,
}
}
if out.is_empty() {
Err(BrokerError::Transport("Vertex catalog is empty".into()))
} else {
Ok(out)
}
}
}
}
async fn invoke(
&self,
provider: CloudProviderId,
context_ref: &str,
model_id: &str,
request: InvokeRequest,
) -> Result<CloudEventStream, BrokerError> {
self.validate_context(provider, context_ref)?;
if !request.tools.is_empty() {
return Err(BrokerError::Denied(
"tools are not yet supported by cloud providers".into(),
));
}
let authorized = self.catalog(provider, context_ref, false).await?;
if !authorized.iter().any(|entry| entry.id == model_id) {
return Err(BrokerError::Denied(
"model is not present in the current provider catalog".into(),
));
}
match provider {
CloudProviderId::AwsBedrock => {
let broker = self.aws().await?;
if request.stream {
use futures::StreamExt;
let stream = broker
.converse_stream(model_id, request)
.await
.map_err(|_| {
BrokerError::Transport(
"AWS invocation failed (details redacted)".into(),
)
})?
.map(|event| {
event
.map_err(|_| {
BrokerError::Transport(
"AWS stream failed (details redacted)".into(),
)
})
.map(|e| match e {
super::aws_bedrock::ConverseEvent::TextDelta(delta) => {
CloudEvent::TextDelta { delta }
}
super::aws_bedrock::ConverseEvent::ToolArguments {
id,
delta,
} => CloudEvent::ToolArguments {
id,
name: None,
delta,
},
super::aws_bedrock::ConverseEvent::Usage(u) => {
CloudEvent::Usage {
input_tokens: u.input_tokens,
output_tokens: u.output_tokens,
}
}
super::aws_bedrock::ConverseEvent::Done => CloudEvent::Done,
})
});
Ok(Box::pin(stream))
} else {
let o = broker.converse(model_id, request).await.map_err(|_| {
BrokerError::Transport("AWS invocation failed (details redacted)".into())
})?;
Ok(Box::pin(futures::stream::iter(vec![
Ok(CloudEvent::TextDelta { delta: o.text }),
Ok(CloudEvent::Usage {
input_tokens: o.usage.input_tokens,
output_tokens: o.usage.output_tokens,
}),
Ok(CloudEvent::Done),
])))
}
}
CloudProviderId::AzureOpenAi => {
let mut state = self.state(provider)?;
let token = self
.azure_request(&mut state, super::azure_openai::AzureAudience::Inference)
.await?;
let endpoint = super::azure_openai::AzureEndpoint::parse(
state["endpoint"]
.as_str()
.ok_or_else(|| BrokerError::Credential("invalid Azure endpoint".into()))?,
)
.map_err(|_| BrokerError::Credential("invalid Azure endpoint".into()))?;
let deployment = model_id
.strip_prefix("azure-openai/")
.ok_or_else(|| BrokerError::Denied("invalid Azure model".into()))?;
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});
let rr = super::azure_openai::responses_request(
&endpoint,
deployment,
&serde_json::to_vec(&body).unwrap(),
)
.map_err(|_| BrokerError::Denied("invalid Azure request".into()))?;
let r = self
.http
.post(rr.url)
.bearer_auth(token)
.header("content-type", "application/json")
.body(rr.body)
.send()
.await
.map_err(|_| BrokerError::Transport("Azure invocation failed".into()))?;
if !r.status().is_success() {
return Err(BrokerError::Transport("Azure invocation rejected".into()));
}
if request.stream {
use futures::StreamExt;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
let terminal = Arc::new(AtomicBool::new(false));
let seen = terminal.clone();
let events = sse_json_stream(r).filter_map(move |item| {
let seen = seen.clone();
async move {
match item {
Ok(Some(v)) => v["delta"]
.as_str()
.or_else(|| v["text"].as_str())
.map(|delta| {
Ok(CloudEvent::TextDelta {
delta: delta.into(),
})
}),
Ok(None) => {
seen.store(true, Ordering::SeqCst);
Some(Ok(CloudEvent::Done))
}
Err(error) => Some(Err(error)),
}
}
});
let stream = events.chain(futures::stream::unfold(
(terminal, false),
|(terminal, emitted)| async move {
if terminal.load(Ordering::SeqCst) || emitted {
None
} else {
Some((
Err(BrokerError::Transport(
"Azure stream ended without [DONE]".into(),
)),
(terminal, true),
))
}
},
));
Ok(Box::pin(stream))
} else {
let text = read_body_capped(r, MAX_CLOUD_STREAM_EVENT_BYTES).await?;
let v: serde_json::Value = serde_json::from_str(&text)
.map_err(|_| BrokerError::Transport("invalid Azure response".into()))?;
let delta = v["output_text"]
.as_str()
.ok_or_else(|| {
BrokerError::Transport("Azure response omitted output".into())
})?
.to_owned();
Ok(Box::pin(futures::stream::iter(vec![
Ok(CloudEvent::TextDelta { delta }),
Ok(CloudEvent::Done),
])))
}
}
CloudProviderId::GoogleVertex => {
let mut state = self.state(provider)?;
let token = self.vertex_request(&mut state).await?;
let config: super::cloud::GoogleVertexConfig =
serde_json::from_value(state["config"].clone())
.map_err(|_| BrokerError::Credential("invalid Vertex state".into()))?;
let model = model_id
.strip_prefix("google-vertex/publishers/google/models/")
.filter(|m| {
m.bytes()
.all(|b| b.is_ascii_alphanumeric() || b"-._".contains(&b))
})
.ok_or_else(|| BrokerError::Denied("invalid Vertex model".into()))?;
let action = if request.stream {
"streamGenerateContent?alt=sse"
} else {
"generateContent"
};
let url=format!("https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:{}",config.location,config.project_id,config.location,model,action);
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<_>>()});
let r = self
.http
.post(url)
.bearer_auth(token)
.json(&body)
.send()
.await
.map_err(|_| BrokerError::Transport("Vertex invocation failed".into()))?;
if !r.status().is_success() {
return Err(BrokerError::Transport("Vertex invocation rejected".into()));
}
if request.stream {
use futures::StreamExt;
let events = sse_json_stream(r).map(|item| {
item.and_then(|value| {
value.ok_or_else(|| {
BrokerError::Transport(
"Vertex stream terminated without a provider event".into(),
)
})
})
});
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
let terminal = Arc::new(AtomicBool::new(false));
let terminal_in_events = terminal.clone();
let stream = events
.flat_map(move |item| {
let mut out = Vec::new();
match item {
Ok(v) => {
if terminal_in_events.load(Ordering::SeqCst) {
out.push(Err(BrokerError::Transport(
"Vertex emitted data after terminal metadata".into(),
)));
return futures::stream::iter(out);
}
for c in v["candidates"].as_array().into_iter().flatten() {
for p in
c["content"]["parts"].as_array().into_iter().flatten()
{
if let Some(delta) = p["text"].as_str() {
out.push(Ok(CloudEvent::TextDelta {
delta: delta.into(),
}));
}
}
if let Some(reason) = c["finishReason"].as_str() {
const VALID: &[&str] = &[
"STOP",
"MAX_TOKENS",
"SAFETY",
"RECITATION",
"OTHER",
"BLOCKLIST",
"PROHIBITED_CONTENT",
"SPII",
"MALFORMED_FUNCTION_CALL",
];
if !VALID.contains(&reason) {
out.push(Err(BrokerError::Transport(
"Vertex returned an invalid finish reason"
.into(),
)));
} else if terminal_in_events
.swap(true, Ordering::SeqCst)
{
out.push(Err(BrokerError::Transport(
"Vertex returned duplicate terminal metadata"
.into(),
)));
} else {
out.push(Ok(CloudEvent::Done));
}
}
}
if let Some(u) = v.get("usageMetadata") {
out.push(Ok(CloudEvent::Usage {
input_tokens: u["promptTokenCount"]
.as_u64()
.unwrap_or(0),
output_tokens: u["candidatesTokenCount"]
.as_u64()
.unwrap_or(0),
}));
}
}
Err(error) => out.push(Err(error)),
}
futures::stream::iter(out)
})
.chain(futures::stream::unfold(
(terminal, false),
|(terminal, emitted)| async move {
if terminal.load(Ordering::SeqCst) || emitted {
None
} else {
Some((
Err(BrokerError::Transport(
"Vertex stream ended without finish metadata".into(),
)),
(terminal, true),
))
}
},
));
Ok(Box::pin(stream))
} else {
let text = read_body_capped(r, MAX_CLOUD_STREAM_EVENT_BYTES).await?;
let v: serde_json::Value = serde_json::from_str(&text)
.map_err(|_| BrokerError::Transport("invalid Vertex response".into()))?;
let mut out = Vec::new();
for c in v["candidates"].as_array().into_iter().flatten() {
for p in c["content"]["parts"].as_array().into_iter().flatten() {
if let Some(delta) = p["text"].as_str() {
out.push(Ok(CloudEvent::TextDelta {
delta: delta.into(),
}));
}
}
}
out.push(Ok(CloudEvent::Done));
Ok(Box::pin(futures::stream::iter(out)))
}
}
}
}
}
const MAX_CLOUD_CATALOG_BODY_BYTES: usize = 2 * 1024 * 1024;
const MAX_CLOUD_CATALOG_ENTRIES: usize = 10_000;
const MAX_CLOUD_STREAM_EVENT_BYTES: usize = 1024 * 1024;
fn sse_json_stream(
resp: reqwest::Response,
) -> Pin<Box<dyn Stream<Item = Result<Option<serde_json::Value>, BrokerError>> + Send>> {
use futures::StreamExt;
let chunks = Box::pin(resp.bytes_stream());
Box::pin(futures::stream::unfold(
(chunks, Vec::<u8>::new(), false),
|(mut chunks, mut buf, done)| async move {
if done {
return None;
}
loop {
if let Some(end) = buf.windows(2).position(|w| w == b"\n\n") {
let frame: Vec<u8> = buf.drain(..end + 2).collect();
let text = match std::str::from_utf8(&frame[..end]) {
Ok(v) => v,
Err(_) => {
return Some((
Err(BrokerError::Transport("invalid SSE encoding".into())),
(chunks, buf, true),
))
}
};
let data = text
.lines()
.filter_map(|line| line.strip_prefix("data:"))
.map(str::trim_start)
.collect::<Vec<_>>()
.join("\n");
if data.is_empty() {
continue;
}
if data == "[DONE]" {
return Some((Ok(None), (chunks, buf, true)));
}
let value = serde_json::from_str(&data)
.map_err(|_| BrokerError::Transport("invalid SSE event".into()));
let terminal = value.is_err();
return Some((value.map(Some), (chunks, buf, terminal)));
}
match chunks.next().await {
Some(Ok(chunk)) if buf.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES => {
buf.extend_from_slice(&chunk)
}
Some(_) => {
return Some((
Err(BrokerError::Transport(
"cloud stream event exceeded limit".into(),
)),
(chunks, buf, true),
))
}
None if buf.is_empty() => return None,
None => {
return Some((
Err(BrokerError::Transport("truncated SSE event".into())),
(chunks, buf, true),
))
}
}
}
},
))
}
async fn cloud_catalog_json(resp: reqwest::Response) -> Result<serde_json::Value, BrokerError> {
let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
serde_json::from_str(&body)
.map_err(|_| BrokerError::Transport("invalid cloud catalog response".into()))
}
pub struct LocalBroker {
http: reqwest::Client,
local_base_url: Option<String>,
anthropic_usage_url: Option<String>,
google_gemini_base_url: Option<String>,
openai_codex_base_url: Option<String>,
muse_base_url: Option<String>,
refresh_endpoints: super::provider::RefreshEndpoints,
stream_idle_timeout: Duration,
request_timeout: Duration,
max_response_bytes: usize,
cloud_backend: Option<Arc<dyn CloudBackend>>,
account_policy: Option<AccountPolicy>,
auto_policy: Option<super::AutoSelectionPolicy>,
auto_policy_cache: std::sync::Mutex<
Option<(
std::path::PathBuf,
Option<std::time::SystemTime>,
super::AutoSelectionPolicy,
)>,
>,
current_accounts: SharedMap<(CredentialRef, SeatIdentity)>,
cooldowns: Arc<std::sync::Mutex<BTreeMap<String, Cooldown>>>,
usage_endpoint_override: Option<String>,
max_snapshot_age: Duration,
snapshots: SharedMap<super::usage::UsageSnapshot>,
usage_reads: Arc<super::usage_reads::UsageReads>,
instance: u64,
policy_cache: Arc<std::sync::Mutex<Option<CachedPolicy>>>,
observations: SharedMap<SlotObservations>,
selections: SharedMap<super::auth_state::SelectionState>,
}
#[derive(Debug, Clone, Default)]
struct SlotObservations {
last_usage: Option<super::auth_state::ObservedOutcome>,
last_usage_ok_ms: Option<u64>,
}
type SharedMap<T> = Arc<std::sync::Mutex<BTreeMap<String, T>>>;
type CachedPolicy = (Option<std::time::SystemTime>, AccountPolicy);
#[derive(Debug, Clone, PartialEq, Eq)]
struct Cooldown {
until_ms: u64,
reason: String,
}
const DEFAULT_COOLDOWN: Duration = Duration::from_secs(15 * 60);
const MAX_COOLDOWN: Duration = Duration::from_secs(8 * 24 * 60 * 60);
pub const DEFAULT_MAX_SNAPSHOT_AGE: Duration = Duration::from_secs(60);
pub const BROKER_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
pub const GEMINI_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
fn next_broker_instance() -> u64 {
static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
}
impl LocalBroker {
pub fn new(http: reqwest::Client) -> Self {
let cloud_backend = Arc::new(ProductionCloudBackend::new(http.clone()));
Self {
http,
local_base_url: None,
anthropic_usage_url: None,
google_gemini_base_url: None,
openai_codex_base_url: None,
muse_base_url: None,
refresh_endpoints: Default::default(),
stream_idle_timeout: BROKER_STREAM_IDLE_TIMEOUT,
request_timeout: PROXY_REQUEST_TIMEOUT,
max_response_bytes: MAX_PROXY_RESPONSE_BYTES,
cloud_backend: Some(cloud_backend),
account_policy: None,
auto_policy: None,
auto_policy_cache: std::sync::Mutex::new(None),
current_accounts: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
cooldowns: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
usage_endpoint_override: None,
max_snapshot_age: DEFAULT_MAX_SNAPSHOT_AGE,
snapshots: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
usage_reads: super::usage_reads::process_table(),
instance: next_broker_instance(),
policy_cache: Arc::new(std::sync::Mutex::new(None)),
observations: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
selections: Arc::new(std::sync::Mutex::new(BTreeMap::new())),
}
}
pub fn with_cloud_backend(mut self, backend: Arc<dyn CloudBackend>) -> Self {
self.cloud_backend = Some(backend);
self
}
pub fn with_account_policy(mut self, policy: AccountPolicy) -> Self {
self.account_policy = Some(policy);
self
}
pub fn with_auto_policy(mut self, policy: super::AutoSelectionPolicy) -> Self {
self.auto_policy = Some(policy);
self
}
fn auto_policy(&self) -> super::AutoSelectionPolicy {
if let Some(policy) = &self.auto_policy {
return policy.clone();
}
let path = crate::config::resolve_read_path("config");
let mtime = std::fs::metadata(&path).and_then(|m| m.modified()).ok();
let mut cache = self
.auto_policy_cache
.lock()
.unwrap_or_else(|e| e.into_inner());
if let Some((cached_path, cached_time, policy)) = cache.as_ref() {
if cached_path == &path && *cached_time == mtime {
return policy.clone();
}
}
let (policy, warnings) =
super::AutoSelectionPolicy::from_config_map(&crate::config::load_config().auth.auto);
for warning in warnings {
tracing::warn!(%warning, "automatic account policy");
}
*cache = Some((path, mtime, policy.clone()));
policy
}
fn current_key(provider: OAuthProviderId, model: Option<&str>) -> String {
format!(
"{}|{}|{}",
storage::auth_file_path().display(),
provider,
model.unwrap_or("")
)
}
#[doc(hidden)]
pub fn with_openai_codex_base_url_for_tests(mut self, base_url: impl Into<String>) -> Self {
self.openai_codex_base_url = Some(base_url.into().trim_end_matches('/').to_string());
self
}
#[doc(hidden)]
pub fn with_muse_base_url_for_tests(mut self, base_url: impl Into<String>) -> Self {
self.muse_base_url = Some(base_url.into().trim_end_matches('/').to_string());
self
}
fn stream_idle_timeout_for(&self, provider_key: &str) -> Duration {
if provider_key == "google-gemini" {
GEMINI_STREAM_IDLE_TIMEOUT.max(self.stream_idle_timeout)
} else {
self.stream_idle_timeout
}
}
#[doc(hidden)]
pub fn with_stream_idle_timeout_for_tests(mut self, idle: Duration) -> Self {
self.stream_idle_timeout = idle;
self
}
#[doc(hidden)]
pub fn with_muse_mint_url_for_tests(mut self, mint_url: impl Into<String>) -> Self {
self.refresh_endpoints =
std::mem::take(&mut self.refresh_endpoints).with_muse_mint_url_for_tests(mint_url);
self
}
#[doc(hidden)]
pub fn with_anthropic_token_url_for_tests(mut self, url: impl Into<String>) -> Self {
self.refresh_endpoints =
std::mem::take(&mut self.refresh_endpoints).with_anthropic_token_url_for_tests(url);
self
}
pub async fn access_token_auto(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
self.vend_auto(provider, model).await
}
#[doc(hidden)]
pub fn with_usage_endpoint_override(mut self, url: impl Into<String>) -> Self {
self.usage_endpoint_override = Some(url.into());
self
}
pub fn with_max_snapshot_age(mut self, age: Duration) -> Self {
self.max_snapshot_age = age;
self
}
fn policy(&self) -> AccountPolicy {
if let Some(policy) = &self.account_policy {
return policy.clone();
}
let mtime = std::fs::metadata(crate::config::resolve_read_path("config"))
.and_then(|m| m.modified())
.ok();
let mut cache = self.policy_cache.lock().unwrap_or_else(|e| e.into_inner());
let config_policy = match cache.as_ref() {
Some((cached_mtime, policy)) if *cached_mtime == mtime => policy.clone(),
_ => {
let (policy, _warnings) = AccountPolicy::from_config_map(
&crate::config::load_config().auth.accounts,
);
*cache = Some((mtime, policy.clone()));
policy
}
};
config_policy.with_env_overlay()
}
fn cooldown_key(&self, cred: &CredentialRef) -> String {
let seat = self
.stored_seat(cred)
.ok()
.map(|s| s.as_str().to_string())
.unwrap_or_else(|| cred.storage_key());
format!(
"{}|{}|{}",
storage::auth_file_path().display(),
cred.provider,
seat
)
}
fn cooldown_until(&self, storage_key: &str, now_ms: u64) -> Option<u64> {
let cred = CredentialRef::parse_storage_key(storage_key)?;
let key = self.cooldown_key(&cred);
self.cooldowns
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&key)
.filter(|c| c.until_ms > now_ms)
.map(|c| c.until_ms)
}
pub fn auth_state(
&self,
query: &super::auth_state::AuthStateQuery,
) -> Result<super::auth_state::AuthStateReport, BrokerError> {
auth_state_report(Some(self), query)
}
fn cooldown_state(
&self,
cred: &CredentialRef,
now_ms: u64,
) -> Option<super::auth_state::CooldownState> {
let key = self.cooldown_key(cred);
self.cooldowns
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&key)
.filter(|c| c.until_ms > now_ms)
.map(|c| super::auth_state::CooldownState {
until_ms: c.until_ms,
reason: super::auth_state::CooldownReason::from_reported(&c.reason),
})
}
fn clear_state(&self, creds: &[CredentialRef], whole_providers: &[OAuthProviderId]) {
let store = storage::auth_file_path().display().to_string();
{
let mut cooldowns = self.cooldowns.lock().unwrap_or_else(|e| e.into_inner());
for cred in creds {
cooldowns.remove(&self.cooldown_key(cred));
cooldowns.remove(&format!("{store}|{}|{}", cred.provider, cred.storage_key()));
}
for provider in whole_providers {
let prefix = format!("{store}|{provider}|");
cooldowns.retain(|k, _| !k.starts_with(&prefix));
}
}
for cred in creds {
self.forget_snapshots(cred);
self.usage_reads.forget(&self.read_key(cred));
}
{
let mut observations = self.observations.lock().unwrap_or_else(|e| e.into_inner());
for cred in creds {
observations.remove(&Self::state_key(cred));
}
}
self.current_accounts
.lock()
.unwrap_or_else(|e| e.into_inner())
.retain(|_, (current, _)| !creds.contains(current));
let mut selections = self.selections.lock().unwrap_or_else(|e| e.into_inner());
for provider in whole_providers {
let prefix = format!("{store}|{provider}|");
selections.retain(|k, _| !k.starts_with(&prefix));
}
}
fn check_slot_exists(&self, cred: &CredentialRef) -> Result<(), BrokerError> {
if let Account::Named(label) = &cred.account {
match storage::load_credential(cred) {
Ok(Some(_)) => {}
Ok(None) => {
return Err(BrokerError::UnknownAccount {
provider: cred.provider.as_str().to_string(),
label: label.as_str().to_string(),
})
}
Err(e) => return Err(BrokerError::Credential(e)),
}
}
Ok(())
}
async fn resolve_and_vend(
&self,
provider: OAuthProviderId,
explicit: Option<Account>,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
let selector = match explicit {
Some(account) => AccountSelector::Account(account),
None => self.policy().selector(provider),
};
match selector {
AccountSelector::Account(account) => {
let credential = CredentialRef::new(provider, account);
let token = self.access_token_for(&credential).await?;
Ok(PinnedToken { credential, token })
}
AccountSelector::Auto => self.vend_auto(provider, model).await,
AccountSelector::Invalid { source, reason } => {
Err(BrokerError::InvalidAccount(format!("{source}: {reason}")))
}
}
}
const MAX_AUTO_VEND_ROUNDS: usize = 2;
async fn vend_auto(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
use super::token::{ensure_fresh_credential_for_seat, SeatVendError};
for _ in 0..Self::MAX_AUTO_VEND_ROUNDS {
let (credential, seat) = match self.select_auto(provider, model).await {
Ok(choice) => choice,
Err(e) => {
let failure = match &e {
BrokerError::NoAccountAvailable { failure, .. } => failure.clone(),
_ => Some(super::quota_policy::SelectionFailure::new(
super::quota_policy::SelectionFailureKind::CredentialStore,
)),
};
self.record_selection(provider, model, None, failure, None);
return Err(e);
}
};
match ensure_fresh_credential_for_seat(
&self.http,
&credential,
&seat,
&self.refresh_endpoints,
)
.await
{
Ok(creds) => {
self.current_accounts
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(
Self::current_key(provider, model),
(credential.clone(), seat),
);
self.record_selection(provider, model, Some(&credential), None, None);
return Ok(PinnedToken {
credential,
token: AccessToken {
token: creds.access,
expires: creds.expires,
},
})
}
Err(SeatVendError::SeatChanged) => {
tracing::info!(credential = %credential, "slot changed after selection; re-selecting");
self.forget_snapshots(&credential);
continue;
}
Err(SeatVendError::Other(f)) => {
let cause = refresh_failure(&f);
tracing::warn!(
credential = %credential,
kind = %cause.kind,
"automatic choice could not be vended: token refresh failed"
);
self.record_selection(provider, model, Some(&credential), None, Some(cause));
return Err(BrokerError::Refresh(auto_vend_failure(&credential, f)));
}
}
}
let failure = super::quota_policy::SelectionFailure::new(
super::quota_policy::SelectionFailureKind::SeatChurn,
);
self.record_selection(provider, model, None, Some(failure.clone()), None);
Err(selection_failed(
provider,
model,
"selected credential kept changing while its capacity was being proven".into(),
failure,
))
}
fn record_selection(
&self,
provider: OAuthProviderId,
model: Option<&str>,
chosen: Option<&CredentialRef>,
failure: Option<super::quota_policy::SelectionFailure>,
vend_failure: Option<super::quota_policy::ObservationFailure>,
) {
let state = super::auth_state::SelectionState {
provider: provider.as_str().to_string(),
model: model.map(str::to_string),
at_ms: crate::epoch_millis(),
chosen: chosen.map(|c| c.account.label_str().to_string()),
failure,
vend_failure,
};
self.selections
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(Self::current_key(provider, model), state);
}
fn seat_key(cred: &CredentialRef, seat: &SeatIdentity) -> String {
format!(
"{}|{}|{}",
storage::auth_file_path().display(),
cred.storage_key(),
seat.as_str()
)
}
fn stored_seat(&self, cred: &CredentialRef) -> Result<SeatIdentity, BrokerError> {
match storage::load_credential(cred).map_err(BrokerError::Credential)? {
Some(creds) => Ok(SeatIdentity::of(cred.provider, &creds)),
None => Err(BrokerError::UnknownAccount {
provider: cred.provider.as_str().to_string(),
label: cred.account.label_str().to_string(),
}),
}
}
fn cached_snapshot(&self, key: &str) -> Option<super::usage::UsageSnapshot> {
let max_age_ms = self.max_snapshot_age.as_millis() as u64;
let now_ms = crate::epoch_millis();
self.snapshots
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(key)
.filter(|s| !s.is_stale(now_ms, max_age_ms))
.cloned()
}
fn read_key(&self, cred: &CredentialRef) -> String {
format!(
"{}|{}|{}",
storage::auth_file_path().display(),
cred.storage_key(),
self.usage_endpoint_override.as_deref().unwrap_or("")
)
}
fn forget_snapshots(&self, cred: &CredentialRef) {
let prefix = format!("|{}|", cred.storage_key());
self.snapshots
.lock()
.unwrap_or_else(|e| e.into_inner())
.retain(|k, _| !k.contains(&prefix));
}
async fn fresh_usage(
&self,
cred: &CredentialRef,
) -> Result<(super::usage::UsageSnapshot, SeatIdentity), super::usage_reads::ReadError> {
let seat = self
.stored_seat(cred)
.map_err(|e| UsageReadError::slot(e).into_read_error())?;
let key = Self::seat_key(cred, &seat);
if let Some(cached) = self.cached_snapshot(&key) {
return Ok((cached, seat));
}
let (snapshot, used) = self.shared_read(cred).await?;
self.forget_snapshots(cred);
let used_key = Self::seat_key(cred, &used);
self.snapshots
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(used_key, snapshot.clone());
Ok((snapshot, used))
}
async fn usage_with_seat(
&self,
cred: &CredentialRef,
) -> Result<(super::usage::UsageSnapshot, SeatIdentity), UsageReadError> {
self.shared_read(cred)
.await
.map_err(|e| UsageReadError::from_read(&e))
}
fn state_key(cred: &CredentialRef) -> String {
format!(
"{}|{}",
storage::auth_file_path().display(),
cred.storage_key()
)
}
fn record_usage(
&self,
cred: &CredentialRef,
failure: Option<super::quota_policy::ObservationFailure>,
) {
let at_ms = crate::epoch_millis();
let mut observations = self.observations.lock().unwrap_or_else(|e| e.into_inner());
let slot = observations.entry(Self::state_key(cred)).or_default();
if failure.is_none() {
slot.last_usage_ok_ms = Some(at_ms);
}
slot.last_usage = Some(super::auth_state::ObservedOutcome { at_ms, failure });
}
async fn shared_read(&self, cred: &CredentialRef) -> super::usage_reads::ReadResult {
use super::usage_reads::{ReadError, ReadPath, USAGE_READ_TARGET};
let key = self.read_key(cred);
let (result, path) = self
.usage_reads
.read(&key, || self.network_read(cred))
.await;
if path != ReadPath::NegativeCacheHit {
self.record_usage(cred, result.as_ref().err().map(ReadError::cause));
}
let outcome = match (path, &result) {
(ReadPath::Network, Ok(_)) => "ok",
(ReadPath::Network, Err(_)) => "failed",
(ReadPath::Joined, _) => "joined",
(ReadPath::NegativeCacheHit, _) => "negative_cache_hit",
};
let now_ms = crate::epoch_millis();
let (kind, status, retry_in_s) = match &result {
Ok(_) => ("none", String::new(), 0),
Err(ReadError::Transient {
failure,
retry_at_ms,
}) => (
failure.kind,
failure.status.map(|s| s.to_string()).unwrap_or_default(),
retry_at_ms.saturating_sub(now_ms) / 1000,
),
Err(ReadError::Other { cause, .. }) => (
cause.kind.as_str(),
cause.http_status.map(|s| s.to_string()).unwrap_or_default(),
0,
),
};
tracing::info!(
target: USAGE_READ_TARGET,
broker = self.instance,
account = %cred.account.label_str(),
outcome = %outcome,
kind = %kind,
status = %status,
retry_in_s,
"usage read"
);
result
}
async fn network_read(&self, cred: &CredentialRef) -> super::usage_reads::Fetched {
use super::usage::{fetch_usage, supports_usage, UsageClient, UsageError, UsageFetchOptions};
use super::usage_reads::Fetched;
if !supports_usage(cred.provider) {
return Fetched::Usage(UsageError::UnsupportedProvider {
provider: cred.provider.as_str().to_string(),
});
}
if let Err(e) = self.check_slot_exists(cred) {
return UsageReadError::slot(e).into_fetched();
}
let creds = match super::token::ensure_fresh_credential_via(
&self.http,
cred,
&self.refresh_endpoints,
)
.await
{
Ok(creds) => creds,
Err(f) => return UsageReadError::refresh(f).into_fetched(),
};
let seat = SeatIdentity::of(cred.provider, &creds);
let client = match UsageClient::new() {
Ok(client) => client,
Err(e) => return Fetched::Usage(e),
};
let opts = UsageFetchOptions {
endpoint_override: self.usage_endpoint_override.clone(),
..UsageFetchOptions::default()
};
match fetch_usage(
&client,
cred.provider,
cred.account.label_str(),
&creds.access,
&opts,
)
.await
{
Ok(snapshot) => Fetched::Ok(snapshot, seat),
Err(e) => Fetched::Usage(e),
}
}
async fn capacity_for(&self, cred: &CredentialRef) -> CandidateRead {
use super::quota_policy::AccountCapacity;
use super::usage_reads::ReadError;
match self.fresh_usage(cred).await {
Ok((snapshot, seat)) => CandidateRead {
capacity: capacity_from_snapshot(cred, &snapshot),
seat: Some(seat),
transient: None,
grace: None,
},
Err(read_err) => {
let err = UsageReadError::from_read(&read_err);
if err.failure.kind.needs_action() {
tracing::warn!(
credential = %cred,
kind = %err.failure.kind,
detail = %action_detail(cred, &err),
"usage read needs action"
);
} else {
tracing::debug!(credential = %cred, cause = %err.failure, "usage read failed");
}
let capacity = AccountCapacity::failed(cred.clone(), err.failure);
let ReadError::Transient {
failure,
retry_at_ms,
} = read_err
else {
return CandidateRead {
capacity,
seat: None,
transient: None,
grace: None,
};
};
let graced = self.stored_seat(cred).ok().and_then(|seat| {
let good = self
.usage_reads
.last_good_for(&self.read_key(cred), &seat)?;
Some((good, seat))
});
let (grace, seat) = match graced {
Some((good, seat)) => (Some(capacity_from_snapshot(cred, &good)), Some(seat)),
None => (None, None),
};
CandidateRead {
capacity,
seat,
transient: Some((failure, retry_at_ms)),
grace,
}
}
}
}
pub async fn plan(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<super::SelectionPlan, BrokerError> {
self.selection_plan(provider, model)
.await
.map(|(plan, ..)| plan)
}
#[allow(clippy::type_complexity)]
async fn selection_plan(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<
(
super::SelectionPlan,
BTreeMap<String, SeatIdentity>,
BTreeMap<String, (super::usage_reads::TransientReadFailure, u64)>,
),
BrokerError,
> {
use super::quota_policy::SelectionRequest;
let summaries = storage::list_accounts(provider).map_err(BrokerError::Credential)?;
let mut candidates = Vec::with_capacity(summaries.len());
let mut seats = BTreeMap::new();
let mut transient = BTreeMap::new();
let mut grace = Vec::new();
let mut seen = std::collections::BTreeSet::new();
let mut exclude = Vec::new();
for summary in &summaries {
let Some(cred) = summary.credential_ref() else {
continue;
};
let read = self.capacity_for(&cred).await;
let (mut capacity, seat) = (read.capacity, read.seat);
if let Some(failed) = read.transient {
transient.insert(cred.storage_key(), failed);
}
require_requested_model(provider, model, &mut capacity);
grace.extend(read.grace.map(|mut good| {
require_requested_model(provider, model, &mut good);
good
}));
if let Some(seat) = seat {
if !seen.insert(seat.as_str().to_string()) {
exclude.push(cred.clone());
}
seats.insert(cred.storage_key(), seat);
}
candidates.push(capacity);
}
let now_ms = crate::epoch_millis();
for capacity in &mut candidates {
capacity.cooldown_until_ms =
self.cooldown_until(&capacity.credential.storage_key(), now_ms);
}
let previous = self
.current_accounts
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&Self::current_key(provider, model))
.cloned();
let current = previous
.filter(|(cred, seat)| seats.get(&cred.storage_key()) == Some(seat))
.map(|(cred, _)| cred);
let policy = self.auto_policy();
let preference: Vec<_> = summaries.iter().map(|s| s.label.clone()).collect();
let request = SelectionRequest {
model,
preference: &preference,
exclude: &exclude,
strategy: policy.strategy_for(provider),
current: current.as_ref(),
urgent_horizon_ms: policy.urgent_horizon_ms,
sticky: policy.sticky,
grace: &grace,
..SelectionRequest::new(provider, now_ms, self.max_snapshot_age.as_millis() as u64)
};
Ok((
super::SelectionPlan::from_candidates(&request, &candidates),
seats,
transient,
))
}
async fn select_auto(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<(CredentialRef, SeatIdentity), BrokerError> {
use super::quota_policy::{Selection, SelectionFailure, SelectionFailureKind};
use super::usage_reads::USAGE_READ_TARGET;
let (plan, mut seats, transient) = match self.selection_plan(provider, model).await {
Ok(planned) => planned,
Err(e) => {
log_selection_failure(
provider,
model,
&SelectionFailure::new(SelectionFailureKind::CredentialStore),
);
return Err(e);
}
};
match plan.selection {
Selection::Selected {
credential,
tier,
budget_reset_ms,
utilization,
..
} => {
let seat = seats.remove(&credential.storage_key());
debug_assert!(
seat.is_some(),
"selected account has no paired seat identity"
);
let Some(seat) = seat else {
return Err(selection_failed(
provider,
model,
"selected account has no paired seat identity".into(),
SelectionFailure::new(SelectionFailureKind::SeatChurn),
));
};
if let Some((failure, retry_at_ms)) = transient.get(&credential.storage_key()) {
tracing::info!(
target: USAGE_READ_TARGET,
broker = self.instance,
account = %credential.account.label_str(),
outcome = %"grace",
kind = %failure.kind,
status = %failure.status.map(|s| s.to_string()).unwrap_or_default(),
retry_in_s = retry_at_ms.saturating_sub(crate::epoch_millis()) / 1000,
"usage read"
);
}
tracing::info!(%credential, ?plan.strategy, tier, ?budget_reset_ms, ?utilization, "automatic subscription choice");
Ok((credential, seat))
}
Selection::NoCapacity {
rejections,
earliest_reset_ms,
} => {
let detail: Vec<_> = rejections
.iter()
.map(|r| format!("{}: {}", r.credential.account, r.describe()))
.collect();
let reset = earliest_reset_ms
.map(|t| format!("; earliest reset at {t}"))
.unwrap_or_default();
let mut failure = super::quota_policy::SelectionFailure::from_rejections(
&rejections,
plan.now_ms,
);
if failure.kind == SelectionFailureKind::UsageReadFailed {
failure.usage_read_retry_ms = rejections
.iter()
.filter_map(|r| transient.get(&r.credential.storage_key()))
.map(|(_, at)| *at)
.min();
}
let reason = if detail.is_empty() {
format!("no {provider} account is logged in")
} else {
format!("{}{reset}", detail.join(", "))
};
Err(selection_failed(provider, model, reason, failure))
}
}
}
pub fn with_local_base_url(http: reqwest::Client, base_url: impl Into<String>) -> Self {
Self {
local_base_url: Some(base_url.into()),
..Self::new(http)
}
}
#[doc(hidden)]
pub fn with_google_gemini_base_url_for_tests(
http: reqwest::Client,
base_url: impl Into<String>,
) -> Self {
Self {
google_gemini_base_url: Some(base_url.into().trim_end_matches('/').to_string()),
..Self::new(http)
}
}
#[doc(hidden)]
pub fn with_anthropic_usage_url(mut self, url: impl Into<String>) -> Self {
self.anthropic_usage_url = Some(url.into());
self
}
#[doc(hidden)]
pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
self.request_timeout = timeout;
self
}
#[doc(hidden)]
pub fn with_max_response_bytes(mut self, cap: usize) -> Self {
self.max_response_bytes = cap;
self
}
fn resolve_static_key(&self, provider: &str) -> Result<String, BrokerError> {
if provider == LOCAL_PROVIDER_KEY {
return Ok(self.resolve_local_key());
}
let spec = static_provider(provider)
.ok_or_else(|| BrokerError::UnknownProvider(provider.to_string()))?;
if let Ok(Some(key)) = storage::load_static_key(provider) {
return Ok(key);
}
if let Some(key) = discover_legacy_static_key(spec) {
if let Err(e) = storage::save_static_key(provider, &key) {
tracing::warn!(provider, "static key migration failed: {e}");
}
return Ok(key);
}
Err(BrokerError::NotConfigured(provider.to_string()))
}
fn resolve_local_key(&self) -> String {
storage::load_static_key(LOCAL_PROVIDER_KEY)
.ok()
.flatten()
.or_else(|| {
crate::config::get_provider_keys()
.get(LOCAL_PROVIDER_KEY)
.filter(|s| !s.is_empty())
.cloned()
})
.or_else(|| {
std::env::var("LOCAL_API_KEY")
.ok()
.filter(|s| !s.is_empty())
})
.unwrap_or_else(|| "local".to_string())
}
fn base_url_for(&self, provider: &str) -> Result<String, BrokerError> {
if provider == LOCAL_PROVIDER_KEY {
if let Some(url) = &self.local_base_url {
return Ok(url.trim_end_matches('/').to_string());
}
return Ok(local_endpoint_url());
}
static_provider(provider)
.map(|s| s.base_url.trim_end_matches('/').to_string())
.ok_or_else(|| BrokerError::UnknownProvider(provider.to_string()))
}
async fn send(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
request.validate()?;
let (provider_key, explicit_account) = request.split_provider()?;
let provider_key = provider_key.as_str();
let oauth_provider: Option<OAuthProviderId> = provider_key.parse().ok();
let mut pinned: Option<PinnedToken> = None;
let muse_env_key = (provider_key == "muse" && explicit_account.is_none())
.then(super::muse::api_key_from_env)
.flatten();
if let Some(provider) = oauth_provider.filter(|_| muse_env_key.is_none()) {
if provider_key != LOCAL_PROVIDER_KEY && static_provider(provider_key).is_none() {
pinned = Some(self.resolve_and_vend(provider, explicit_account, None).await?);
}
}
let bearer = |pinned: &Option<PinnedToken>| -> Result<String, BrokerError> {
pinned
.as_ref()
.map(|p| p.token.token.clone())
.ok_or_else(|| BrokerError::UnknownProvider(request.provider.clone()))
};
let (key, base) = if provider_key == "xai-auth" {
(bearer(&pinned)?, "https://api.x.ai/v1".to_string())
} else if provider_key == "github-copilot" {
(
bearer(&pinned)?,
super::github_copilot_models_base_url().to_string(),
)
} else if provider_key == "google-gemini" {
let base = self
.google_gemini_base_url
.clone()
.unwrap_or_else(|| GOOGLE_GEMINI_CODE_ASSIST_BASE_URL.to_string());
(bearer(&pinned)?, base)
} else if provider_key == "openai-codex" {
let base = self
.openai_codex_base_url
.clone()
.unwrap_or_else(|| OPENAI_CODEX_BACKEND_BASE_URL.to_string());
(bearer(&pinned)?, base)
} else if provider_key == "kimi-code" {
(bearer(&pinned)?, super::kimi_code::API_BASE_URL.to_string())
} else if provider_key == "muse" {
let key = match &muse_env_key {
Some(key) => key.clone(),
None => bearer(&pinned)?,
};
let base = self
.muse_base_url
.clone()
.unwrap_or_else(|| super::muse::API_ORIGIN.to_string());
(key, base)
} else if provider_key == "anthropic" {
(bearer(&pinned)?, "https://api.anthropic.com".to_string())
} else {
(
self.resolve_static_key(provider_key)?,
self.base_url_for(provider_key)?,
)
};
let response = self
.dispatch(request, provider_key, &key, &base, &pinned)
.await?;
if provider_key == "muse" && response.status() == reqwest::StatusCode::UNAUTHORIZED {
if let Some(rejected) = pinned.as_ref() {
tracing::info!(credential = %rejected.credential, "muse key rejected (401); re-minting once");
match self.recover_rejected_token(&rejected.credential, &key).await {
Ok(fresh) => {
let pinned = Some(PinnedToken {
credential: rejected.credential.clone(),
token: fresh,
});
let key = bearer(&pinned)?;
return self
.dispatch(request, provider_key, &key, &base, &pinned)
.await;
}
Err(BrokerError::Refresh(f))
if f.kind.is_retryable()
|| f.kind == super::token::RefreshFailureKind::ActionRequired =>
{
tracing::warn!(
credential = %rejected.credential,
reason = %f,
"muse key re-mint failed; returning a retryable error"
);
return Err(BrokerError::Refresh(f));
}
Err(e) => {
tracing::warn!(
credential = %rejected.credential,
reason = %e,
"muse key re-mint refused or rejected; returning the upstream 401"
);
}
}
}
}
Ok(response)
}
async fn dispatch(
&self,
request: &ProxyRequest,
provider_key: &str,
key: &str,
base: &str,
pinned: &Option<PinnedToken>,
) -> Result<reqwest::Response, BrokerError> {
let url = format!("{base}{}", request.path);
let mut builder = if provider_key == "google-gemini"
|| provider_key == "openai-codex"
|| provider_key == "anthropic"
|| provider_key == "kimi-code"
|| provider_key == "muse"
{
let no_redirect = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(Duration::from_secs(10))
.read_timeout(self.stream_idle_timeout_for(provider_key))
.build()
.map_err(|e| BrokerError::Transport(format!("client build failed: {e}")))?;
match request.method {
ProxyMethod::Get => no_redirect.get(&url),
ProxyMethod::Post => no_redirect.post(&url),
}
} else {
match request.method {
ProxyMethod::Get => self.http.get(&url),
ProxyMethod::Post => self.http.post(&url),
}
};
builder = builder.bearer_auth(key);
if provider_key == "openai-codex" {
let cred = pinned
.as_ref()
.map(|p| p.credential.clone())
.unwrap_or_else(|| CredentialRef::default_for(OAuthProviderId::OpenAiCodex));
let account_id = super::extract_codex_account_id(key)
.or_else(|| {
storage::load_credential(&cred)
.ok()
.flatten()
.and_then(|c| c.account_id)
.filter(|s| !s.trim().is_empty())
})
.ok_or_else(|| {
BrokerError::Credential(
"openai-codex credential is missing chatgpt account id".into(),
)
})?;
builder = builder
.header("chatgpt-account-id", account_id)
.header("originator", "synaps")
.header("OpenAI-Beta", "responses=experimental")
.header("accept", "application/json");
}
if provider_key == "anthropic" {
for (name, value) in anthropic_oauth_catalog_request_headers() {
builder = builder.header(*name, *value);
}
}
if provider_key == "github-copilot" {
for (name, value) in super::github_copilot_models_request_headers() {
builder = builder.header(*name, *value);
}
if request.path != "/models" {
builder = builder
.header("Openai-Intent", "conversation-edits")
.header("X-Initiator", "agent");
}
}
if provider_key == "google-gemini" {
if request.stream && request.path == "/v1internal:streamGenerateContent" {
builder = builder.query(&[("alt", "sse")]);
}
builder = builder.header("user-agent", "SynapsCLI/0.6.0 (google-gemini)");
}
if provider_key == "muse" {
builder = builder.header("user-agent", super::muse::USER_AGENT);
}
if provider_key == "kimi-code" {
let device_id = super::kimi_code::load_or_create_device_id();
for (name, value) in super::kimi_code::identity_request_headers(&device_id) {
builder = builder.header(name, value);
}
}
if let Some(bytes) = &request.body_bytes {
builder = builder
.header("content-type", "application/json")
.body(bytes.clone());
} else if let Some(body) = &request.body {
builder = builder.json(body);
}
if request.stream {
builder = builder.header("accept", "text/event-stream");
} else {
builder = builder.timeout(self.request_timeout);
}
builder.send().await.map_err(|e| {
if e.is_connect() && provider_key == LOCAL_PROVIDER_KEY {
BrokerError::Transport(format!(
"can't reach local endpoint at {url} — is Ollama/LM Studio running?"
))
} else {
BrokerError::Transport(format!("request to {} failed: {e}", request.provider))
}
})
}
}
pub fn capacity_from_snapshot(
cred: &CredentialRef,
snapshot: &super::usage::UsageSnapshot,
) -> super::quota_policy::AccountCapacity {
use super::quota_policy::{
AccountCapacity, ModelAvailability, ModelState, QuotaObservation, WindowLimit,
};
use super::usage::{Availability, WindowScope};
let mut windows: Vec<WindowLimit> = snapshot
.windows
.iter()
.filter_map(|w| {
let models = match &w.scope {
WindowScope::Account => None,
WindowScope::Model { model } => Some(vec![model.clone()]),
WindowScope::Feature { .. } => return None,
};
Some(WindowLimit {
id: w.id.clone(),
duration_ms: w.duration_secs.map(|s| s.saturating_mul(1000)),
used_percent: w.used_percent.valid(),
limit_reached: w.limit_reached,
resets_at_ms: w.reset_at,
models,
})
})
.collect();
if snapshot.limit_reached == Some(true) {
windows.push(WindowLimit {
id: "limit_reached".into(),
duration_ms: None,
used_percent: None,
limit_reached: Some(true),
resets_at_ms: snapshot.earliest_reset_at(),
models: None,
});
}
let models = if snapshot.model_availability.is_empty() {
None
} else {
Some(
snapshot
.model_availability
.iter()
.map(|m| ModelAvailability {
model: m.model.clone(),
state: match m.availability {
Availability::Available => ModelState::Available,
Availability::Exhausted => ModelState::Exhausted,
Availability::Unknown => ModelState::Unknown,
},
})
.collect(),
)
};
AccountCapacity {
credential: cred.clone(),
observed_at_ms: Some(snapshot.observed_at),
observation: QuotaObservation::Ok { windows, models },
cooldown_until_ms: None,
failure: None,
}
}
#[derive(Debug)]
struct UsageReadError {
error: BrokerError,
failure: super::quota_policy::ObservationFailure,
}
struct CandidateRead {
capacity: super::quota_policy::AccountCapacity,
seat: Option<SeatIdentity>,
transient: Option<(super::usage_reads::TransientReadFailure, u64)>,
grace: Option<super::quota_policy::AccountCapacity>,
}
impl UsageReadError {
fn from_read(e: &super::usage_reads::ReadError) -> Self {
Self {
error: e.to_broker_error(),
failure: e.cause(),
}
}
fn into_read_error(self) -> super::usage_reads::ReadError {
super::usage_reads::ReadError::Other {
error: self.error,
cause: self.failure,
}
}
fn into_fetched(self) -> super::usage_reads::Fetched {
super::usage_reads::Fetched::Failed {
error: self.error,
cause: self.failure,
}
}
fn refresh(f: super::token::RefreshFailure) -> Self {
Self {
failure: refresh_failure(&f),
error: BrokerError::Credential(f.into()),
}
}
fn slot(error: BrokerError) -> Self {
use super::quota_policy::{AccountFailureKind as K, ObservationFailure};
let kind = match &error {
BrokerError::UnknownAccount { .. } => K::CredentialMissing,
_ => K::CredentialStore,
};
Self {
failure: ObservationFailure::new(kind),
error,
}
}
#[cfg(test)]
fn usage(e: super::usage::UsageError) -> Self {
let table = super::usage_reads::UsageReads::default();
match table.record_at("slot", super::usage_reads::Fetched::Usage(e), 0) {
Err(read) => Self::from_read(&read),
Ok(_) => unreachable!("a usage error is never a reading"),
}
}
}
fn require_requested_model(
provider: OAuthProviderId,
model: Option<&str>,
capacity: &mut super::quota_policy::AccountCapacity,
) {
use super::quota_policy::{model_matches, ModelAvailability, ModelState, QuotaObservation};
if provider != OAuthProviderId::OpenAiCodex {
return;
}
if let (Some(model), QuotaObservation::Ok { models, .. }) = (model, &mut capacity.observation) {
let availability = models.get_or_insert_with(Vec::new);
if !availability.iter().any(|m| model_matches(&m.model, model)) {
availability.push(ModelAvailability {
model: model.to_string(),
state: ModelState::Unknown,
});
}
}
}
fn action_detail(cred: &CredentialRef, err: &UsageReadError) -> String {
use super::quota_policy::AccountFailureKind as K;
match err.failure.kind {
K::CredentialStore => format!(
"the stored login of '{}' could not be read from {}; repair or restore the \
credential store",
cred.account.label_str(),
storage::auth_file_path().display()
),
_ => err.error.to_string(),
}
}
fn refresh_failure(f: &super::token::RefreshFailure) -> super::quota_policy::ObservationFailure {
let failure = refresh_kind_failure(f.kind);
match f.http_status() {
Some(status) => failure.with_status(status),
None => failure,
}
}
fn refresh_kind_failure(
kind: super::token::RefreshFailureKind,
) -> super::quota_policy::ObservationFailure {
use super::quota_policy::{AccountFailureKind as K, ObservationFailure};
use super::token::RefreshFailureKind as R;
match kind {
R::Transient => ObservationFailure::new(K::RefreshTransient),
R::LockBusy => ObservationFailure::new(K::RefreshLockBusy),
R::Other => ObservationFailure::new(K::RefreshOther),
R::Cooldown { retry_in_secs } => {
ObservationFailure::new(K::RefreshOther).with_retry_after(Some(retry_in_secs))
}
R::Credential => ObservationFailure::new(K::RefreshRejected),
R::ActionRequired => ObservationFailure::new(K::RefreshActionRequired),
}
}
fn auto_vend_failure(
credential: &CredentialRef,
f: super::token::RefreshFailure,
) -> super::token::RefreshFailure {
let label = credential.account.label_str();
let mut message = format!("account '{label}': {}", f.message());
if f.kind == super::token::RefreshFailureKind::Credential && !message.contains("synaps login") {
let account = if credential.account.is_default() {
String::new()
} else {
format!(" --account {label}")
};
message.push_str(&format!(
"; log it in again with `synaps login --provider {}{account}`",
credential.provider
));
}
super::token::RefreshFailure::new(f.kind, message).with_http_status(f.http_status())
}
fn selection_failed(
provider: OAuthProviderId,
model: Option<&str>,
reason: String,
failure: super::quota_policy::SelectionFailure,
) -> BrokerError {
log_selection_failure(provider, model, &failure);
BrokerError::NoAccountAvailable {
provider: provider.to_string(),
reason,
failure: Some(failure),
}
}
fn log_selection_failure(
provider: OAuthProviderId,
model: Option<&str>,
failure: &super::quota_policy::SelectionFailure,
) {
tracing::warn!(
provider = %provider,
model = model.unwrap_or("any"),
kind = %failure.kind,
accounts = %failure.accounts_field(),
earliest_reset_ms = ?failure.earliest_reset_ms,
"auth_selection_failed"
);
}
async fn read_body_capped(resp: reqwest::Response, cap: usize) -> Result<String, BrokerError> {
use futures::StreamExt;
let mut buf: Vec<u8> = Vec::new();
let mut stream = resp.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk =
chunk.map_err(|e| BrokerError::Transport(format!("failed to read response: {e}")))?;
if buf.len() + chunk.len() > cap {
return Err(BrokerError::Transport(format!(
"response body exceeded the {cap}-byte broker buffering limit"
)));
}
buf.extend_from_slice(&chunk);
}
String::from_utf8(buf)
.map_err(|_| BrokerError::Transport("response body was not valid UTF-8".into()))
}
const MAX_PROXY_ERROR_CLASSIFY_BYTES: usize = 8 * 1024;
async fn read_error_body_for_classification(resp: reqwest::Response, cap: usize) -> Option<String> {
use futures::StreamExt;
let mut buf: Vec<u8> = Vec::with_capacity(cap.min(1024));
let mut stream = resp.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk.ok()?;
let room = cap.saturating_sub(buf.len());
if chunk.len() >= room {
buf.extend_from_slice(&chunk[..room]);
return String::from_utf8(buf).ok();
}
buf.extend_from_slice(&chunk);
}
String::from_utf8(buf).ok()
}
fn discover_legacy_static_key(spec: &StaticProviderSpec) -> Option<String> {
if let Some(v) = crate::config::get_provider_keys().get(spec.key) {
if !v.is_empty() {
return Some(v.clone());
}
}
spec.env_vars
.iter()
.find_map(|var| std::env::var(var).ok().filter(|v| !v.is_empty()))
}
pub fn static_key_configured(provider: &str) -> bool {
if provider == LOCAL_PROVIDER_KEY {
return true; }
let Some(spec) = static_provider(provider) else {
return false;
};
storage::load_static_key(provider).ok().flatten().is_some()
|| discover_legacy_static_key(spec).is_some()
}
pub fn static_key_status(provider: &str) -> StaticKeyStatus {
let stored = storage::load_static_key(provider)
.ok()
.flatten()
.or_else(|| {
crate::config::get_provider_keys()
.get(provider)
.filter(|s| !s.is_empty())
.cloned()
});
if let Some(key) = stored {
return StaticKeyStatus::Configured {
masked: mask_key(&key),
};
}
if let Some(spec) = static_provider(provider) {
if spec
.env_vars
.iter()
.any(|var| std::env::var(var).is_ok_and(|v| !v.is_empty()))
{
return StaticKeyStatus::FromEnv;
}
}
StaticKeyStatus::NotSet
}
fn mask_key(key: &str) -> String {
let chars: Vec<char> = key.chars().collect();
if chars.len() <= 8 {
return "…".to_string();
}
let head: String = chars[..4].iter().collect();
let tail: String = chars[chars.len() - 4..].iter().collect();
format!("{head}…{tail}")
}
pub fn configured_static_provider_keys() -> std::collections::BTreeSet<String> {
STATIC_PROVIDERS
.iter()
.filter(|s| static_key_configured(s.key))
.map(|s| s.key.to_string())
.collect()
}
pub fn local_endpoint_url() -> String {
local_endpoint_config()
.unwrap_or_else(|| LOCAL_DEFAULT_BASE_URL.to_string())
.trim_end_matches('/')
.to_string()
}
pub fn local_endpoint_config() -> Option<String> {
crate::config::get_provider_keys()
.get("local.url")
.filter(|s| !s.is_empty())
.cloned()
.or_else(|| {
std::env::var("LOCAL_ENDPOINT")
.ok()
.filter(|s| !s.is_empty())
})
}
pub fn local_model_ids() -> Vec<String> {
crate::config::get_provider_keys()
.get("local.models")
.map(|value| {
value
.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect()
})
.unwrap_or_default()
}
pub fn anthropic_credential_available() -> bool {
oauth_provider_logged_in(OAuthProviderId::Anthropic)
|| std::env::var("ANTHROPIC_API_KEY").is_ok_and(|v| !v.is_empty())
}
pub fn oauth_provider_logged_in(provider: OAuthProviderId) -> bool {
oauth_provider_logged_in_at(
&storage::auth_file_path(),
&AccountPolicy::from_environment(),
provider,
)
}
pub(crate) fn oauth_provider_logged_in_at(
path: &std::path::Path,
policy: &AccountPolicy,
provider: OAuthProviderId,
) -> bool {
let Ok(inventory) = storage::list_accounts_detailed_at(path, Some(provider)) else {
return false;
};
match policy.selector(provider) {
AccountSelector::Account(Account::Named(label)) => inventory
.accounts
.iter()
.any(|row| row.label == label.as_str()),
AccountSelector::Account(Account::Default)
| AccountSelector::Auto
| AccountSelector::Invalid { .. } => !inventory.accounts.is_empty(),
}
}
#[async_trait]
impl CredentialBroker for LocalBroker {
async fn recover_rejected_token(
&self,
cred: &CredentialRef,
rejected: &str,
) -> Result<AccessToken, BrokerError> {
self.check_slot_exists(cred)?;
let creds = super::token::recover_rejected_credential(
&self.http,
cred,
rejected,
&self.refresh_endpoints,
)
.await
.map_err(BrokerError::Refresh)?;
Ok(AccessToken {
token: creds.access,
expires: creds.expires,
})
}
async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
Ok(self.access_token_pinned(provider).await?.token)
}
fn account_selector(&self, provider: OAuthProviderId) -> AccountSelector {
self.policy().selector(provider)
}
async fn access_token_for(&self, cred: &CredentialRef) -> Result<AccessToken, BrokerError> {
self.check_slot_exists(cred)?;
let creds =
super::token::ensure_fresh_credential_via(&self.http, cred, &self.refresh_endpoints)
.await
.map_err(|f| match f.kind {
super::token::RefreshFailureKind::ActionRequired => BrokerError::Refresh(f),
_ => BrokerError::Credential(f.into()),
})?;
Ok(AccessToken {
token: creds.access,
expires: creds.expires,
})
}
async fn access_token_pinned(
&self,
provider: OAuthProviderId,
) -> Result<PinnedToken, BrokerError> {
self.access_token_pinned_for(provider, None).await
}
async fn access_token_pinned_for(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
self.resolve_and_vend(provider, None, model).await
}
async fn accounts(&self, provider: OAuthProviderId) -> Result<Vec<AccountSummary>, BrokerError> {
let mut rows = storage::list_accounts(provider).map_err(BrokerError::Credential)?;
let selector = self.policy().selector(provider);
let now_ms = crate::epoch_millis();
for row in &mut rows {
row.selected = matches!(&selector, AccountSelector::Account(a) if a.label_str() == row.label);
if let Some(cred) = row.credential_ref() {
row.cooldown_until = self.cooldown_until(&cred.storage_key(), now_ms);
}
}
Ok(rows)
}
async fn usage(&self, cred: &CredentialRef) -> Result<super::usage::UsageSnapshot, BrokerError> {
self.usage_with_seat(cred)
.await
.map(|(snapshot, _seat)| snapshot)
.map_err(|e| e.error)
}
async fn report_cooldown(
&self,
cred: &CredentialRef,
until_ms: Option<u64>,
reason: &str,
) -> Result<(), BrokerError> {
let now_ms = crate::epoch_millis();
let max_until = now_ms.saturating_add(MAX_COOLDOWN.as_millis() as u64);
let until_ms = until_ms
.filter(|t| *t > now_ms)
.unwrap_or(now_ms + DEFAULT_COOLDOWN.as_millis() as u64)
.min(max_until);
let reason = crate::truncate_str(reason, 64).to_string();
tracing::info!(credential = %cred, until_ms, reason = %reason, "account cooldown reported");
let key = self.cooldown_key(cred);
self.cooldowns
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(key.clone(), Cooldown { until_ms, reason });
self.forget_snapshots(cred);
self.usage_reads.forget_last_good(&self.read_key(cred));
Ok(())
}
async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
let resp = self.send(&request).await?;
let status = resp.status().as_u16();
let body = read_body_capped(resp, self.max_response_bytes).await?;
Ok(ProxyResponse { status, body })
}
async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
let mut request = request;
request.stream = true;
let resp = self.send(&request).await?;
let status = resp.status();
if !status.is_success() {
let label = read_error_body_for_classification(resp, MAX_PROXY_ERROR_CLASSIFY_BYTES)
.await
.as_deref()
.and_then(crate::error::vetted_proxy_error_label);
tracing::warn!(
provider = %request.provider,
status = status.as_u16(),
label = label.unwrap_or("-"),
"provider request failed"
);
let label_suffix = label.map(|l| format!(" [{l}]")).unwrap_or_default();
return Err(BrokerError::Transport(format!(
"provider request failed: {status}{label_suffix}"
)));
}
use futures::StreamExt;
let stream = resp
.bytes_stream()
.map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
Ok(Box::pin(stream))
}
async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
let token = self.access_token(OAuthProviderId::Anthropic).await?;
let url = self
.anthropic_usage_url
.as_deref()
.unwrap_or(ANTHROPIC_USAGE_URL);
let resp = self
.http
.get(url)
.timeout(self.request_timeout)
.bearer_auth(&token.token)
.header("anthropic-beta", ANTHROPIC_OAUTH_BETA)
.send()
.await
.map_err(|e| BrokerError::Transport(format!("usage request failed: {e}")))?;
let status = resp.status();
if !status.is_success() {
drop(resp);
return Err(BrokerError::Transport(format!(
"usage request failed: {status}"
)));
}
let body = read_body_capped(resp, self.max_response_bytes).await?;
serde_json::from_str(&body)
.map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
}
async fn cloud_catalog(
&self,
provider: CloudProviderId,
context_ref: &str,
allow_stale: bool,
) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
self.cloud_backend
.as_ref()
.ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
.catalog(provider, context_ref, allow_stale)
.await
}
async fn cloud_invoke(
&self,
provider: CloudProviderId,
context_ref: &str,
model_id: &str,
request: InvokeRequest,
) -> Result<CloudEventStream, BrokerError> {
self.cloud_backend
.as_ref()
.ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
.invoke(provider, context_ref, model_id, request)
.await
}
async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
let mut out = Vec::new();
let policy = self.policy();
for descriptor in super::provider::registry().iter() {
let accounts = self.accounts(descriptor.id).await.unwrap_or_default();
let configured = match policy.selector(descriptor.id) {
AccountSelector::Account(account) => accounts
.iter()
.any(|a| a.label == account.label_str()),
AccountSelector::Auto => !accounts.is_empty(),
AccountSelector::Invalid { .. } => false,
};
out.push(ProviderStatus {
key: descriptor.id.as_str().to_string(),
name: descriptor.display_name.to_string(),
kind: CredentialKind::OAuth,
configured,
accounts,
});
}
for spec in STATIC_PROVIDERS {
out.push(ProviderStatus {
key: spec.key.to_string(),
name: spec.name.to_string(),
kind: CredentialKind::StaticKey,
configured: static_key_configured(spec.key),
accounts: Vec::new(),
});
}
out.push(ProviderStatus {
key: LOCAL_PROVIDER_KEY.to_string(),
name: "Local endpoint".to_string(),
kind: CredentialKind::LocalEndpoint,
configured: true,
accounts: Vec::new(),
});
Ok(out)
}
}
pub struct RemoteBroker {
http: reqwest::Client,
endpoint: String,
machine_token: super::MachineToken,
cache: super::TokenCache,
account_policy: Option<AccountPolicy>,
}
async fn route_missing(resp: reqwest::Response) -> bool {
let json = resp
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.is_some_and(|v| v.contains("json"));
!json
&& read_body_capped(resp, 1024)
.await
.is_ok_and(|body| body.trim().is_empty())
}
impl RemoteBroker {
async fn predates_usage_snapshot(&self) -> Result<bool, BrokerError> {
let resp = self
.machine_token
.send(|t| {
self.http
.get(format!("{}/usage/snapshot", self.endpoint))
.query(&[("provider", "anthropic"), ("account", "")])
.bearer_auth(t)
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => Err(BrokerError::Unauthorized),
404 => Ok(route_missing(resp).await),
_ => Ok(false),
}
}
pub fn new(
endpoint: impl Into<String>,
machine_token: impl Into<String>,
http: reqwest::Client,
cache: super::TokenCache,
) -> Self {
let endpoint = endpoint.into().trim_end_matches('/').to_string();
let http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(Duration::from_secs(10))
.timeout(Duration::from_secs(90))
.build()
.unwrap_or(http);
let machine_token = super::MachineToken::for_endpoint(&endpoint, machine_token);
Self {
http,
endpoint,
machine_token,
cache,
account_policy: None,
}
}
pub fn with_account_policy(mut self, policy: AccountPolicy) -> Self {
self.account_policy = Some(policy);
self
}
fn policy(&self) -> AccountPolicy {
self.account_policy
.clone()
.unwrap_or_else(AccountPolicy::from_environment)
}
fn fetcher(&self) -> super::BrokerClient {
super::BrokerClient::with_machine_token(
self.endpoint.clone(),
self.machine_token.clone(),
self.http.clone(),
)
}
fn scope(&self) -> String {
super::credential_source::source_scope(&self.endpoint, &self.machine_token.get())
}
async fn resolve_pinned(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
match self.policy().selector(provider) {
AccountSelector::Account(account) => {
let credential = CredentialRef::new(provider, account);
let token = self.access_token_for(&credential).await?;
Ok(PinnedToken { credential, token })
}
AccountSelector::Auto => self.fetch_auto(provider, model).await,
AccountSelector::Invalid { source, reason } => {
Err(BrokerError::InvalidAccount(format!("{source}: {reason}")))
}
}
}
async fn fetch_auto(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
let mut query: Vec<(&str, &str)> = vec![
("provider", provider.as_str()),
("account", super::account::AUTO_ACCOUNT_NAME),
];
if let Some(model) = model {
query.push(("model", model));
}
let resp = self
.machine_token
.send(|t| {
self.http
.get(format!("{}/token", self.endpoint))
.query(&query)
.bearer_auth(t)
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => return Err(BrokerError::Unauthorized),
503 => {
drop(resp);
return Err(BrokerError::NoAccountAvailable {
provider: provider.as_str().to_string(),
reason: "broker reported no account with proven capacity".into(),
failure: Some(super::quota_policy::SelectionFailure::new(
super::quota_policy::SelectionFailureKind::RemoteNoCapacity,
)),
});
}
s if !(200..300).contains(&s) => {
drop(resp);
return Err(BrokerError::Transport(format!("broker returned HTTP {s}")));
}
_ => {}
}
#[derive(Deserialize)]
struct AutoToken {
access_token: String,
expires: u64,
#[serde(default)]
ttl_ms: Option<u64>,
account: Option<String>,
}
let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
let tok: AutoToken = serde_json::from_str(&body)
.map_err(|e| BrokerError::Transport(format!("invalid broker token response: {e}")))?;
let account = tok
.account
.as_deref()
.ok_or_else(|| BrokerError::Transport("broker did not report the selected account".into()))
.and_then(|a| Account::parse(a).map_err(BrokerError::InvalidAccount))?;
if tok.access_token.is_empty() {
return Err(BrokerError::Transport("broker returned an empty access_token".into()));
}
let expires = match tok.ttl_ms {
Some(ttl) => crate::epoch_millis().saturating_add(ttl),
None => tok.expires,
};
Ok(PinnedToken {
credential: CredentialRef::new(provider, account),
token: AccessToken {
token: tok.access_token,
expires,
},
})
}
async fn proxy_failure(resp: reqwest::Response, fallback: String) -> BrokerError {
if resp.status() != reqwest::StatusCode::BAD_GATEWAY {
return BrokerError::Transport(fallback);
}
let Ok(body) = read_body_capped(resp, 4096).await else {
return BrokerError::Transport(fallback);
};
let text = serde_json::from_str::<serde_json::Value>(&body)
.ok()
.and_then(|v| v.get("error")?.as_str().map(str::to_string))
.unwrap_or_default();
use super::token::{RefreshFailure, RefreshFailureKind as K};
for kind in [K::Transient, K::LockBusy, K::Other] {
let prefix = BrokerError::Refresh(RefreshFailure::new(kind, "")).to_string();
if text.starts_with(&prefix) {
return BrokerError::Refresh(RefreshFailure::new(
kind,
"the remote credential broker could not refresh this credential",
));
}
}
BrokerError::Transport(fallback)
}
async fn post_proxy(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
request.validate()?;
let resp = self
.machine_token
.send(|t| {
self.http
.post(format!("{}/proxy", self.endpoint))
.bearer_auth(t)
.json(request)
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => Err(BrokerError::Unauthorized),
s @ (400 | 403) => {
drop(resp);
Err(BrokerError::Denied(format!(
"broker rejected the request (HTTP {s})"
)))
}
_ => Ok(resp),
}
}
}
#[async_trait]
impl CredentialBroker for RemoteBroker {
async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
Ok(self.access_token_pinned(provider).await?.token)
}
fn account_selector(&self, provider: OAuthProviderId) -> AccountSelector {
self.policy().selector(provider)
}
async fn access_token_for(&self, cred: &CredentialRef) -> Result<AccessToken, BrokerError> {
let fetcher = self.fetcher();
let tok = super::credential_source::resolve_remote_credential(
&fetcher,
&self.cache,
&self.scope(),
cred,
super::DEFAULT_MARGIN_MS,
)
.await
.map_err(|e| e.into_broker_error(cred))?;
Ok(AccessToken {
token: tok.access_token,
expires: tok.expires,
})
}
async fn access_token_pinned(
&self,
provider: OAuthProviderId,
) -> Result<PinnedToken, BrokerError> {
self.resolve_pinned(provider, None).await
}
async fn access_token_pinned_for(
&self,
provider: OAuthProviderId,
model: Option<&str>,
) -> Result<PinnedToken, BrokerError> {
self.resolve_pinned(provider, model).await
}
async fn accounts(&self, provider: OAuthProviderId) -> Result<Vec<AccountSummary>, BrokerError> {
let caps = self.capabilities().await?;
let accounts = caps
.into_iter()
.find(|c| c.key == provider.as_str())
.map(|c| c.accounts)
.unwrap_or_default();
if accounts.is_empty() && self.predates_usage_snapshot().await? {
return Err(BrokerError::BrokerTooOld("account listing".into()));
}
Ok(accounts)
}
async fn usage(&self, cred: &CredentialRef) -> Result<super::usage::UsageSnapshot, BrokerError> {
let resp = self
.machine_token
.send(|t| {
self.http
.get(format!("{}/usage/snapshot", self.endpoint))
.query(&[
("provider", cred.provider.as_str()),
("account", cred.account.label_str()),
])
.bearer_auth(t)
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => Err(BrokerError::Unauthorized),
404 => {
if !route_missing(resp).await {
return Err(BrokerError::UnknownAccount {
provider: cred.provider.as_str().to_string(),
label: cred.account.label_str().to_string(),
});
}
if cred.provider != OAuthProviderId::Anthropic || !cred.account.is_default() {
return Err(BrokerError::BrokerTooOld("per-account usage".into()));
}
let body = self.anthropic_usage().await?.to_string();
super::usage::parse_anthropic_usage(
&body,
cred.account.label_str(),
crate::epoch_millis(),
)
.map_err(super::usage::UsageError::into_broker_error)
}
s if !(200..300).contains(&s) => {
drop(resp);
Err(BrokerError::Transport(format!(
"broker usage snapshot returned HTTP {s}"
)))
}
_ => {
let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
let snapshot: super::usage::UsageSnapshot = serde_json::from_str(&body)
.map_err(|e| BrokerError::Transport(format!("invalid usage snapshot: {e}")))?;
if snapshot.provider != cred.provider.as_str()
|| snapshot.account != cred.account.label_str()
{
return Err(BrokerError::Transport(
"broker returned a usage snapshot for a different account".into(),
));
}
Ok(snapshot)
}
}
}
async fn report_cooldown(
&self,
cred: &CredentialRef,
until_ms: Option<u64>,
reason: &str,
) -> Result<(), BrokerError> {
let resp = self
.machine_token
.send(|t| {
self.http
.post(format!("{}/accounts/cooldown", self.endpoint))
.bearer_auth(t)
.json(&serde_json::json!({
"provider": cred.provider.as_str(),
"account": cred.account.label_str(),
"until_ms": until_ms,
"reason": crate::truncate_str(reason, 64),
}))
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => Err(BrokerError::Unauthorized),
s if !(200..300).contains(&s) => {
Err(BrokerError::Transport(format!("broker returned HTTP {s}")))
}
_ => Ok(()),
}
}
async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
let mut request = request;
request.stream = false;
let resp = self.post_proxy(&request).await?;
if !resp.status().is_success() {
let fallback = format!("broker proxy returned HTTP {}", resp.status());
return Err(Self::proxy_failure(resp, fallback).await);
}
let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
serde_json::from_str::<ProxyResponse>(&body)
.map_err(|e| BrokerError::Transport(format!("invalid broker proxy response: {e}")))
}
async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
let mut request = request;
request.stream = true;
let resp = self.post_proxy(&request).await?;
let status = resp.status();
if !status.is_success() {
let fallback = format!("broker proxy stream failed: {status}");
return Err(Self::proxy_failure(resp, fallback).await);
}
use futures::StreamExt;
let stream = resp
.bytes_stream()
.map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
Ok(Box::pin(stream))
}
async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
let resp = self
.machine_token
.send(|t| {
self.http
.get(format!("{}/usage", self.endpoint))
.bearer_auth(t)
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => Err(BrokerError::Unauthorized),
s if !(200..300).contains(&s) => {
drop(resp);
Err(BrokerError::Transport(format!(
"broker usage returned HTTP {s}"
)))
}
_ => {
let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
serde_json::from_str(&body)
.map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
}
}
}
async fn cloud_catalog(
&self,
provider: CloudProviderId,
context_ref: &str,
allow_stale: bool,
) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
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()))?;
if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
return Err(BrokerError::Unauthorized);
}
if !resp.status().is_success() {
return Err(BrokerError::Transport(format!(
"cloud catalog returned HTTP {}",
resp.status()
)));
}
let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
let entries: Vec<CloudCatalogEntry> = serde_json::from_str(&body)
.map_err(|e| BrokerError::Transport(format!("invalid cloud catalog: {e}")))?;
if entries.len() > MAX_CLOUD_CATALOG_ENTRIES
|| entries.iter().any(|e| {
e.provider != provider || !e.context_ref.starts_with("ctx-") || e.id.len() > 512
})
{
return Err(BrokerError::Transport(
"broker returned an invalid cloud catalog".into(),
));
}
Ok(entries)
}
async fn cloud_invoke(
&self,
provider: CloudProviderId,
context_ref: &str,
model_id: &str,
request: InvokeRequest,
) -> Result<CloudEventStream, BrokerError> {
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()))?;
if !resp.status().is_success() {
return Err(if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
BrokerError::Unauthorized
} else {
BrokerError::Transport(format!("cloud invoke returned HTTP {}", resp.status()))
});
}
use futures::StreamExt;
let chunks = Box::pin(resp.bytes_stream());
let stream = futures::stream::unfold(
(chunks, Vec::<u8>::new(), false, false),
|(mut chunks, mut buffer, done, failed)| async move {
if failed {
return None;
}
loop {
if let Some(end) = buffer.iter().position(|b| *b == b'\n') {
let line: Vec<u8> = buffer.drain(..=end).collect();
let parsed = serde_json::from_slice::<CloudEvent>(&line[..end])
.map_err(|_| BrokerError::Transport("invalid cloud event".into()));
let event = match parsed {
Ok(CloudEvent::Done) if done => Err(BrokerError::Transport(
"duplicate cloud terminal event".into(),
)),
Ok(CloudEvent::Done) => Ok(CloudEvent::Done),
Ok(_) if done => Err(BrokerError::Transport(
"cloud data followed terminal event".into(),
)),
other => other,
};
let next_done = done || matches!(event, Ok(CloudEvent::Done));
let next_failed = event.is_err();
return Some((event, (chunks, buffer, next_done, next_failed)));
}
match chunks.next().await {
Some(Ok(chunk))
if buffer.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES =>
{
buffer.extend_from_slice(&chunk)
}
Some(Ok(_)) => {
return Some((
Err(BrokerError::Transport("cloud event exceeded limit".into())),
(chunks, Vec::new(), done, true),
))
}
Some(Err(_)) => {
return Some((
Err(BrokerError::Transport("broker stream failed".into())),
(chunks, Vec::new(), done, true),
))
}
None if buffer.is_empty() && done => return None,
None if buffer.is_empty() => {
return Some((
Err(BrokerError::Transport(
"cloud stream ended without terminal event".into(),
)),
(chunks, Vec::new(), done, true),
))
}
None => {
return Some((
Err(BrokerError::Transport("truncated cloud event".into())),
(chunks, Vec::new(), done, true),
))
}
}
}
},
);
Ok(Box::pin(stream))
}
async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
let resp = self
.machine_token
.send(|t| {
self.http
.get(format!("{}/capabilities", self.endpoint))
.bearer_auth(t)
})
.await
.map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
match resp.status().as_u16() {
401 => Err(BrokerError::Unauthorized),
s if !(200..300).contains(&s) => {
Err(BrokerError::Transport(format!("broker returned HTTP {s}")))
}
_ => {
let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
serde_json::from_str::<Vec<ProviderStatus>>(&body).map_err(|e| {
BrokerError::Transport(format!("invalid capabilities response: {e}"))
})
}
}
}
}
pub fn preflight_cloud_capability(
provider: CloudProviderId,
needs_tools: bool,
) -> Result<(), BrokerError> {
if needs_tools && !provider.supports_tools() {
return Err(BrokerError::UnsupportedCapability {
provider: provider.to_string(),
capability: "tools".into(),
});
}
Ok(())
}
pub fn local_auth_state(
cache: &super::TokenCache,
query: &super::auth_state::AuthStateQuery,
live: bool,
) -> Result<super::auth_state::AuthStateReport, BrokerError> {
let key = (
storage::auth_file_path(),
crate::config::resolve_read_path("config"),
);
let broker = cache
.local_brokers
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&key)
.cloned();
let mut report = auth_state_report(broker.as_deref(), query)?;
report.live = live;
Ok(report)
}
fn auth_state_report(
broker: Option<&LocalBroker>,
query: &super::auth_state::AuthStateQuery,
) -> Result<super::auth_state::AuthStateReport, BrokerError> {
use super::auth_state::{
AccountState, AuthStateReport, IdentityKind, ObservedOutcome, AUTH_STATE_SCHEMA_VERSION,
};
let provider = match query.provider.as_deref() {
Some(p) => Some(
p.parse::<OAuthProviderId>()
.map_err(|_| BrokerError::UnknownProvider(p.to_string()))?,
),
None => None,
};
let mut summaries = match provider {
Some(p) => storage::list_accounts(p),
None => storage::list_all_accounts(),
}
.map_err(BrokerError::Credential)?;
if let Some(label) = query.account.as_deref() {
summaries.retain(|s| s.label == label);
if summaries.is_empty() {
return Err(BrokerError::UnknownAccount {
provider: provider.map_or("any provider", |p| p.as_str()).to_string(),
label: label.to_string(),
});
}
}
let creds: Vec<CredentialRef> = summaries
.iter()
.filter_map(AccountSummary::credential_ref)
.collect();
let mut cleared = Vec::new();
if query.clear {
if let Some(broker) = broker {
let whole: Vec<OAuthProviderId> = match (&query.account, provider) {
(Some(_), _) => Vec::new(),
(None, Some(p)) => vec![p],
(None, None) => super::provider::DESCRIPTORS.iter().map(|d| d.id).collect(),
};
broker.clear_state(&creds, &whole);
}
cleared = creds
.iter()
.map(|c| format!("{}:{}", c.provider, c.account.label_str()))
.collect();
tracing::info!(
accounts = %cleared.join(","),
broker_active = broker.is_some(),
"auth_state_cleared"
);
}
let store = storage::auth_file_path();
let now_ms = crate::epoch_millis();
let mut accounts = Vec::with_capacity(summaries.len());
for summary in &summaries {
let Some(cred) = summary.credential_ref() else {
continue;
};
let observed = broker
.and_then(|b| {
b.observations
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&LocalBroker::state_key(&cred))
.cloned()
})
.unwrap_or_default();
accounts.push(AccountState {
provider: summary.provider.clone(),
label: summary.label.clone(),
identity: if summary.seat_fingerprint.is_some() {
IdentityKind::Stable
} else {
IdentityKind::Rotating
},
access_expires_ms: (summary.expires != 0).then_some(summary.expires),
cooldown: broker.and_then(|b| b.cooldown_state(&cred, now_ms)),
last_usage: observed.last_usage,
last_usage_ok_ms: observed.last_usage_ok_ms,
last_refresh: super::token::last_refresh_exchange(&store, &cred.storage_key()).map(
|x| ObservedOutcome {
at_ms: x.at_ms,
failure: x.failure.map(|kind| {
let failure = refresh_kind_failure(kind);
match x.http_status {
Some(status) => failure.with_status(status),
None => failure,
}
}),
},
),
});
}
let store_prefix = format!("{}|", store.display());
let mut selections: Vec<_> = broker
.map(|b| {
b.selections
.lock()
.unwrap_or_else(|e| e.into_inner())
.iter()
.filter(|(k, s)| {
k.starts_with(&store_prefix)
&& provider.is_none_or(|p| s.provider == p.as_str())
&& query
.account
.as_deref()
.is_none_or(|label| selection_involves(s, label))
})
.map(|(_, s)| s.clone())
.collect()
})
.unwrap_or_default();
selections.sort_by_key(|s| std::cmp::Reverse(s.at_ms));
Ok(AuthStateReport {
schema_version: AUTH_STATE_SCHEMA_VERSION,
live: false,
broker_active: broker.is_some(),
generated_at_ms: now_ms,
store: store.display().to_string(),
accounts,
selections,
cleared,
})
}
fn selection_involves(s: &super::auth_state::SelectionState, label: &str) -> bool {
s.chosen.as_deref() == Some(label)
|| s.failure
.as_ref()
.is_some_and(|f| f.accounts.iter().any(|a| a.account == label))
}
pub fn broker_from_source(
source: &super::CredentialSource,
cache: &super::TokenCache,
http: reqwest::Client,
) -> Arc<dyn CredentialBroker> {
match source {
super::CredentialSource::Local => {
let key = (
storage::auth_file_path(),
crate::config::resolve_read_path("config"),
);
let mut brokers = cache
.local_brokers
.lock()
.unwrap_or_else(|e| e.into_inner());
brokers
.entry(key)
.or_insert_with(|| Arc::new(LocalBroker::new(http)))
.clone()
}
super::CredentialSource::Remote {
endpoint,
machine_token,
} => Arc::new(RemoteBroker::new(
endpoint.clone(),
machine_token.clone(),
http,
cache.clone(),
)),
}
}
static GLOBAL_BROKER: std::sync::RwLock<Option<Arc<dyn CredentialBroker>>> =
std::sync::RwLock::new(None);
static GLOBAL_BROKER_INSTALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub fn set_global_broker(broker: Arc<dyn CredentialBroker>) {
*GLOBAL_BROKER.write().expect("broker registry poisoned") = Some(broker);
let n = GLOBAL_BROKER_INSTALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
if crate::core::memstat::mem_trace_enabled() {
tracing::info!(target: "agent_core::memstat", installs = n, "global broker installed");
}
}
pub fn global_broker_install_count() -> u64 {
GLOBAL_BROKER_INSTALLS.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn global_broker() -> Arc<dyn CredentialBroker> {
if let Some(b) = GLOBAL_BROKER
.read()
.expect("broker registry poisoned")
.clone()
{
return b;
}
let default: Arc<dyn CredentialBroker> = Arc::new(LocalBroker::new(
reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(10))
.timeout(std::time::Duration::from_secs(300))
.build()
.unwrap_or_default(),
));
let mut guard = GLOBAL_BROKER.write().expect("broker registry poisoned");
if let Some(b) = guard.clone() {
return b;
}
*guard = Some(default.clone());
default
}
pub async fn broker_access_token(
provider: OAuthProviderId,
source: &super::CredentialSource,
cache: &super::TokenCache,
http: &reqwest::Client,
) -> Result<AccessToken, BrokerError> {
broker_from_source(source, cache, http.clone())
.access_token(provider)
.await
}
pub fn static_key_status_map() -> BTreeMap<String, StaticKeyStatus> {
STATIC_PROVIDERS
.iter()
.map(|s| (s.key.to_string(), static_key_status(s.key)))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn login_availability_handles_named_only_and_explicit_policy() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("auth.json");
let provider = OAuthProviderId::Anthropic;
let empty = AccountPolicy::new();
assert!(!oauth_provider_logged_in_at(&path, &empty, provider));
std::fs::write(&path, serde_json::json!({
"anthropic@work": {"type":"oauth", "access":"synthetic-access",
"refresh":"synthetic-refresh", "expires":1}
}).to_string()).unwrap();
assert!(oauth_provider_logged_in_at(&path, &empty, provider));
for selector in [AccountSelector::Auto,
AccountSelector::parse_or_invalid("../bad", "test"),
AccountSelector::Account(Account::named("work").unwrap())] {
let policy = AccountPolicy::new().with(provider, selector);
assert!(oauth_provider_logged_in_at(&path, &policy, provider));
}
let missing = AccountPolicy::new().with(provider,
AccountSelector::Account(Account::named("missing").unwrap()));
assert!(!oauth_provider_logged_in_at(&path, &missing, provider));
assert!(!oauth_provider_logged_in_at(&path, &empty, OAuthProviderId::OpenAiCodex));
std::fs::write(&path, "{broken").unwrap();
assert!(!oauth_provider_logged_in_at(&path, &empty, provider));
}
#[test]
fn preflight_rejects_tool_requiring_cloud_routes_with_typed_error() {
for provider in [
CloudProviderId::AzureOpenAi,
CloudProviderId::AwsBedrock,
CloudProviderId::GoogleVertex,
] {
let err = preflight_cloud_capability(provider, true)
.expect_err("tool-requiring cloud route must fail pre-flight");
assert_eq!(
err,
BrokerError::UnsupportedCapability {
provider: provider.to_string(),
capability: "tools".into(),
}
);
assert!(err.to_string().contains("text-only"));
assert_eq!(preflight_cloud_capability(provider, false), Ok(()));
}
}
async fn spawn_ndjson(body: &'static str) -> String {
use axum::{body::Body, routing::post, Router};
use bytes::Bytes;
let app = Router::new().route(
"/cloud/invoke",
post(move || async move {
let chunks = body
.as_bytes()
.chunks(3)
.map(|chunk| Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(chunk)));
axum::response::Response::builder()
.header("content-type", "application/x-ndjson")
.body(Body::from_stream(futures::stream::iter(chunks)))
.unwrap()
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
format!("http://{addr}")
}
async fn remote_events(body: &'static str) -> Vec<Result<CloudEvent, BrokerError>> {
use futures::StreamExt;
let broker = RemoteBroker::new(
spawn_ndjson(body).await,
"opaque-machine-token",
reqwest::Client::new(),
super::super::TokenCache::new(),
);
broker
.cloud_invoke(
CloudProviderId::GoogleVertex,
"ctx-opaque",
"google-vertex/publishers/google/models/test",
InvokeRequest {
messages: vec![],
tools: vec![],
stream: true,
options: Default::default(),
},
)
.await
.unwrap()
.collect()
.await
}
#[tokio::test]
async fn remote_ndjson_requires_exactly_one_terminal_and_stops_after_error() {
let valid =
remote_events("{\"type\":\"text_delta\",\"delta\":\"hi\"}\n{\"type\":\"done\"}\n")
.await;
assert_eq!(valid.len(), 2);
assert!(matches!(valid[1], Ok(CloudEvent::Done)));
for body in [
"{\"type\":\"text_delta\",\"delta\":\"hi\"}\n",
"{\"type\":\"done\"}\n{\"type\":\"done\"}\n",
"{\"type\":\"done\"}\n{\"type\":\"text_delta\",\"delta\":\"late\"}\n",
] {
let events = remote_events(body).await;
assert_eq!(events.iter().filter(|event| event.is_err()).count(), 1);
assert!(events.last().unwrap().is_err());
}
}
#[test]
fn proxy_rejects_unlisted_same_host_paths() {
for (provider, path) in [
("groq", "/v1/keys"),
("groq", "/admin"),
("openrouter", "/api/v1/auth/keys"),
("local", "/audio/speech"),
] {
let req = ProxyRequest {
provider: provider.into(),
method: ProxyMethod::Get,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
match req.validate() {
Err(BrokerError::Denied(msg)) => {
assert!(msg.contains("allowlist"), "got: {msg}")
}
other => panic!("{provider} {path} must be denied, got {other:?}"),
}
}
for path in ["/models", "/chat/completions"] {
let req = ProxyRequest {
provider: "groq".into(),
method: ProxyMethod::Get,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(req.validate().is_ok(), "{path} must be allowed");
}
}
#[test]
fn proxy_rejects_oversize_request_body() {
let req = ProxyRequest {
provider: "groq".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: Some(serde_json::json!({
"blob": "x".repeat(MAX_PROXY_REQUEST_BYTES + 1)
})),
stream: false,
body_bytes: None,
};
match req.validate() {
Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
other => panic!("oversize body must be denied, got {other:?}"),
}
}
#[test]
fn proxy_rejects_oversize_body_bytes_without_json_body() {
let req = ProxyRequest {
provider: "groq".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: false,
body_bytes: Some(bytes::Bytes::from(vec![b'x'; MAX_PROXY_REQUEST_BYTES + 1])),
};
match req.validate() {
Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
other => panic!("oversize body_bytes must be denied, got {other:?}"),
}
}
#[test]
fn post_json_exact_sets_coherent_body_and_bytes() {
let value = serde_json::json!({"model": "m", "stream": true});
let (req, digest_bytes) =
ProxyRequest::post_json_exact("groq", "/chat/completions", value.clone(), true)
.expect("serializable body must construct");
assert_eq!(req.method, ProxyMethod::Post);
assert!(req.stream);
assert_eq!(req.body.as_ref(), Some(&value));
let stored = req.body_bytes.as_ref().expect("bytes must be set");
assert_eq!(stored, &digest_bytes, "caller digest bytes == wire bytes");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(stored).unwrap(),
value
);
req.validate().expect("coherent request must validate");
}
#[test]
fn mismatched_body_bytes_rejected_by_debug_validation() {
let req = ProxyRequest {
provider: "groq".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: Some(serde_json::json!({"model": "claimed"})),
stream: true,
body_bytes: Some(bytes::Bytes::from_static(b"{\"model\":\"actually-sent\"}")),
};
match req.validate() {
Err(BrokerError::Denied(msg)) => {
assert!(msg.contains("body_bytes"), "got: {msg}");
}
other => panic!("incoherent handoff must be denied, got {other:?}"),
}
}
#[test]
fn proxy_request_validation_fails_closed() {
let ok = ProxyRequest {
provider: "groq".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(ok.validate().is_ok());
let unknown = ProxyRequest {
provider: "evil".into(),
..ok.clone()
};
assert!(matches!(
unknown.validate(),
Err(BrokerError::UnknownProvider(_))
));
let absolute = ProxyRequest {
path: "https://evil.example/x".into(),
..ok.clone()
};
assert!(matches!(absolute.validate(), Err(BrokerError::Denied(_))));
let traversal = ProxyRequest {
path: "/../secrets".into(),
..ok.clone()
};
assert!(matches!(traversal.validate(), Err(BrokerError::Denied(_))));
let relative = ProxyRequest {
path: "chat/completions".into(),
..ok
};
assert!(matches!(relative.validate(), Err(BrokerError::Denied(_))));
}
#[test]
fn proxy_allows_only_pinned_kimi_code_paths() {
let chat = ProxyRequest {
provider: "kimi-code".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: true,
body_bytes: None,
};
assert!(chat.validate().is_ok(), "kimi-code chat must be allowed");
for path in ["/models", "/me", "/usages"] {
let get = ProxyRequest {
provider: "kimi-code".into(),
method: ProxyMethod::Get,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(get.validate().is_ok(), "kimi-code GET {path} must pass");
let post = ProxyRequest {
method: ProxyMethod::Post,
..get
};
assert!(
matches!(post.validate(), Err(BrokerError::Denied(_))),
"kimi-code POST {path} must be denied"
);
}
for path in ["/feedback", "/search", "/fetch", "/completions"] {
let denied = ProxyRequest {
provider: "kimi-code".into(),
method: ProxyMethod::Post,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(denied.validate(), Err(BrokerError::Denied(_))),
"kimi-code {path} must be denied"
);
}
let static_kimi = ProxyRequest {
provider: "kimi".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(static_kimi.validate().is_ok());
}
#[test]
fn proxy_rejects_oauth_providers_except_codex_catalog() {
let claude = ProxyRequest {
provider: "claude".into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(claude.validate(), Err(BrokerError::UnknownProvider(_))),
"claude alias must not be proxyable"
);
let anth = ProxyRequest {
provider: "anthropic".into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(anth.validate(), Err(BrokerError::Denied(_))),
"anthropic /models must be denied"
);
let bare = ProxyRequest {
provider: "openai-codex".into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(bare.validate(), Err(BrokerError::Denied(_))),
"openai-codex /models without client_version must be denied"
);
}
#[test]
fn proxy_allows_only_pinned_openai_codex_models_path() {
let allowed = ProxyRequest {
provider: "openai-codex".into(),
method: ProxyMethod::Get,
path: "/codex/models?client_version=0.153.3".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
allowed.validate().is_ok(),
"exact codex models path must be allowed"
);
for path in [
"/codex/models",
"/codex/models?client_version=",
"/codex/models?foo=1",
"/codex/models?client_version=0.153.3&extra=1",
"/codex/responses",
"/models",
"/models?client_version=0.153.3",
"/backend-api/codex/models?client_version=0.153.3",
"/v1/models?client_version=0.153.3",
"/chat/completions",
] {
let req = ProxyRequest {
provider: "openai-codex".into(),
method: ProxyMethod::Get,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(req.validate(), Err(BrokerError::Denied(_))),
"openai-codex path {path} must be denied"
);
}
let post = ProxyRequest {
provider: "openai-codex".into(),
method: ProxyMethod::Post,
path: "/codex/models?client_version=0.153.3".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(post.validate(), Err(BrokerError::Denied(_))),
"POST openai-codex models must be denied"
);
}
#[test]
fn proxy_allows_only_pinned_github_copilot_paths() {
let models = ProxyRequest {
provider: "github-copilot".into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(models.validate().is_ok());
let chat = ProxyRequest {
provider: "github-copilot".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(chat.validate().is_ok());
let responses = ProxyRequest {
provider: "github-copilot".into(),
method: ProxyMethod::Post,
path: "/responses".into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(responses.validate().is_ok());
for path in ["/v1/messages", "/models?x=1", "/embeddings"] {
let request = ProxyRequest {
provider: "github-copilot".into(),
method: ProxyMethod::Post,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(matches!(request.validate(), Err(BrokerError::Denied(_))));
}
}
#[test]
fn proxy_allows_only_pinned_anthropic_models_path() {
for path in [
"/v1/models",
"/v1/models?limit=100",
"/v1/models?limit=100&after_id=claude-opus-4-7",
] {
let req = ProxyRequest {
provider: "anthropic".into(),
method: ProxyMethod::Get,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(req.validate().is_ok(), "{path} must be allowed");
}
for path in [
"/v1/messages",
"/models",
"/v1/models?foo=1",
"/v1/models?limit=",
"/api/oauth/usage",
"/v1/models?limit=100&evil=1",
"/v1/models?limit=100&limit=50",
"/v1/models?after_id=a&after_id=b",
"/v1/models?limit=100&after_id=a&limit=50",
] {
let req = ProxyRequest {
provider: "anthropic".into(),
method: ProxyMethod::Get,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(req.validate(), Err(BrokerError::Denied(_))),
"anthropic path {path} must be denied"
);
}
}
#[test]
fn anthropic_oauth_catalog_headers_are_bearer_version_beta_not_x_api_key() {
let headers = anthropic_oauth_catalog_request_headers();
let names: Vec<_> = headers.iter().map(|(n, _)| *n).collect();
assert!(names.contains(&"anthropic-version"));
assert!(names.contains(&"anthropic-beta"));
assert!(names.contains(&"accept"));
assert!(!names.iter().any(|n| n.eq_ignore_ascii_case("x-api-key")));
assert!(!names
.iter()
.any(|n| n.eq_ignore_ascii_case("authorization")));
let beta = headers
.iter()
.find(|(n, _)| *n == "anthropic-beta")
.map(|(_, v)| *v)
.expect("beta");
assert_eq!(beta, "oauth-2025-04-20");
}
#[test]
fn anthropic_models_path_round_trip_and_rejects_duplicates() {
assert!(is_allowed_anthropic_path("/v1/models"));
assert!(is_allowed_anthropic_path("/v1/models?limit=100"));
assert!(is_allowed_anthropic_path(
"/v1/models?limit=100&after_id=claude-opus-4-7"
));
assert!(!is_allowed_anthropic_path("/v1/models?limit=100&limit=50"));
assert!(!is_allowed_anthropic_path(
"/v1/models?after_id=a&after_id=b"
));
let page0 = "/v1/models?limit=100";
let page1 = "/v1/models?limit=100&after_id=model-1";
assert!(is_allowed_anthropic_path(page0));
assert!(is_allowed_anthropic_path(page1));
}
#[test]
fn proxy_allows_only_pinned_google_gemini_paths() {
for path in [
"/v1internal:loadCodeAssist",
"/v1internal:onboardUser",
"/v1internal:streamGenerateContent",
"/v1internal:countTokens",
"/v1internal/operations/op-12345",
] {
let req = ProxyRequest {
provider: "google-gemini".into(),
method: ProxyMethod::Post,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(req.validate().is_ok(), "{path} must be allowed");
}
for path in [
"/v1internal:listExperiments",
"/v1internal:fetchAdminControls",
"/v1internal:setCodeAssistGlobalUserSetting",
"/v1internal:generateContent",
"/v2/models",
"/v1beta/models",
"/",
"/v1internal:",
] {
let req = ProxyRequest {
provider: "google-gemini".into(),
method: ProxyMethod::Post,
path: path.into(),
body: None,
stream: false,
body_bytes: None,
};
assert!(
matches!(req.validate(), Err(BrokerError::Denied(_))),
"{path} must be denied"
);
}
}
#[test]
fn access_token_type_has_no_refresh_field() {
let t: AccessToken =
serde_json::from_str(r#"{"token":"sk-x","expires":123,"refresh":"MUST-NOT-EXIST"}"#)
.unwrap();
let round = serde_json::to_value(&t).unwrap();
assert_eq!(round.get("refresh"), None);
assert_eq!(round["token"], "sk-x");
}
#[test]
fn broker_error_display_never_echoes_values() {
let e = BrokerError::NotConfigured("groq".into());
assert!(!format!("{e}").contains("sk-"));
let e = BrokerError::Unauthorized;
let shown = format!("{e}");
assert!(shown.starts_with("broker rejected machine auth"), "{shown}");
assert!(shown.contains("auth.machine_token") && shown.contains("SYNAPS_MACHINE_TOKEN"), "{shown}");
assert!(shown.contains("auth-broker.token") && shown.contains("restart"), "{shown}");
}
#[test]
fn mask_key_is_short_and_lossy() {
assert_eq!(mask_key("gsk-live-1234567890abcdef"), "gsk-…cdef");
assert_eq!(mask_key("short"), "…");
}
#[test]
fn proxy_request_serde_roundtrip() {
let req = ProxyRequest {
provider: "openrouter".into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: Some(serde_json::json!({"a": 1})),
stream: true,
body_bytes: None,
};
let json = serde_json::to_string(&req).unwrap();
let back: ProxyRequest = serde_json::from_str(&json).unwrap();
assert_eq!(back.provider, "openrouter");
assert!(back.stream);
assert!(matches!(back.method, ProxyMethod::Get));
}
async fn spawn_upstream(app: axum::Router) -> String {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{addr}")
}
#[tokio::test]
async fn local_broker_streams_local_endpoint_and_applies_key() {
use axum::routing::post;
let seen_auth = Arc::new(std::sync::Mutex::new(String::new()));
let seen = seen_auth.clone();
let app = axum::Router::new().route(
"/chat/completions",
post(move |headers: axum::http::HeaderMap| {
let seen = seen.clone();
async move {
*seen.lock().unwrap() = headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.unwrap_or("")
.to_string();
(
[("content-type", "text/event-stream")],
"data: {\"x\":1}\n\ndata: [DONE]\n\n",
)
}
}),
);
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
let mut stream = broker
.proxy_stream(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: Some(serde_json::json!({"model": "m"})),
stream: true,
body_bytes: None,
})
.await
.expect("stream must open");
use futures::StreamExt;
let mut collected = Vec::new();
while let Some(chunk) = stream.next().await {
collected.extend_from_slice(&chunk.unwrap());
}
let text = String::from_utf8(collected).unwrap();
assert!(text.contains("data: {\"x\":1}"));
assert!(text.contains("[DONE]"));
assert_eq!(&*seen_auth.lock().unwrap(), "Bearer local");
}
#[tokio::test]
async fn local_broker_stream_error_is_typed_and_keyless() {
use axum::routing::post;
let app = axum::Router::new().route(
"/chat/completions",
post(|| async {
(
axum::http::StatusCode::UNAUTHORIZED,
"{\"error\":\"bad key\"}",
)
}),
);
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
let err = match broker
.proxy_stream(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: true,
body_bytes: None,
})
.await
{
Err(e) => e,
Ok(_) => panic!("non-2xx upstream must not yield a stream"),
};
let msg = format!("{err}");
assert!(msg.contains("401"), "got: {msg}");
assert!(
!msg.to_lowercase().contains("bearer"),
"no auth material in errors"
);
}
#[tokio::test]
async fn local_broker_proxy_returns_status_and_body() {
use axum::routing::get;
let app =
axum::Router::new().route("/models", get(|| async { "{\"data\":[{\"id\":\"m1\"}]}" }));
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
let resp = broker
.proxy(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
})
.await
.unwrap();
assert_eq!(resp.status, 200);
assert!(resp.body.contains("m1"));
}
#[tokio::test]
async fn local_broker_proxy_rejects_oversize_response() {
use axum::routing::get;
let app = axum::Router::new().route("/models", get(|| async { "x".repeat(4096) }));
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
.with_max_response_bytes(64);
let err = broker
.proxy(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
})
.await
.expect_err("oversize body must be rejected");
let msg = format!("{err}");
assert!(msg.contains("limit"), "got: {msg}");
}
#[tokio::test]
async fn local_broker_stream_error_is_status_only() {
use axum::routing::post;
let app = axum::Router::new().route(
"/chat/completions",
post(|| async {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
format!("\x1b[2J{}", "E".repeat(64 * 1024)),
)
}),
);
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
let err = broker
.proxy_stream(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: true,
body_bytes: None,
})
.await
.err()
.expect("non-2xx upstream must not yield a stream");
let msg = format!("{err}");
assert!(
msg.len() < 700,
"error must be bounded, got {} bytes",
msg.len()
);
assert!(!msg.contains('\x1b'), "control chars must be stripped");
assert!(msg.contains("500"));
}
#[tokio::test]
async fn local_broker_buffered_request_times_out() {
use axum::routing::get;
let app = axum::Router::new().route(
"/models",
get(|| async {
tokio::time::sleep(Duration::from_secs(30)).await;
"too late"
}),
);
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
.with_request_timeout(Duration::from_millis(200));
let started = std::time::Instant::now();
let err = broker
.proxy(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Get,
path: "/models".into(),
body: None,
stream: false,
body_bytes: None,
})
.await
.expect_err("hung upstream must time out");
assert!(matches!(err, BrokerError::Transport(_)));
assert!(
started.elapsed() < Duration::from_secs(5),
"timeout must be enforced promptly"
);
}
#[tokio::test]
async fn global_broker_defaults_to_in_process_local() {
let broker = global_broker();
let caps = broker
.capabilities()
.await
.expect("in-process broker needs no daemon");
assert!(caps.iter().any(|c| c.key == LOCAL_PROVIDER_KEY));
let again = global_broker();
assert!(again.capabilities().await.is_ok());
}
const HOSTILE_SENTINEL: &str = "ZX9-HOSTILE-SENTINEL-7Q";
fn hostile_body() -> String {
format!(
"{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} \
{{\\\"messages\\\":[{{\\\"content\\\":\\\"my secret prompt\\\"}}],\
\\\"tools\\\":[{{\\\"input_schema\\\":{{}}}}]}}\"}}}}"
)
}
fn hostile_broker_json() -> String {
format!(
"{{\"error\":\"ECHOED {HOSTILE_SENTINEL} input_schema \
my secret prompt\"}}"
)
}
fn assert_no_hostile_bytes(msg: &str) {
assert!(
!msg.contains(HOSTILE_SENTINEL),
"sentinel leaked into error: {msg}"
);
assert!(!msg.contains("ECHOED"), "echoed body leaked: {msg}");
assert!(
!msg.contains("input_schema"),
"request-shaped data leaked: {msg}"
);
assert!(
!msg.contains("secret prompt"),
"prompt content leaked: {msg}"
);
}
#[tokio::test]
async fn local_broker_stream_error_omits_upstream_body() {
use axum::routing::post;
let app = axum::Router::new().route(
"/chat/completions",
post(|| async {
(
axum::http::StatusCode::TOO_MANY_REQUESTS,
format!("\x1b[2J{}", hostile_body()),
)
}),
);
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
let err = broker
.proxy_stream(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: true,
body_bytes: None,
})
.await
.err()
.expect("non-2xx upstream must not yield a stream");
assert!(matches!(err, BrokerError::Transport(_)));
let msg = format!("{err}");
assert!(
msg.contains("provider request failed: 429 Too Many Requests"),
"status prefix must survive for engine classification: {msg}"
);
assert!(
!msg.contains('\x1b'),
"escapes must not reach errors: {msg}"
);
assert_no_hostile_bytes(&msg);
}
async fn stream_error_message(status: axum::http::StatusCode, body: String) -> String {
use axum::routing::post;
let app = axum::Router::new().route(
"/chat/completions",
post(move || async move { (status, body) }),
);
let url = spawn_upstream(app).await;
let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
let err = broker
.proxy_stream(ProxyRequest {
provider: LOCAL_PROVIDER_KEY.into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream: true,
body_bytes: None,
})
.await
.err()
.expect("non-2xx upstream must not yield a stream");
assert!(matches!(err, BrokerError::Transport(_)));
format!("{err}")
}
fn after_marker(msg: &str) -> &str {
let marker = "provider request failed: ";
let idx = msg.find(marker).expect("marker present");
&msg[idx + marker.len()..]
}
#[tokio::test]
async fn local_broker_stream_403_kimi_quota_surfaces_vetted_label_only() {
let kimi_body = "{\"error\":{\"type\":\"access_terminated_error\",\"message\":\
\"You've reached your weekly (7-day) usage limit. Your quota will \
reset when the current 7-day window ends. \"}}";
let msg = stream_error_message(axum::http::StatusCode::FORBIDDEN, kimi_body.into()).await;
assert!(
msg.ends_with("provider request failed: 403 Forbidden [access_terminated_error]"),
"expected exact labelled prefix: {msg}"
);
assert!(!msg.contains("weekly"), "body text leaked: {msg}");
assert!(!msg.contains("quota"), "body text leaked: {msg}");
assert!(!msg.contains('{'), "raw JSON leaked: {msg}");
assert!(
!after_marker(&msg).contains(':'),
"suffix must not contain ':' (engine truncates there): {msg}"
);
}
#[tokio::test]
async fn local_broker_stream_429_unvetted_type_yields_no_label() {
let body = format!(
"{{\"error\":{{\"type\":\"<script>ECHO:secret\",\"message\":\"ECHOED {HOSTILE_SENTINEL}\"}}}}"
);
let msg = stream_error_message(axum::http::StatusCode::TOO_MANY_REQUESTS, body).await;
assert!(
msg.ends_with("provider request failed: 429 Too Many Requests"),
"no label expected: {msg}"
);
assert!(!msg.contains('['), "unexpected label: {msg}");
assert!(!msg.contains("script"), "hostile type leaked: {msg}");
assert!(!msg.contains("secret"), "hostile type leaked: {msg}");
assert!(!after_marker(&msg).contains(':'), "got: {msg}");
assert_no_hostile_bytes(&msg);
}
#[tokio::test]
async fn local_broker_stream_html_body_yields_no_label() {
let body =
format!("<html><body>ECHOED {HOSTILE_SENTINEL} access_terminated_error</body></html>");
let msg = stream_error_message(axum::http::StatusCode::BAD_GATEWAY, body).await;
assert!(
msg.ends_with("provider request failed: 502 Bad Gateway"),
"no label expected: {msg}"
);
assert!(!msg.contains("html"), "body leaked: {msg}");
assert!(!msg.contains('['), "unexpected label: {msg}");
assert_no_hostile_bytes(&msg);
}
#[tokio::test]
async fn local_broker_stream_oversized_body_yields_no_label_without_panic() {
let pad = "x".repeat(MAX_PROXY_ERROR_CLASSIFY_BYTES * 4);
let body = format!(
"{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} {pad}\",\"type\":\"access_terminated_error\"}}}}"
);
let msg = stream_error_message(axum::http::StatusCode::FORBIDDEN, body).await;
assert!(
msg.ends_with("provider request failed: 403 Forbidden"),
"no label expected for oversized body: {msg}"
);
assert!(!msg.contains('['), "unexpected label: {msg}");
assert!(
msg.len() < 200,
"oversized body leaked: {} bytes",
msg.len()
);
assert_no_hostile_bytes(&msg);
}
#[tokio::test]
async fn read_error_body_for_classification_truncates_and_never_fails() {
use axum::routing::get;
let app = axum::Router::new()
.route("/big", get(|| async { "a".repeat(100) }))
.route("/utf8", get(|| async { format!("{}é", "a".repeat(9)) }));
let url = spawn_upstream(app).await;
let client = reqwest::Client::new();
let resp = client.get(format!("{url}/big")).send().await.unwrap();
let got = read_error_body_for_classification(resp, 10).await;
assert_eq!(got.as_deref(), Some("aaaaaaaaaa"));
let resp = client.get(format!("{url}/utf8")).send().await.unwrap();
let got = read_error_body_for_classification(resp, 10).await;
assert_eq!(got, None);
let resp = client.get(format!("{url}/utf8")).send().await.unwrap();
let got = read_error_body_for_classification(resp, 1024).await;
assert_eq!(got.as_deref(), Some("aaaaaaaaaé"));
}
#[tokio::test]
#[serial_test::serial]
async fn local_broker_usage_error_omits_upstream_body() {
use axum::routing::get;
let app = axum::Router::new().route(
"/usage",
get(|| async {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
hostile_body(),
)
}),
);
let url = spawn_upstream(app).await;
let dir = tempfile::tempdir().unwrap();
let base = dir.path().join(".synaps-cli");
std::fs::create_dir_all(&base).unwrap();
let creds = crate::auth::OAuthCredentials {
auth_type: "oauth".into(),
refresh: String::new(),
access: "test-access-token".into(),
expires: crate::epoch_millis() + 3_600_000,
account_id: None,
};
super::super::storage::save_provider_auth_at_test_hook(
&base.join("auth.json"),
"anthropic",
&creds,
)
.unwrap();
let old_base = std::env::var("SYNAPS_BASE_DIR").ok();
let old_home = std::env::var("HOME").ok();
std::env::remove_var("SYNAPS_BASE_DIR");
std::env::set_var("HOME", dir.path());
let broker = LocalBroker::new(reqwest::Client::new())
.with_anthropic_usage_url(format!("{url}/usage"));
let result = broker.anthropic_usage().await;
match old_home {
Some(v) => std::env::set_var("HOME", v),
None => std::env::remove_var("HOME"),
}
if let Some(v) = old_base {
std::env::set_var("SYNAPS_BASE_DIR", v);
}
let err = result.expect_err("non-2xx usage response must fail");
assert!(matches!(err, BrokerError::Transport(_)));
let msg = format!("{err}");
assert!(
msg.contains("usage request failed") && msg.contains("500"),
"operation + status must survive: {msg}"
);
assert_no_hostile_bytes(&msg);
}
async fn spawn_remote_broker_error(status: u16) -> RemoteBroker {
use axum::routing::{get, post};
let code = axum::http::StatusCode::from_u16(status).unwrap();
let app = axum::Router::new()
.route(
"/proxy",
post(move || async move { (code, hostile_broker_json()) }),
)
.route(
"/usage",
get(move || async move { (code, hostile_broker_json()) }),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
RemoteBroker::new(
format!("http://{addr}"),
"opaque-machine-token",
reqwest::Client::new(),
super::super::TokenCache::new(),
)
}
fn groq_request(stream: bool) -> ProxyRequest {
ProxyRequest {
provider: "groq".into(),
method: ProxyMethod::Post,
path: "/chat/completions".into(),
body: None,
stream,
body_bytes: None,
}
}
#[tokio::test]
async fn remote_broker_denied_omits_broker_error_body() {
for status in [400u16, 403] {
let broker = spawn_remote_broker_error(status).await;
let err = broker
.proxy(groq_request(false))
.await
.expect_err("4xx broker response must fail");
assert!(
matches!(err, BrokerError::Denied(_)),
"typed category must survive, got {err:?}"
);
let msg = format!("{err}");
assert!(msg.contains(&status.to_string()), "status lost: {msg}");
assert_no_hostile_bytes(&msg);
}
}
#[test]
fn gemini_gets_the_longer_stream_idle_timeout() {
let broker = LocalBroker::new(reqwest::Client::new());
assert_eq!(broker.stream_idle_timeout_for("google-gemini"), Duration::from_secs(300));
for key in ["muse", "kimi-code", "openai-codex", "anthropic"] {
assert_eq!(broker.stream_idle_timeout_for(key), Duration::from_secs(90), "{key}");
}
let short = LocalBroker::new(reqwest::Client::new())
.with_stream_idle_timeout_for_tests(Duration::from_millis(10));
assert_eq!(short.stream_idle_timeout_for("muse"), Duration::from_millis(10));
assert!(short.stream_idle_timeout_for("google-gemini") >= Duration::from_secs(300));
}
async fn spawn_remote_broker_body(status: u16, body: String) -> RemoteBroker {
use axum::routing::post;
let code = axum::http::StatusCode::from_u16(status).unwrap();
let app = axum::Router::new().route(
"/proxy",
post(move || {
let body = body.clone();
async move { (code, body) }
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
RemoteBroker::new(
format!("http://{addr}"),
"opaque-machine-token",
reqwest::Client::new(),
super::super::TokenCache::new(),
)
}
#[tokio::test]
async fn remote_broker_502_keeps_a_retryable_refresh_failure() {
use super::super::token::{RefreshFailure, RefreshFailureKind as K};
for kind in [K::Transient, K::LockBusy, K::Other] {
let daemon_text = BrokerError::Refresh(RefreshFailure::new(
kind,
format!("HTTP error 503 {HOSTILE_SENTINEL} ECHOED input_schema"),
))
.to_string();
let body = serde_json::json!({ "error": daemon_text }).to_string();
let broker = spawn_remote_broker_body(502, body.clone()).await;
let stream_err = broker.proxy_stream(groq_request(true)).await.err().unwrap();
let buffered_err = broker.proxy(groq_request(false)).await.unwrap_err();
for err in [stream_err, buffered_err] {
let BrokerError::Refresh(f) = &err else {
panic!("{kind:?}: expected Refresh, got {err:?}");
};
assert_eq!(f.kind, kind);
let msg = err.to_string();
assert!(msg.contains("retry shortly"), "{msg}");
assert_no_hostile_bytes(&msg);
}
}
let broker = spawn_remote_broker_body(502, hostile_broker_json()).await;
let err = broker.proxy_stream(groq_request(true)).await.err().unwrap();
assert!(matches!(err, BrokerError::Transport(_)), "{err:?}");
assert_no_hostile_bytes(&err.to_string());
}
#[tokio::test]
async fn remote_broker_stream_error_omits_broker_body() {
let broker = spawn_remote_broker_error(500).await;
let err = broker
.proxy_stream(groq_request(true))
.await
.err()
.expect("non-2xx broker response must not yield a stream");
assert!(matches!(err, BrokerError::Transport(_)));
let msg = format!("{err}");
assert!(
msg.contains("broker proxy stream failed") && msg.contains("500"),
"operation + status must survive: {msg}"
);
assert_no_hostile_bytes(&msg);
}
#[tokio::test]
async fn remote_broker_usage_error_omits_broker_body() {
let broker = spawn_remote_broker_error(502).await;
let err = broker
.anthropic_usage()
.await
.expect_err("non-2xx broker response must fail");
assert!(matches!(err, BrokerError::Transport(_)));
let msg = format!("{err}");
assert!(
msg.contains("broker usage returned HTTP 502"),
"operation + status must survive: {msg}"
);
assert_no_hostile_bytes(&msg);
}
#[tokio::test]
async fn capabilities_expose_no_secret_fields() {
let broker = LocalBroker::new(reqwest::Client::new());
let caps = broker.capabilities().await.unwrap();
assert!(caps
.iter()
.any(|c| c.key == "anthropic" && c.kind == CredentialKind::OAuth));
assert!(caps
.iter()
.any(|c| c.key == "groq" && c.kind == CredentialKind::StaticKey));
let json = serde_json::to_string(&caps).unwrap();
for field in ["key\":", "name\":", "kind\":", "configured\":"] {
assert!(json.contains(field));
}
assert!(!json.contains("refresh"));
assert!(!json.contains("access"));
}
#[test]
fn auto_vend_failure_names_the_chosen_account_and_its_login_command() {
use super::super::token::{RefreshFailure, RefreshFailureKind as R};
let named = CredentialRef::new(
OAuthProviderId::Anthropic,
Account::named("claude3").unwrap(),
);
let refused =
auto_vend_failure(&named, RefreshFailure::new(R::Credential, "invalid_grant"));
assert_eq!(refused.kind, R::Credential);
assert_eq!(
refused.message(),
"account 'claude3': invalid_grant; log it in again with \
`synaps login --provider anthropic --account claude3`"
);
let default = CredentialRef::new(OAuthProviderId::Anthropic, Account::Default);
let refused = auto_vend_failure(&default, RefreshFailure::new(R::Credential, "revoked"));
assert_eq!(
refused.message(),
"account 'default': revoked; log it in again with \
`synaps login --provider anthropic`"
);
let advised = auto_vend_failure(
&named,
RefreshFailure::new(
R::Credential,
"slot removed; run `synaps login --account claude3`",
),
);
assert_eq!(advised.message().matches("synaps login").count(), 1);
for kind in [R::Transient, R::LockBusy, R::Other, R::ActionRequired] {
let f = auto_vend_failure(&named, RefreshFailure::new(kind, "HTTP 503"));
assert_eq!(f.kind, kind);
assert_eq!(f.message(), "account 'claude3': HTTP 503");
}
}
#[test]
fn refresh_failure_kinds_map_to_observation_kinds() {
use super::super::quota_policy::{AccountFailureKind as K, QuotaObservation};
use super::super::token::{RefreshFailure, RefreshFailureKind as R};
for (kind, expected, transient, login) in [
(R::Transient, K::RefreshTransient, true, false),
(R::LockBusy, K::RefreshLockBusy, true, false),
(R::Other, K::RefreshOther, true, false),
(
R::Cooldown { retry_in_secs: 42 },
K::RefreshOther,
true,
false,
),
(R::Credential, K::RefreshRejected, false, true),
(R::ActionRequired, K::RefreshActionRequired, false, false),
] {
let read = UsageReadError::refresh(RefreshFailure::new(kind, "why"));
assert_eq!(read.failure.kind, expected, "{kind:?}");
assert_eq!(read.failure.kind.is_transient(), transient, "{kind:?}");
assert_eq!(read.failure.kind.needs_login(), login, "{kind:?}");
assert_eq!(
read.error,
BrokerError::Credential("why".into()),
"{kind:?}"
);
let observation = read.failure.observation();
if login {
assert_eq!(observation, QuotaObservation::AuthError);
} else {
assert_eq!(observation, QuotaObservation::Unknown);
}
}
let cooldown =
refresh_failure(&RefreshFailure::new(R::Cooldown { retry_in_secs: 42 }, "x"));
assert_eq!(cooldown.retry_after_s, Some(42));
let gone = UsageReadError::slot(BrokerError::UnknownAccount {
provider: "anthropic".into(),
label: "x".into(),
});
assert_eq!(gone.failure.kind, K::CredentialMissing);
assert!(gone.failure.kind.needs_login());
let store = UsageReadError::slot(BrokerError::Credential("unreadable".into()));
assert_eq!(store.failure.kind, K::CredentialStore);
assert!(!store.failure.kind.needs_login() && !store.failure.kind.is_transient());
}
#[test]
fn usage_errors_map_to_observation_kinds_with_status() {
use super::super::quota_policy::AccountFailureKind as K;
use super::super::usage::UsageError as U;
for (err, kind, status) in [
(
U::RateLimited {
retry_after_s: Some(9),
},
K::UsageRateLimited,
Some(429),
),
(
U::Unauthorized { status: 403 },
K::UsageUnauthorized,
Some(403),
),
(
U::UpstreamStatus {
status: 503,
retry_after_s: None,
},
K::UsageUpstream,
Some(503),
),
(U::Redirected { status: 302 }, K::UsageUpstream, Some(302)),
(U::Timeout, K::UsageTimeout, None),
(
U::Transport {
detail: "connect".into(),
},
K::UsageTransport,
None,
),
(U::Malformed { detail: "x".into() }, K::UsageMalformed, None),
(U::BodyTooLarge { cap: 1 }, K::UsageMalformed, None),
(U::AccountIdUnavailable, K::UsageUnauthorized, None),
(
U::UnsupportedProvider {
provider: "p".into(),
},
K::Unsupported,
None,
),
] {
let read = UsageReadError::usage(err.clone());
assert_eq!(read.failure.kind, kind, "{err:?}");
assert_eq!(read.failure.http_status, status, "{err:?}");
assert_eq!(read.error, err.clone().into_broker_error(), "{err:?}");
}
assert_eq!(
UsageReadError::usage(U::RateLimited {
retry_after_s: Some(9)
})
.failure
.retry_after_s,
Some(9)
);
}
}