1use std::collections::BTreeMap;
23use std::pin::Pin;
24use std::sync::Arc;
25use std::time::Duration;
26
27use async_trait::async_trait;
28use futures::Stream;
29use serde::{Deserialize, Serialize};
30
31use super::cloud::{CloudProviderId, InvokeRequest};
32use super::provider::OAuthProviderId;
33use super::static_providers::{
34 allowed_proxy_paths, static_provider, StaticProviderSpec, LOCAL_DEFAULT_BASE_URL,
35 LOCAL_PROVIDER_KEY, STATIC_PROVIDERS,
36};
37use super::{load_provider_auth, storage};
38
39pub const MAX_PROXY_REQUEST_BYTES: usize = 2 * 1024 * 1024;
48pub const MAX_PROXY_RESPONSE_BYTES: usize = 8 * 1024 * 1024;
50pub const MAX_UPSTREAM_ERROR_BYTES: usize = 2 * 1024;
54pub const PROXY_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
56
57const ANTHROPIC_USAGE_URL: &str = "https://api.anthropic.com/api/oauth/usage";
60const ANTHROPIC_OAUTH_BETA: &str = "oauth-2025-04-20";
61
62pub const GOOGLE_GEMINI_CODE_ASSIST_BASE_URL: &str = "https://cloudcode-pa.googleapis.com";
66
67pub(crate) fn is_allowed_google_gemini_path(path: &str) -> bool {
71 matches!(
72 path,
73 "/v1internal:loadCodeAssist"
74 | "/v1internal:onboardUser"
75 | "/v1internal:streamGenerateContent"
76 | "/v1internal:countTokens"
77 ) || path.starts_with("/v1internal/operations/")
78}
79
80pub const OPENAI_CODEX_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
82
83pub(crate) const OPENAI_CODEX_MODELS_PATH: &str = "/codex/models";
92
93pub(crate) fn is_allowed_openai_codex_path(path: &str) -> bool {
94 let Some((base, query)) = path.split_once('?') else {
95 return false;
96 };
97 if base != OPENAI_CODEX_MODELS_PATH {
98 return false;
99 }
100 let mut client_version: Option<&str> = None;
101 for pair in query.split('&') {
102 let Some((k, v)) = pair.split_once('=') else {
103 return false;
104 };
105 if k != "client_version" || client_version.is_some() {
106 return false;
107 }
108 if v.trim().is_empty() {
109 return false;
110 }
111 client_version = Some(v);
112 }
113 client_version.is_some()
114}
115
116pub(crate) fn is_allowed_anthropic_path(path: &str) -> bool {
121 if path == "/v1/models" {
122 return true;
123 }
124 let Some((base, query)) = path.split_once('?') else {
125 return false;
126 };
127 if base != "/v1/models" {
128 return false;
129 }
130 if query.is_empty() {
131 return false;
132 }
133 let mut saw_limit = false;
134 let mut saw_after_id = false;
135 for pair in query.split('&') {
136 let Some((k, v)) = pair.split_once('=') else {
137 return false;
138 };
139 if v.is_empty() {
140 return false;
141 }
142 match k {
143 "limit" if !saw_limit => saw_limit = true,
144 "after_id" if !saw_after_id => saw_after_id = true,
145 _ => return false,
146 }
147 }
148 true
151}
152
153pub(crate) fn anthropic_oauth_catalog_request_headers() -> &'static [(&'static str, &'static str)] {
157 &[
158 ("anthropic-version", "2023-06-01"),
159 ("anthropic-beta", ANTHROPIC_OAUTH_BETA),
160 ("accept", "application/json"),
161 ]
162}
163
164#[derive(Debug, Clone, PartialEq, Eq)]
168pub enum BrokerError {
169 UnknownProvider(String),
171 RegistrationRequired {
173 provider: String,
174 remediation: String,
175 },
176 NotConfigured(String),
178 Unauthorized,
180 Denied(String),
182 UnsupportedCapability {
185 provider: String,
186 capability: String,
187 },
188 Transport(String),
190 Credential(String),
192}
193
194impl std::fmt::Display for BrokerError {
195 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
196 match self {
197 Self::UnknownProvider(p) => write!(f, "unknown provider: {p}"),
198 Self::RegistrationRequired {
199 provider,
200 remediation,
201 } => {
202 write!(f, "registration required for '{provider}': {remediation}")
203 }
204 Self::NotConfigured(p) => write!(
205 f,
206 "no credential configured for '{p}'. Run `synaps login` to add one."
207 ),
208 Self::Unauthorized => write!(f, "broker rejected machine auth"),
209 Self::Denied(msg) => write!(f, "broker denied request: {msg}"),
210 Self::UnsupportedCapability {
211 provider,
212 capability,
213 } => write!(
214 f,
215 "cloud provider '{provider}' is text-only: {capability} are not supported yet \
216 (no credential was used and no network request was made)"
217 ),
218 Self::Transport(msg) => write!(f, "broker transport error: {msg}"),
219 Self::Credential(msg) => write!(f, "credential error: {msg}"),
220 }
221 }
222}
223
224impl std::error::Error for BrokerError {}
225
226#[derive(Debug, Clone, Serialize, Deserialize)]
233pub struct AccessToken {
234 pub token: String,
235 pub expires: u64,
237}
238
239#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
241#[serde(rename_all = "lowercase")]
242pub enum ProxyMethod {
243 Get,
244 Post,
245}
246
247#[derive(Debug, Clone, Serialize, Deserialize)]
255pub struct ProxyRequest {
256 pub provider: String,
258 pub method: ProxyMethod,
259 pub path: String,
261 #[serde(default, skip_serializing_if = "Option::is_none")]
263 pub body: Option<serde_json::Value>,
264 #[serde(default)]
266 pub stream: bool,
267 #[serde(skip)]
276 pub body_bytes: Option<bytes::Bytes>,
277}
278
279impl ProxyRequest {
280 pub fn validate(&self) -> Result<(), BrokerError> {
283 if self.provider != LOCAL_PROVIDER_KEY
284 && self.provider != "xai-auth"
285 && self.provider != "github-copilot"
286 && self.provider != "google-gemini"
287 && self.provider != "openai-codex"
288 && self.provider != "anthropic"
289 && self.provider != "kimi-code"
290 && static_provider(&self.provider).is_none()
291 {
292 return Err(BrokerError::UnknownProvider(self.provider.clone()));
293 }
294 if !self.path.starts_with('/') {
295 return Err(BrokerError::Denied("proxy path must be relative".into()));
296 }
297 if self.path.contains("://") || self.path.contains("..") {
298 return Err(BrokerError::Denied(
299 "proxy path is not a plain relative path".into(),
300 ));
301 }
302 let oauth_path_allowed = self.provider == "xai-auth" && self.path == "/responses"
306 || self.provider == "github-copilot"
307 && matches!(
308 self.path.as_str(),
309 "/models" | "/chat/completions" | "/responses"
310 )
311 || self.provider == "google-gemini" && is_allowed_google_gemini_path(&self.path)
312 || self.provider == "kimi-code"
315 && (self.path == "/chat/completions"
316 || self.method == ProxyMethod::Get
317 && matches!(self.path.as_str(), "/models" | "/me" | "/usages"))
318 || self.provider == "openai-codex"
319 && self.method == ProxyMethod::Get
320 && is_allowed_openai_codex_path(&self.path)
321 || self.provider == "anthropic"
322 && self.method == ProxyMethod::Get
323 && is_allowed_anthropic_path(&self.path);
324 if !oauth_path_allowed && !allowed_proxy_paths(&self.provider).contains(&self.path.as_str())
325 {
326 return Err(BrokerError::Denied(format!(
327 "proxy path '{}' is not in the provider's endpoint allowlist",
328 self.path
329 )));
330 }
331 let size = match (&self.body_bytes, &self.body) {
336 (Some(bytes), _) => Some(bytes.len()),
338 (None, Some(body)) => Some(
339 serde_json::to_vec(body)
340 .map(|v| v.len())
341 .unwrap_or(usize::MAX),
342 ),
343 (None, None) => None,
344 };
345 if let Some(size) = size {
346 if size > MAX_PROXY_REQUEST_BYTES {
347 return Err(BrokerError::Denied(format!(
348 "request body exceeds the {MAX_PROXY_REQUEST_BYTES}-byte broker limit"
349 )));
350 }
351 }
352 #[cfg(any(test, debug_assertions))]
358 if let (Some(bytes), Some(body)) = (&self.body_bytes, &self.body) {
359 let parsed = serde_json::from_slice::<serde_json::Value>(bytes).ok();
360 if parsed.as_ref() != Some(body) {
361 return Err(BrokerError::Denied(
362 "body_bytes does not match the JSON body (exact-byte handoff incoherent)"
363 .into(),
364 ));
365 }
366 }
367 Ok(())
368 }
369
370 pub fn post_json_exact(
379 provider: impl Into<String>,
380 path: impl Into<String>,
381 body: serde_json::Value,
382 stream: bool,
383 ) -> Result<(Self, bytes::Bytes), BrokerError> {
384 let bytes =
385 bytes::Bytes::from(serde_json::to_vec(&body).map_err(|e| {
386 BrokerError::Denied(format!("request body failed to serialize: {e}"))
387 })?);
388 let request = Self {
389 provider: provider.into(),
390 method: ProxyMethod::Post,
391 path: path.into(),
392 body: Some(body),
393 stream,
394 body_bytes: Some(bytes.clone()),
395 };
396 Ok((request, bytes))
397 }
398}
399
400#[derive(Debug, Clone, Serialize, Deserialize)]
402pub struct ProxyResponse {
403 pub status: u16,
404 pub body: String,
405}
406
407pub type ProxyByteStream = Pin<Box<dyn Stream<Item = Result<bytes::Bytes, BrokerError>> + Send>>;
409
410#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
412#[serde(rename_all = "snake_case")]
413pub enum CredentialKind {
414 OAuth,
415 StaticKey,
416 LocalEndpoint,
417}
418
419#[derive(Debug, Clone, Serialize, Deserialize)]
421pub struct ProviderStatus {
422 pub key: String,
423 pub name: String,
424 pub kind: CredentialKind,
425 pub configured: bool,
427}
428
429#[derive(Debug, Clone, PartialEq, Eq)]
431pub enum StaticKeyStatus {
432 NotSet,
433 Configured {
436 masked: String,
437 },
438 FromEnv,
440}
441
442impl StaticKeyStatus {
443 pub fn is_configured(&self) -> bool {
444 !matches!(self, StaticKeyStatus::NotSet)
445 }
446}
447
448#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
452pub struct CloudCatalogEntry {
453 pub provider: CloudProviderId,
454 pub id: String,
455 pub display_name: String,
456 pub context_ref: String,
458 pub context_label: String,
459 pub stale: bool,
460 pub fetched_at: u64,
462}
463
464#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
466#[serde(tag = "type", rename_all = "snake_case")]
467pub enum CloudEvent {
468 TextDelta {
469 delta: String,
470 },
471 ToolArguments {
472 id: String,
473 name: Option<String>,
474 delta: String,
475 },
476 Usage {
477 input_tokens: u64,
478 output_tokens: u64,
479 },
480 Done,
481}
482pub type CloudEventStream = Pin<Box<dyn Stream<Item = Result<CloudEvent, BrokerError>> + Send>>;
483
484#[async_trait]
488pub trait CredentialBroker: Send + Sync {
489 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError>;
491
492 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError>;
494
495 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError>;
499
500 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError>;
507
508 async fn cloud_catalog(
510 &self,
511 provider: CloudProviderId,
512 context_ref: &str,
513 allow_stale: bool,
514 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
515 let _ = (context_ref, allow_stale);
516 Err(BrokerError::NotConfigured(provider.to_string()))
517 }
518
519 async fn cloud_invoke(
522 &self,
523 provider: CloudProviderId,
524 context_ref: &str,
525 model_id: &str,
526 request: InvokeRequest,
527 ) -> Result<CloudEventStream, BrokerError> {
528 let _ = (context_ref, model_id, request);
529 Err(BrokerError::NotConfigured(provider.to_string()))
530 }
531
532 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError>;
534}
535
536#[async_trait]
539pub trait CloudBackend: Send + Sync {
540 async fn catalog(
541 &self,
542 provider: CloudProviderId,
543 context_ref: &str,
544 allow_stale: bool,
545 ) -> Result<Vec<CloudCatalogEntry>, BrokerError>;
546 async fn invoke(
547 &self,
548 provider: CloudProviderId,
549 context_ref: &str,
550 model_id: &str,
551 request: InvokeRequest,
552 ) -> Result<CloudEventStream, BrokerError>;
553}
554
555#[derive(Clone)]
556struct ProductionCloudBackend {
557 http: reqwest::Client,
560 refresh_lock: Arc<tokio::sync::Mutex<()>>,
561}
562impl ProductionCloudBackend {
563 fn new(_http: reqwest::Client) -> Self {
564 let http = reqwest::Client::builder()
565 .redirect(reqwest::redirect::Policy::none())
566 .connect_timeout(Duration::from_secs(10))
567 .timeout(Duration::from_secs(90))
568 .build()
569 .expect("cloud HTTPS client configuration is valid");
570 Self {
571 http,
572 refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
573 }
574 }
575
576 fn state(&self, provider: CloudProviderId) -> Result<serde_json::Value, BrokerError> {
577 storage::load_cloud_state(provider.as_str())
578 .map_err(BrokerError::Credential)?
579 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))
580 }
581
582 fn context_ref(&self, provider: CloudProviderId) -> Result<String, BrokerError> {
583 use sha2::{Digest, Sha256};
584 let state = self.state(provider)?;
585 let public = serde_json::to_vec(&state["config"])
586 .map_err(|_| BrokerError::Credential("invalid cloud context".into()))?;
587 let digest = Sha256::digest([provider.as_str().as_bytes(), &public].concat());
588 Ok(format!(
589 "ctx-{}",
590 digest[..16]
591 .iter()
592 .map(|b| format!("{b:02x}"))
593 .collect::<String>()
594 ))
595 }
596
597 fn validate_context(
598 &self,
599 provider: CloudProviderId,
600 supplied: &str,
601 ) -> Result<String, BrokerError> {
602 let opaque = self.context_ref(provider)?;
603 if supplied == provider.as_str() || supplied == opaque {
606 Ok(opaque)
607 } else {
608 Err(BrokerError::Denied(
609 "cloud context does not match stored provider".into(),
610 ))
611 }
612 }
613
614 async fn aws(
615 &self,
616 ) -> Result<super::aws_bedrock::AwsBedrockBroker<super::aws_bedrock::AwsHttpApi>, BrokerError>
617 {
618 #[derive(Deserialize)]
619 struct State {
620 config: super::cloud::AwsBedrockConfig,
621 access_key: String,
622 secret_key: String,
623 session_token: String,
624 expires_at: u64,
625 registered_client: serde_json::Value,
626 sso_access_token: String,
627 sso_refresh_token: Option<String>,
628 sso_expires_at: u64,
629 }
630 let decode = |value| {
631 serde_json::from_value::<State>(value)
632 .map_err(|_| BrokerError::Credential("invalid AWS broker state".into()))
633 };
634 let mut raw = self.state(CloudProviderId::AwsBedrock)?;
635 let now = std::time::SystemTime::now()
636 .duration_since(std::time::UNIX_EPOCH)
637 .unwrap_or_default()
638 .as_millis() as u64;
639 let mut s = decode(raw.clone())?;
640 if s.expires_at <= now + 60_000 {
641 let _guard = self.refresh_lock.lock().await;
644 raw = self.state(CloudProviderId::AwsBedrock)?;
645 s = decode(raw.clone())?;
646 if s.expires_at <= now + 60_000 {
647 let api =
648 super::aws_bedrock::AwsHttpApi::new(self.http.clone(), &s.config.sso_region);
649 if s.sso_expires_at <= now + 60_000 {
650 let refresh = s
651 .sso_refresh_token
652 .as_deref()
653 .filter(|v| !v.is_empty())
654 .ok_or_else(|| {
655 BrokerError::Credential("AWS SSO session expired; login again".into())
656 })?;
657 let client = super::aws_bedrock::RegisteredClient::new(
658 s.registered_client["id"].as_str().ok_or_else(|| {
659 BrokerError::Credential("invalid AWS client registration".into())
660 })?,
661 s.registered_client["secret"].as_str().ok_or_else(|| {
662 BrokerError::Credential("invalid AWS client registration".into())
663 })?,
664 s.registered_client["expires_at"].as_u64().ok_or_else(|| {
665 BrokerError::Credential("invalid AWS client registration".into())
666 })?,
667 );
668 if client.expires_at <= now / 1000 {
669 return Err(BrokerError::Credential(
670 "AWS client registration expired; login again".into(),
671 ));
672 }
673 use super::aws_bedrock::AwsApi;
674 let token = api
675 .create_token(
676 &client,
677 &s.config.sso_region,
678 super::aws_bedrock::TokenGrant::RefreshToken(refresh),
679 )
680 .await
681 .map_err(|_| {
682 BrokerError::Credential("AWS SSO refresh rejected; login again".into())
683 })?;
684 raw["sso_access_token"] = token.access().into();
685 if let Some(rotated) = token.refresh() {
686 raw["sso_refresh_token"] = rotated.into();
687 }
688 raw["sso_expires_at"] = (now + token.expires_in * 1000).into();
689 storage::save_cloud_state("aws-bedrock", &raw)
694 .map_err(BrokerError::Credential)?;
695 s = decode(raw.clone())?;
696 }
697 use super::aws_bedrock::AwsApi;
698 let credentials = api
699 .get_role_credentials(
700 &s.config.sso_region,
701 &s.sso_access_token,
702 &s.config.account_id,
703 &s.config.role_name,
704 )
705 .await
706 .map_err(|_| {
707 BrokerError::Credential("AWS role refresh rejected; login again".into())
708 })?;
709 raw["access_key"] = credentials.access_key().into();
710 raw["secret_key"] = credentials.secret_key().into();
711 raw["session_token"] = credentials.session_token().into();
712 raw["expires_at"] = credentials.expires_at.into();
713 storage::save_cloud_state("aws-bedrock", &raw).map_err(BrokerError::Credential)?;
714 s = decode(raw)?;
715 }
716 }
717 let api = super::aws_bedrock::AwsHttpApi::new(self.http.clone(), &s.config.sso_region);
718 Ok(super::aws_bedrock::AwsBedrockBroker::from_credentials(
719 api,
720 s.config,
721 super::aws_bedrock::RoleCredentials::new(
722 s.access_key,
723 s.secret_key,
724 s.session_token,
725 s.expires_at,
726 ),
727 ))
728 }
729 async fn azure_request(
730 &self,
731 state: &mut serde_json::Value,
732 audience: super::azure_openai::AzureAudience,
733 ) -> Result<String, BrokerError> {
734 let now = std::time::SystemTime::now()
735 .duration_since(std::time::UNIX_EPOCH)
736 .unwrap_or_default()
737 .as_millis() as u64;
738 let key = match audience {
739 super::azure_openai::AzureAudience::Arm => "arm",
740 super::azure_openai::AzureAudience::Inference => "inference",
741 };
742 if state[key]["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
743 return state[key]["access_token"]
744 .as_str()
745 .map(str::to_owned)
746 .ok_or_else(|| BrokerError::Credential("invalid Azure token state".into()));
747 }
748 let _refresh = self.refresh_lock.lock().await;
751 *state = self.state(CloudProviderId::AzureOpenAi)?;
752 if state[key]["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
753 return state[key]["access_token"]
754 .as_str()
755 .map(str::to_owned)
756 .ok_or_else(|| BrokerError::Credential("invalid Azure token state".into()));
757 }
758 let config: super::cloud::AzureOpenAiConfig =
759 serde_json::from_value(state["config"].clone())
760 .map_err(|_| BrokerError::Credential("invalid Azure broker state".into()))?;
761 let client_id = state["client_id"]
762 .as_str()
763 .ok_or_else(|| BrokerError::Credential("invalid Azure client registration".into()))?;
764 let refresh = state["refresh_token"]
765 .as_str()
766 .ok_or_else(|| BrokerError::Credential("invalid Azure refresh state".into()))?;
767 let reg = super::azure_openai::AzureRegistration::production(Some(client_id.into()))
768 .map_err(|e| BrokerError::Credential(e.to_string()))?;
769 let r = super::azure_openai::refresh_request(&config, ®, audience, refresh)
770 .map_err(|e| BrokerError::Credential(e.to_string()))?;
771 let response = self
772 .http
773 .post(r.url)
774 .form(&r.form)
775 .send()
776 .await
777 .map_err(|_| BrokerError::Transport("Azure refresh failed".into()))?;
778 if !response.status().is_success() {
779 return Err(BrokerError::Credential(
780 "Azure refresh rejected; login again".into(),
781 ));
782 }
783 let wire: serde_json::Value = response
784 .json()
785 .await
786 .map_err(|_| BrokerError::Transport("invalid Azure token response".into()))?;
787 let access = wire["access_token"]
788 .as_str()
789 .filter(|v| !v.is_empty())
790 .ok_or_else(|| BrokerError::Transport("invalid Azure token response".into()))?
791 .to_owned();
792 if let Some(r) = wire["refresh_token"].as_str().filter(|v| !v.is_empty()) {
793 state["refresh_token"] = r.into();
794 }
795 state[key] = serde_json::json!({"access_token": access, "expires_at": now + wire["expires_in"].as_u64().unwrap_or(3600) * 1000});
796 storage::save_cloud_state("azure-openai", state).map_err(BrokerError::Credential)?;
797 Ok(access)
798 }
799
800 async fn vertex_request(&self, state: &mut serde_json::Value) -> Result<String, BrokerError> {
801 let now = std::time::SystemTime::now()
802 .duration_since(std::time::UNIX_EPOCH)
803 .unwrap_or_default()
804 .as_millis() as u64;
805 if state["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
806 return state["access_token"]
807 .as_str()
808 .map(str::to_owned)
809 .ok_or_else(|| BrokerError::Credential("invalid Vertex token state".into()));
810 }
811 let _refresh = self.refresh_lock.lock().await;
812 *state = self.state(CloudProviderId::GoogleVertex)?;
813 if state["expires_at"].as_u64().unwrap_or(0) > now + 60_000 {
814 return state["access_token"]
815 .as_str()
816 .map(str::to_owned)
817 .ok_or_else(|| BrokerError::Credential("invalid Vertex token state".into()));
818 }
819 let client_id = state["client_id"]
820 .as_str()
821 .ok_or_else(|| BrokerError::Credential("invalid Vertex registration".into()))?;
822 let refresh = state["refresh_token"]
823 .as_str()
824 .ok_or_else(|| BrokerError::Credential("invalid Vertex refresh state".into()))?;
825 let response = self
826 .http
827 .post(super::google_vertex::TOKEN_URL)
828 .form(&[
829 ("client_id", client_id),
830 ("grant_type", "refresh_token"),
831 ("refresh_token", refresh),
832 ])
833 .send()
834 .await
835 .map_err(|_| BrokerError::Transport("Vertex refresh failed".into()))?;
836 if !response.status().is_success() {
837 return Err(BrokerError::Credential(
838 "Vertex refresh rejected; login again".into(),
839 ));
840 }
841 let wire: serde_json::Value = response
842 .json()
843 .await
844 .map_err(|_| BrokerError::Transport("invalid Vertex token response".into()))?;
845 let access = wire["access_token"]
846 .as_str()
847 .filter(|v| !v.is_empty())
848 .ok_or_else(|| BrokerError::Transport("invalid Vertex token response".into()))?
849 .to_owned();
850 if let Some(r) = wire["refresh_token"].as_str().filter(|v| !v.is_empty()) {
851 state["refresh_token"] = r.into();
852 }
853 state["access_token"] = access.clone().into();
854 state["expires_at"] = (now + wire["expires_in"].as_u64().unwrap_or(3600) * 1000).into();
855 storage::save_cloud_state("google-vertex", state).map_err(BrokerError::Credential)?;
856 Ok(access)
857 }
858}
859
860#[async_trait]
861impl CloudBackend for ProductionCloudBackend {
862 async fn catalog(
863 &self,
864 provider: CloudProviderId,
865 context_ref: &str,
866 _allow_stale: bool,
867 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
868 let opaque_context = self.validate_context(provider, context_ref)?;
869 let fetched_at = std::time::SystemTime::now()
870 .duration_since(std::time::UNIX_EPOCH)
871 .unwrap_or_default()
872 .as_millis() as u64;
873 match provider {
874 CloudProviderId::AwsBedrock => {
875 let broker = self.aws().await?;
876 let label = broker.public_context().bedrock_region;
877 broker
878 .catalog()
879 .await
880 .map_err(|_| {
881 BrokerError::Transport("AWS catalog failed (details redacted)".into())
882 })
883 .map(|xs| {
884 xs.into_iter()
885 .map(|x| CloudCatalogEntry {
886 provider,
887 id: x.id,
888 display_name: x.display_name,
889 context_ref: opaque_context.clone(),
890 context_label: label.clone(),
891 stale: false,
892 fetched_at,
893 })
894 .collect()
895 })
896 }
897 CloudProviderId::AzureOpenAi => {
898 let mut state = self.state(provider)?;
899 let config: super::cloud::AzureOpenAiConfig =
900 serde_json::from_value(state["config"].clone()).map_err(|_| {
901 BrokerError::Credential("invalid Azure broker state".into())
902 })?;
903 let token = self
904 .azure_request(&mut state, super::azure_openai::AzureAudience::Arm)
905 .await?;
906 let mut discovery =
907 super::azure_openai::DeploymentDiscovery::new(config.clone(), 20, 1000);
908 let mut url = Some(discovery.initial_url());
909 while let Some(next) = url {
910 let r = self
911 .http
912 .get(next)
913 .bearer_auth(&token)
914 .send()
915 .await
916 .map_err(|_| BrokerError::Transport("Azure catalog failed".into()))?;
917 if !r.status().is_success() {
918 return Err(BrokerError::Transport("Azure catalog rejected".into()));
919 }
920 let body = read_body_capped(r, MAX_CLOUD_CATALOG_BODY_BYTES)
921 .await
922 .map_err(|_| {
923 BrokerError::Transport("Azure catalog exceeded body limit".into())
924 })?;
925 url = discovery
926 .accept_page(&body)
927 .map_err(|_| BrokerError::Transport("invalid Azure catalog".into()))?;
928 }
929 discovery
930 .finish()
931 .map_err(|_| BrokerError::Transport("Azure catalog unavailable".into()))
932 .map(|xs| {
933 xs.into_iter()
934 .map(|x| CloudCatalogEntry {
935 provider,
936 id: x.id,
937 display_name: x.display_name,
938 context_ref: opaque_context.clone(),
939 context_label: format!(
940 "{}/{}",
941 config.resource_name, config.resource_group
942 ),
943 stale: false,
944 fetched_at,
945 })
946 .collect()
947 })
948 }
949 CloudProviderId::GoogleVertex => {
950 let mut state = self.state(provider)?;
951 let token = self.vertex_request(&mut state).await?;
952 let config: super::cloud::GoogleVertexConfig =
953 serde_json::from_value(state["config"].clone()).map_err(|_| {
954 BrokerError::Credential("invalid Vertex broker state".into())
955 })?;
956 let mut page: Option<String> = None;
957 let mut out = Vec::new();
958 let mut seen = std::collections::HashSet::new();
959 for _ in 0..20 {
960 let mut url=format!("https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models",config.location,config.project_id,config.location);
961 if let Some(p) = &page {
962 let mut parsed = url::Url::parse(&url).map_err(|_| {
963 BrokerError::Transport("invalid Vertex catalog URL".into())
964 })?;
965 parsed.query_pairs_mut().append_pair("pageToken", p);
966 url = parsed.into();
967 }
968 let r = self
969 .http
970 .get(url)
971 .bearer_auth(&token)
972 .send()
973 .await
974 .map_err(|_| BrokerError::Transport("Vertex catalog failed".into()))?;
975 if !r.status().is_success() {
976 return Err(BrokerError::Transport("Vertex catalog rejected".into()));
977 }
978 let v = cloud_catalog_json(r).await?;
979 for m in v["publisherModels"].as_array().into_iter().flatten() {
980 if let Some(name) = m["name"]
981 .as_str()
982 .filter(|n| n.starts_with("publishers/google/models/"))
983 {
984 if out.len() >= MAX_CLOUD_CATALOG_ENTRIES {
985 return Err(BrokerError::Transport(
986 "Vertex catalog exceeded entry limit".into(),
987 ));
988 }
989 if !out.iter().any(|entry: &CloudCatalogEntry| {
990 entry.id == format!("google-vertex/{name}")
991 }) {
992 out.push(CloudCatalogEntry {
993 provider,
994 id: format!("google-vertex/{name}"),
995 display_name: m["displayName"].as_str().unwrap_or(name).into(),
996 context_ref: opaque_context.clone(),
997 context_label: format!(
998 "{}/{}",
999 config.project_id, config.location
1000 ),
1001 stale: false,
1002 fetched_at,
1003 });
1004 }
1005 }
1006 }
1007 page = v["nextPageToken"].as_str().map(str::to_owned);
1008 match &page {
1009 Some(p) if seen.insert(p.clone()) => {}
1010 Some(_) => {
1011 return Err(BrokerError::Transport("Vertex pagination loop".into()))
1012 }
1013 None => break,
1014 }
1015 }
1016 if out.is_empty() {
1017 Err(BrokerError::Transport("Vertex catalog is empty".into()))
1018 } else {
1019 Ok(out)
1020 }
1021 }
1022 }
1023 }
1024 async fn invoke(
1025 &self,
1026 provider: CloudProviderId,
1027 context_ref: &str,
1028 model_id: &str,
1029 request: InvokeRequest,
1030 ) -> Result<CloudEventStream, BrokerError> {
1031 self.validate_context(provider, context_ref)?;
1032 if !request.tools.is_empty() {
1033 return Err(BrokerError::Denied(
1034 "tools are not yet supported by cloud providers".into(),
1035 ));
1036 }
1037 let authorized = self.catalog(provider, context_ref, false).await?;
1038 if !authorized.iter().any(|entry| entry.id == model_id) {
1039 return Err(BrokerError::Denied(
1040 "model is not present in the current provider catalog".into(),
1041 ));
1042 }
1043 match provider {
1044 CloudProviderId::AwsBedrock => {
1045 let broker = self.aws().await?;
1046 if request.stream {
1047 use futures::StreamExt;
1048 let stream = broker
1049 .converse_stream(model_id, request)
1050 .await
1051 .map_err(|_| {
1052 BrokerError::Transport(
1053 "AWS invocation failed (details redacted)".into(),
1054 )
1055 })?
1056 .map(|event| {
1057 event
1058 .map_err(|_| {
1059 BrokerError::Transport(
1060 "AWS stream failed (details redacted)".into(),
1061 )
1062 })
1063 .map(|e| match e {
1064 super::aws_bedrock::ConverseEvent::TextDelta(delta) => {
1065 CloudEvent::TextDelta { delta }
1066 }
1067 super::aws_bedrock::ConverseEvent::ToolArguments {
1068 id,
1069 delta,
1070 } => CloudEvent::ToolArguments {
1071 id,
1072 name: None,
1073 delta,
1074 },
1075 super::aws_bedrock::ConverseEvent::Usage(u) => {
1076 CloudEvent::Usage {
1077 input_tokens: u.input_tokens,
1078 output_tokens: u.output_tokens,
1079 }
1080 }
1081 super::aws_bedrock::ConverseEvent::Done => CloudEvent::Done,
1082 })
1083 });
1084 Ok(Box::pin(stream))
1085 } else {
1086 let o = broker.converse(model_id, request).await.map_err(|_| {
1087 BrokerError::Transport("AWS invocation failed (details redacted)".into())
1088 })?;
1089 Ok(Box::pin(futures::stream::iter(vec![
1090 Ok(CloudEvent::TextDelta { delta: o.text }),
1091 Ok(CloudEvent::Usage {
1092 input_tokens: o.usage.input_tokens,
1093 output_tokens: o.usage.output_tokens,
1094 }),
1095 Ok(CloudEvent::Done),
1096 ])))
1097 }
1098 }
1099 CloudProviderId::AzureOpenAi => {
1100 let mut state = self.state(provider)?;
1101 let token = self
1102 .azure_request(&mut state, super::azure_openai::AzureAudience::Inference)
1103 .await?;
1104 let endpoint = super::azure_openai::AzureEndpoint::parse(
1105 state["endpoint"]
1106 .as_str()
1107 .ok_or_else(|| BrokerError::Credential("invalid Azure endpoint".into()))?,
1108 )
1109 .map_err(|_| BrokerError::Credential("invalid Azure endpoint".into()))?;
1110 let deployment = model_id
1111 .strip_prefix("azure-openai/")
1112 .ok_or_else(|| BrokerError::Denied("invalid Azure model".into()))?;
1113 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});
1114 let rr = super::azure_openai::responses_request(
1115 &endpoint,
1116 deployment,
1117 &serde_json::to_vec(&body).unwrap(),
1118 )
1119 .map_err(|_| BrokerError::Denied("invalid Azure request".into()))?;
1120 let r = self
1121 .http
1122 .post(rr.url)
1123 .bearer_auth(token)
1124 .header("content-type", "application/json")
1125 .body(rr.body)
1126 .send()
1127 .await
1128 .map_err(|_| BrokerError::Transport("Azure invocation failed".into()))?;
1129 if !r.status().is_success() {
1130 return Err(BrokerError::Transport("Azure invocation rejected".into()));
1131 }
1132 if request.stream {
1133 use futures::StreamExt;
1134 use std::sync::{
1135 atomic::{AtomicBool, Ordering},
1136 Arc,
1137 };
1138 let terminal = Arc::new(AtomicBool::new(false));
1139 let seen = terminal.clone();
1140 let events = sse_json_stream(r).filter_map(move |item| {
1141 let seen = seen.clone();
1142 async move {
1143 match item {
1144 Ok(Some(v)) => v["delta"]
1145 .as_str()
1146 .or_else(|| v["text"].as_str())
1147 .map(|delta| {
1148 Ok(CloudEvent::TextDelta {
1149 delta: delta.into(),
1150 })
1151 }),
1152 Ok(None) => {
1153 seen.store(true, Ordering::SeqCst);
1154 Some(Ok(CloudEvent::Done))
1155 }
1156 Err(error) => Some(Err(error)),
1157 }
1158 }
1159 });
1160 let stream = events.chain(futures::stream::unfold(
1161 (terminal, false),
1162 |(terminal, emitted)| async move {
1163 if terminal.load(Ordering::SeqCst) || emitted {
1164 None
1165 } else {
1166 Some((
1167 Err(BrokerError::Transport(
1168 "Azure stream ended without [DONE]".into(),
1169 )),
1170 (terminal, true),
1171 ))
1172 }
1173 },
1174 ));
1175 Ok(Box::pin(stream))
1176 } else {
1177 let text = read_body_capped(r, MAX_CLOUD_STREAM_EVENT_BYTES).await?;
1178 let v: serde_json::Value = serde_json::from_str(&text)
1179 .map_err(|_| BrokerError::Transport("invalid Azure response".into()))?;
1180 let delta = v["output_text"]
1181 .as_str()
1182 .ok_or_else(|| {
1183 BrokerError::Transport("Azure response omitted output".into())
1184 })?
1185 .to_owned();
1186 Ok(Box::pin(futures::stream::iter(vec![
1187 Ok(CloudEvent::TextDelta { delta }),
1188 Ok(CloudEvent::Done),
1189 ])))
1190 }
1191 }
1192 CloudProviderId::GoogleVertex => {
1193 let mut state = self.state(provider)?;
1194 let token = self.vertex_request(&mut state).await?;
1195 let config: super::cloud::GoogleVertexConfig =
1196 serde_json::from_value(state["config"].clone())
1197 .map_err(|_| BrokerError::Credential("invalid Vertex state".into()))?;
1198 let model = model_id
1199 .strip_prefix("google-vertex/publishers/google/models/")
1200 .filter(|m| {
1201 m.bytes()
1202 .all(|b| b.is_ascii_alphanumeric() || b"-._".contains(&b))
1203 })
1204 .ok_or_else(|| BrokerError::Denied("invalid Vertex model".into()))?;
1205 let action = if request.stream {
1206 "streamGenerateContent?alt=sse"
1207 } else {
1208 "generateContent"
1209 };
1210 let url=format!("https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:{}",config.location,config.project_id,config.location,model,action);
1211 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<_>>()});
1212 let r = self
1213 .http
1214 .post(url)
1215 .bearer_auth(token)
1216 .json(&body)
1217 .send()
1218 .await
1219 .map_err(|_| BrokerError::Transport("Vertex invocation failed".into()))?;
1220 if !r.status().is_success() {
1221 return Err(BrokerError::Transport("Vertex invocation rejected".into()));
1222 }
1223 if request.stream {
1224 use futures::StreamExt;
1225 let events = sse_json_stream(r).map(|item| {
1226 item.and_then(|value| {
1227 value.ok_or_else(|| {
1228 BrokerError::Transport(
1229 "Vertex stream terminated without a provider event".into(),
1230 )
1231 })
1232 })
1233 });
1234 use std::sync::{
1235 atomic::{AtomicBool, Ordering},
1236 Arc,
1237 };
1238 let terminal = Arc::new(AtomicBool::new(false));
1239 let terminal_in_events = terminal.clone();
1240 let stream = events
1241 .flat_map(move |item| {
1242 let mut out = Vec::new();
1243 match item {
1244 Ok(v) => {
1245 if terminal_in_events.load(Ordering::SeqCst) {
1246 out.push(Err(BrokerError::Transport(
1247 "Vertex emitted data after terminal metadata".into(),
1248 )));
1249 return futures::stream::iter(out);
1250 }
1251 for c in v["candidates"].as_array().into_iter().flatten() {
1252 for p in
1253 c["content"]["parts"].as_array().into_iter().flatten()
1254 {
1255 if let Some(delta) = p["text"].as_str() {
1256 out.push(Ok(CloudEvent::TextDelta {
1257 delta: delta.into(),
1258 }));
1259 }
1260 }
1261 if let Some(reason) = c["finishReason"].as_str() {
1262 const VALID: &[&str] = &[
1263 "STOP",
1264 "MAX_TOKENS",
1265 "SAFETY",
1266 "RECITATION",
1267 "OTHER",
1268 "BLOCKLIST",
1269 "PROHIBITED_CONTENT",
1270 "SPII",
1271 "MALFORMED_FUNCTION_CALL",
1272 ];
1273 if !VALID.contains(&reason) {
1274 out.push(Err(BrokerError::Transport(
1275 "Vertex returned an invalid finish reason"
1276 .into(),
1277 )));
1278 } else if terminal_in_events
1279 .swap(true, Ordering::SeqCst)
1280 {
1281 out.push(Err(BrokerError::Transport(
1282 "Vertex returned duplicate terminal metadata"
1283 .into(),
1284 )));
1285 } else {
1286 out.push(Ok(CloudEvent::Done));
1287 }
1288 }
1289 }
1290 if let Some(u) = v.get("usageMetadata") {
1291 out.push(Ok(CloudEvent::Usage {
1292 input_tokens: u["promptTokenCount"]
1293 .as_u64()
1294 .unwrap_or(0),
1295 output_tokens: u["candidatesTokenCount"]
1296 .as_u64()
1297 .unwrap_or(0),
1298 }));
1299 }
1300 }
1301 Err(error) => out.push(Err(error)),
1302 }
1303 futures::stream::iter(out)
1304 })
1305 .chain(futures::stream::unfold(
1306 (terminal, false),
1307 |(terminal, emitted)| async move {
1308 if terminal.load(Ordering::SeqCst) || emitted {
1309 None
1310 } else {
1311 Some((
1312 Err(BrokerError::Transport(
1313 "Vertex stream ended without finish metadata".into(),
1314 )),
1315 (terminal, true),
1316 ))
1317 }
1318 },
1319 ));
1320 Ok(Box::pin(stream))
1321 } else {
1322 let text = read_body_capped(r, MAX_CLOUD_STREAM_EVENT_BYTES).await?;
1323 let v: serde_json::Value = serde_json::from_str(&text)
1324 .map_err(|_| BrokerError::Transport("invalid Vertex response".into()))?;
1325 let mut out = Vec::new();
1326 for c in v["candidates"].as_array().into_iter().flatten() {
1327 for p in c["content"]["parts"].as_array().into_iter().flatten() {
1328 if let Some(delta) = p["text"].as_str() {
1329 out.push(Ok(CloudEvent::TextDelta {
1330 delta: delta.into(),
1331 }));
1332 }
1333 }
1334 }
1335 out.push(Ok(CloudEvent::Done));
1336 Ok(Box::pin(futures::stream::iter(out)))
1337 }
1338 }
1339 }
1340 }
1341}
1342
1343const MAX_CLOUD_CATALOG_BODY_BYTES: usize = 2 * 1024 * 1024;
1344const MAX_CLOUD_CATALOG_ENTRIES: usize = 10_000;
1345const MAX_CLOUD_STREAM_EVENT_BYTES: usize = 1024 * 1024;
1346
1347fn sse_json_stream(
1351 resp: reqwest::Response,
1352) -> Pin<Box<dyn Stream<Item = Result<Option<serde_json::Value>, BrokerError>> + Send>> {
1353 use futures::StreamExt;
1354 let chunks = Box::pin(resp.bytes_stream());
1355 Box::pin(futures::stream::unfold(
1356 (chunks, Vec::<u8>::new(), false),
1357 |(mut chunks, mut buf, done)| async move {
1358 if done {
1359 return None;
1360 }
1361 loop {
1362 if let Some(end) = buf.windows(2).position(|w| w == b"\n\n") {
1363 let frame: Vec<u8> = buf.drain(..end + 2).collect();
1364 let text = match std::str::from_utf8(&frame[..end]) {
1365 Ok(v) => v,
1366 Err(_) => {
1367 return Some((
1368 Err(BrokerError::Transport("invalid SSE encoding".into())),
1369 (chunks, buf, true),
1370 ))
1371 }
1372 };
1373 let data = text
1374 .lines()
1375 .filter_map(|line| line.strip_prefix("data:"))
1376 .map(str::trim_start)
1377 .collect::<Vec<_>>()
1378 .join("\n");
1379 if data.is_empty() {
1380 continue;
1381 }
1382 if data == "[DONE]" {
1383 return Some((Ok(None), (chunks, buf, true)));
1384 }
1385 let value = serde_json::from_str(&data)
1386 .map_err(|_| BrokerError::Transport("invalid SSE event".into()));
1387 let terminal = value.is_err();
1388 return Some((value.map(Some), (chunks, buf, terminal)));
1389 }
1390 match chunks.next().await {
1391 Some(Ok(chunk)) if buf.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES => {
1392 buf.extend_from_slice(&chunk)
1393 }
1394 Some(_) => {
1395 return Some((
1396 Err(BrokerError::Transport(
1397 "cloud stream event exceeded limit".into(),
1398 )),
1399 (chunks, buf, true),
1400 ))
1401 }
1402 None if buf.is_empty() => return None,
1403 None => {
1404 return Some((
1405 Err(BrokerError::Transport("truncated SSE event".into())),
1406 (chunks, buf, true),
1407 ))
1408 }
1409 }
1410 }
1411 },
1412 ))
1413}
1414
1415async fn cloud_catalog_json(resp: reqwest::Response) -> Result<serde_json::Value, BrokerError> {
1417 let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
1418 serde_json::from_str(&body)
1419 .map_err(|_| BrokerError::Transport("invalid cloud catalog response".into()))
1420}
1421
1422pub struct LocalBroker {
1426 http: reqwest::Client,
1427 local_base_url: Option<String>,
1429 anthropic_usage_url: Option<String>,
1431 google_gemini_base_url: Option<String>,
1433 request_timeout: Duration,
1435 max_response_bytes: usize,
1437 cloud_backend: Option<Arc<dyn CloudBackend>>,
1438}
1439
1440impl LocalBroker {
1441 pub fn new(http: reqwest::Client) -> Self {
1442 let cloud_backend = Arc::new(ProductionCloudBackend::new(http.clone()));
1443 Self {
1444 http,
1445 local_base_url: None,
1446 anthropic_usage_url: None,
1447 google_gemini_base_url: None,
1448 request_timeout: PROXY_REQUEST_TIMEOUT,
1449 max_response_bytes: MAX_PROXY_RESPONSE_BYTES,
1450 cloud_backend: Some(cloud_backend),
1451 }
1452 }
1453
1454 pub fn with_cloud_backend(mut self, backend: Arc<dyn CloudBackend>) -> Self {
1455 self.cloud_backend = Some(backend);
1456 self
1457 }
1458
1459 pub fn with_local_base_url(http: reqwest::Client, base_url: impl Into<String>) -> Self {
1461 Self {
1462 local_base_url: Some(base_url.into()),
1463 ..Self::new(http)
1464 }
1465 }
1466
1467 #[doc(hidden)]
1471 pub fn with_google_gemini_base_url_for_tests(
1472 http: reqwest::Client,
1473 base_url: impl Into<String>,
1474 ) -> Self {
1475 Self {
1476 google_gemini_base_url: Some(base_url.into().trim_end_matches('/').to_string()),
1477 ..Self::new(http)
1478 }
1479 }
1480
1481 #[doc(hidden)]
1483 pub fn with_anthropic_usage_url(mut self, url: impl Into<String>) -> Self {
1484 self.anthropic_usage_url = Some(url.into());
1485 self
1486 }
1487
1488 #[doc(hidden)]
1490 pub fn with_request_timeout(mut self, timeout: Duration) -> Self {
1491 self.request_timeout = timeout;
1492 self
1493 }
1494
1495 #[doc(hidden)]
1497 pub fn with_max_response_bytes(mut self, cap: usize) -> Self {
1498 self.max_response_bytes = cap;
1499 self
1500 }
1501
1502 fn resolve_static_key(&self, provider: &str) -> Result<String, BrokerError> {
1506 if provider == LOCAL_PROVIDER_KEY {
1507 return Ok(self.resolve_local_key());
1508 }
1509 let spec = static_provider(provider)
1510 .ok_or_else(|| BrokerError::UnknownProvider(provider.to_string()))?;
1511 if let Ok(Some(key)) = storage::load_static_key(provider) {
1512 return Ok(key);
1513 }
1514 if let Some(key) = discover_legacy_static_key(spec) {
1515 if let Err(e) = storage::save_static_key(provider, &key) {
1518 tracing::warn!(provider, "static key migration failed: {e}");
1519 }
1520 return Ok(key);
1521 }
1522 Err(BrokerError::NotConfigured(provider.to_string()))
1523 }
1524
1525 fn resolve_local_key(&self) -> String {
1526 storage::load_static_key(LOCAL_PROVIDER_KEY)
1527 .ok()
1528 .flatten()
1529 .or_else(|| {
1530 crate::config::get_provider_keys()
1531 .get(LOCAL_PROVIDER_KEY)
1532 .filter(|s| !s.is_empty())
1533 .cloned()
1534 })
1535 .or_else(|| {
1536 std::env::var("LOCAL_API_KEY")
1537 .ok()
1538 .filter(|s| !s.is_empty())
1539 })
1540 .unwrap_or_else(|| "local".to_string())
1541 }
1542
1543 fn base_url_for(&self, provider: &str) -> Result<String, BrokerError> {
1544 if provider == LOCAL_PROVIDER_KEY {
1545 if let Some(url) = &self.local_base_url {
1546 return Ok(url.trim_end_matches('/').to_string());
1547 }
1548 return Ok(local_endpoint_url());
1549 }
1550 static_provider(provider)
1551 .map(|s| s.base_url.trim_end_matches('/').to_string())
1552 .ok_or_else(|| BrokerError::UnknownProvider(provider.to_string()))
1553 }
1554
1555 async fn send(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
1556 request.validate()?;
1557 let (key, base) = if request.provider == "xai-auth" {
1558 let token = self.access_token(OAuthProviderId::Xai).await?;
1559 (token.token, "https://api.x.ai/v1".to_string())
1560 } else if request.provider == "github-copilot" {
1561 let token = self.access_token(OAuthProviderId::GitHubCopilot).await?;
1564 (
1565 token.token,
1566 super::github_copilot_models_base_url().to_string(),
1567 )
1568 } else if request.provider == "google-gemini" {
1569 let token = self.access_token(OAuthProviderId::GoogleGemini).await?;
1572 let base = self
1573 .google_gemini_base_url
1574 .clone()
1575 .unwrap_or_else(|| GOOGLE_GEMINI_CODE_ASSIST_BASE_URL.to_string());
1576 (token.token, base)
1577 } else if request.provider == "openai-codex" {
1578 let token = self.access_token(OAuthProviderId::OpenAiCodex).await?;
1581 (token.token, OPENAI_CODEX_BACKEND_BASE_URL.to_string())
1582 } else if request.provider == "kimi-code" {
1583 let token = self.access_token(OAuthProviderId::KimiCode).await?;
1587 (token.token, super::kimi_code::API_BASE_URL.to_string())
1588 } else if request.provider == "anthropic" {
1589 let token = self.access_token(OAuthProviderId::Anthropic).await?;
1592 (token.token, "https://api.anthropic.com".to_string())
1593 } else {
1594 (
1595 self.resolve_static_key(&request.provider)?,
1596 self.base_url_for(&request.provider)?,
1597 )
1598 };
1599 let url = format!("{base}{}", request.path);
1600 let mut builder = if request.provider == "google-gemini"
1603 || request.provider == "openai-codex"
1604 || request.provider == "anthropic"
1605 || request.provider == "kimi-code"
1606 {
1607 let no_redirect = reqwest::Client::builder()
1608 .redirect(reqwest::redirect::Policy::none())
1609 .connect_timeout(Duration::from_secs(10))
1610 .build()
1611 .map_err(|e| BrokerError::Transport(format!("client build failed: {e}")))?;
1612 match request.method {
1613 ProxyMethod::Get => no_redirect.get(&url),
1614 ProxyMethod::Post => no_redirect.post(&url),
1615 }
1616 } else {
1617 match request.method {
1618 ProxyMethod::Get => self.http.get(&url),
1619 ProxyMethod::Post => self.http.post(&url),
1620 }
1621 };
1622 builder = builder.bearer_auth(&key);
1623 if request.provider == "openai-codex" {
1624 let account_id = super::load_provider_auth("openai-codex")
1627 .ok()
1628 .flatten()
1629 .and_then(|c| c.account_id)
1630 .filter(|s| !s.trim().is_empty())
1631 .or_else(|| super::extract_codex_account_id(&key))
1632 .ok_or_else(|| {
1633 BrokerError::Credential(
1634 "openai-codex credential is missing chatgpt account id".into(),
1635 )
1636 })?;
1637 builder = builder
1638 .header("chatgpt-account-id", account_id)
1639 .header("originator", "synaps")
1640 .header("OpenAI-Beta", "responses=experimental")
1641 .header("accept", "application/json");
1642 }
1643 if request.provider == "anthropic" {
1644 for (name, value) in anthropic_oauth_catalog_request_headers() {
1647 builder = builder.header(*name, *value);
1648 }
1649 }
1650 if request.provider == "github-copilot" {
1651 for (name, value) in super::github_copilot_models_request_headers() {
1652 builder = builder.header(*name, *value);
1653 }
1654 if request.path != "/models" {
1655 builder = builder
1656 .header("Openai-Intent", "conversation-edits")
1657 .header("X-Initiator", "agent");
1658 }
1659 }
1660 if request.provider == "google-gemini" {
1661 if request.stream && request.path == "/v1internal:streamGenerateContent" {
1665 builder = builder.query(&[("alt", "sse")]);
1666 }
1667 builder = builder.header("user-agent", "SynapsCLI/0.6.0 (google-gemini)");
1668 }
1669 if request.provider == "kimi-code" {
1670 let device_id = super::kimi_code::load_or_create_device_id();
1674 for (name, value) in super::kimi_code::identity_request_headers(&device_id) {
1675 builder = builder.header(name, value);
1676 }
1677 }
1678 if let Some(bytes) = &request.body_bytes {
1679 builder = builder
1683 .header("content-type", "application/json")
1684 .body(bytes.clone());
1685 } else if let Some(body) = &request.body {
1686 builder = builder.json(body);
1687 }
1688 if request.stream {
1689 builder = builder.header("accept", "text/event-stream");
1690 } else {
1691 builder = builder.timeout(self.request_timeout);
1695 }
1696 builder.send().await.map_err(|e| {
1697 if e.is_connect() && request.provider == LOCAL_PROVIDER_KEY {
1698 BrokerError::Transport(format!(
1699 "can't reach local endpoint at {url} — is Ollama/LM Studio running?"
1700 ))
1701 } else {
1702 BrokerError::Transport(format!("request to {} failed: {e}", request.provider))
1704 }
1705 })
1706 }
1707}
1708
1709async fn read_body_capped(resp: reqwest::Response, cap: usize) -> Result<String, BrokerError> {
1714 use futures::StreamExt;
1715 let mut buf: Vec<u8> = Vec::new();
1716 let mut stream = resp.bytes_stream();
1717 while let Some(chunk) = stream.next().await {
1718 let chunk =
1719 chunk.map_err(|e| BrokerError::Transport(format!("failed to read response: {e}")))?;
1720 if buf.len() + chunk.len() > cap {
1721 return Err(BrokerError::Transport(format!(
1722 "response body exceeded the {cap}-byte broker buffering limit"
1723 )));
1724 }
1725 buf.extend_from_slice(&chunk);
1726 }
1727 String::from_utf8(buf)
1728 .map_err(|_| BrokerError::Transport("response body was not valid UTF-8".into()))
1729}
1730
1731const MAX_PROXY_ERROR_CLASSIFY_BYTES: usize = 8 * 1024;
1735
1736async fn read_error_body_for_classification(resp: reqwest::Response, cap: usize) -> Option<String> {
1741 use futures::StreamExt;
1742 let mut buf: Vec<u8> = Vec::with_capacity(cap.min(1024));
1743 let mut stream = resp.bytes_stream();
1744 while let Some(chunk) = stream.next().await {
1745 let chunk = chunk.ok()?;
1746 let room = cap.saturating_sub(buf.len());
1747 if chunk.len() >= room {
1748 buf.extend_from_slice(&chunk[..room]);
1749 return String::from_utf8(buf).ok();
1752 }
1753 buf.extend_from_slice(&chunk);
1754 }
1755 String::from_utf8(buf).ok()
1756}
1757
1758fn discover_legacy_static_key(spec: &StaticProviderSpec) -> Option<String> {
1761 if let Some(v) = crate::config::get_provider_keys().get(spec.key) {
1762 if !v.is_empty() {
1763 return Some(v.clone());
1764 }
1765 }
1766 spec.env_vars
1767 .iter()
1768 .find_map(|var| std::env::var(var).ok().filter(|v| !v.is_empty()))
1769}
1770
1771pub fn static_key_configured(provider: &str) -> bool {
1774 if provider == LOCAL_PROVIDER_KEY {
1775 return true; }
1777 let Some(spec) = static_provider(provider) else {
1778 return false;
1779 };
1780 storage::load_static_key(provider).ok().flatten().is_some()
1781 || discover_legacy_static_key(spec).is_some()
1782}
1783
1784pub fn static_key_status(provider: &str) -> StaticKeyStatus {
1788 let stored = storage::load_static_key(provider)
1789 .ok()
1790 .flatten()
1791 .or_else(|| {
1792 crate::config::get_provider_keys()
1793 .get(provider)
1794 .filter(|s| !s.is_empty())
1795 .cloned()
1796 });
1797 if let Some(key) = stored {
1798 return StaticKeyStatus::Configured {
1799 masked: mask_key(&key),
1800 };
1801 }
1802 if let Some(spec) = static_provider(provider) {
1803 if spec
1804 .env_vars
1805 .iter()
1806 .any(|var| std::env::var(var).is_ok_and(|v| !v.is_empty()))
1807 {
1808 return StaticKeyStatus::FromEnv;
1809 }
1810 }
1811 StaticKeyStatus::NotSet
1812}
1813
1814fn mask_key(key: &str) -> String {
1816 let chars: Vec<char> = key.chars().collect();
1817 if chars.len() <= 8 {
1818 return "…".to_string();
1819 }
1820 let head: String = chars[..4].iter().collect();
1821 let tail: String = chars[chars.len() - 4..].iter().collect();
1822 format!("{head}…{tail}")
1823}
1824
1825pub fn configured_static_provider_keys() -> std::collections::BTreeSet<String> {
1827 STATIC_PROVIDERS
1828 .iter()
1829 .filter(|s| static_key_configured(s.key))
1830 .map(|s| s.key.to_string())
1831 .collect()
1832}
1833
1834pub fn local_endpoint_url() -> String {
1837 local_endpoint_config()
1838 .unwrap_or_else(|| LOCAL_DEFAULT_BASE_URL.to_string())
1839 .trim_end_matches('/')
1840 .to_string()
1841}
1842
1843pub fn local_endpoint_config() -> Option<String> {
1845 crate::config::get_provider_keys()
1846 .get("local.url")
1847 .filter(|s| !s.is_empty())
1848 .cloned()
1849 .or_else(|| {
1850 std::env::var("LOCAL_ENDPOINT")
1851 .ok()
1852 .filter(|s| !s.is_empty())
1853 })
1854}
1855
1856pub fn local_model_ids() -> Vec<String> {
1858 crate::config::get_provider_keys()
1859 .get("local.models")
1860 .map(|value| {
1861 value
1862 .split(',')
1863 .map(|s| s.trim().to_string())
1864 .filter(|s| !s.is_empty())
1865 .collect()
1866 })
1867 .unwrap_or_default()
1868}
1869
1870pub fn anthropic_credential_available() -> bool {
1873 let oauth = load_provider_auth(OAuthProviderId::Anthropic.as_str())
1874 .ok()
1875 .flatten()
1876 .map(|c| c.auth_type == "oauth" && !c.access.is_empty())
1877 .unwrap_or(false);
1878 oauth || std::env::var("ANTHROPIC_API_KEY").is_ok_and(|v| !v.is_empty())
1879}
1880
1881pub fn oauth_provider_logged_in(provider: OAuthProviderId) -> bool {
1885 load_provider_auth(provider.as_str())
1886 .ok()
1887 .flatten()
1888 .is_some_and(|creds| {
1889 creds.auth_type == "oauth" && (!creds.refresh.is_empty() || !creds.access.is_empty())
1890 })
1891}
1892
1893#[async_trait]
1894impl CredentialBroker for LocalBroker {
1895 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
1896 let creds = super::ensure_fresh_provider_token(&self.http, provider)
1897 .await
1898 .map_err(BrokerError::Credential)?;
1899 Ok(AccessToken {
1901 token: creds.access,
1902 expires: creds.expires,
1903 })
1904 }
1905
1906 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
1907 let resp = self.send(&request).await?;
1908 let status = resp.status().as_u16();
1909 let body = read_body_capped(resp, self.max_response_bytes).await?;
1910 Ok(ProxyResponse { status, body })
1911 }
1912
1913 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
1914 let mut request = request;
1915 request.stream = true;
1916 let resp = self.send(&request).await?;
1917 let status = resp.status();
1918 if !status.is_success() {
1919 let label = read_error_body_for_classification(resp, MAX_PROXY_ERROR_CLASSIFY_BYTES)
1929 .await
1930 .as_deref()
1931 .and_then(crate::error::vetted_proxy_error_label);
1932 tracing::warn!(
1933 provider = %request.provider,
1934 status = status.as_u16(),
1935 label = label.unwrap_or("-"),
1936 "provider request failed"
1937 );
1938 let label_suffix = label.map(|l| format!(" [{l}]")).unwrap_or_default();
1939 return Err(BrokerError::Transport(format!(
1940 "provider request failed: {status}{label_suffix}"
1941 )));
1942 }
1943 use futures::StreamExt;
1944 let stream = resp
1945 .bytes_stream()
1946 .map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
1947 Ok(Box::pin(stream))
1948 }
1949
1950 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
1951 let token = self.access_token(OAuthProviderId::Anthropic).await?;
1954 let url = self
1955 .anthropic_usage_url
1956 .as_deref()
1957 .unwrap_or(ANTHROPIC_USAGE_URL);
1958 let resp = self
1959 .http
1960 .get(url)
1961 .timeout(self.request_timeout)
1962 .bearer_auth(&token.token)
1963 .header("anthropic-beta", ANTHROPIC_OAUTH_BETA)
1964 .send()
1965 .await
1966 .map_err(|e| BrokerError::Transport(format!("usage request failed: {e}")))?;
1967 let status = resp.status();
1968 if !status.is_success() {
1969 drop(resp);
1972 return Err(BrokerError::Transport(format!(
1973 "usage request failed: {status}"
1974 )));
1975 }
1976 let body = read_body_capped(resp, self.max_response_bytes).await?;
1977 serde_json::from_str(&body)
1978 .map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
1979 }
1980
1981 async fn cloud_catalog(
1982 &self,
1983 provider: CloudProviderId,
1984 context_ref: &str,
1985 allow_stale: bool,
1986 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
1987 self.cloud_backend
1988 .as_ref()
1989 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
1990 .catalog(provider, context_ref, allow_stale)
1991 .await
1992 }
1993
1994 async fn cloud_invoke(
1995 &self,
1996 provider: CloudProviderId,
1997 context_ref: &str,
1998 model_id: &str,
1999 request: InvokeRequest,
2000 ) -> Result<CloudEventStream, BrokerError> {
2001 self.cloud_backend
2002 .as_ref()
2003 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
2004 .invoke(provider, context_ref, model_id, request)
2005 .await
2006 }
2007
2008 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
2009 let mut out = Vec::new();
2010 for descriptor in super::provider::registry().iter() {
2011 out.push(ProviderStatus {
2012 key: descriptor.id.as_str().to_string(),
2013 name: descriptor.display_name.to_string(),
2014 kind: CredentialKind::OAuth,
2015 configured: load_provider_auth(descriptor.id.as_str())
2016 .ok()
2017 .flatten()
2018 .is_some(),
2019 });
2020 }
2021 for spec in STATIC_PROVIDERS {
2022 out.push(ProviderStatus {
2023 key: spec.key.to_string(),
2024 name: spec.name.to_string(),
2025 kind: CredentialKind::StaticKey,
2026 configured: static_key_configured(spec.key),
2027 });
2028 }
2029 out.push(ProviderStatus {
2030 key: LOCAL_PROVIDER_KEY.to_string(),
2031 name: "Local endpoint".to_string(),
2032 kind: CredentialKind::LocalEndpoint,
2033 configured: true,
2034 });
2035 Ok(out)
2036 }
2037}
2038
2039pub struct RemoteBroker {
2045 http: reqwest::Client,
2046 endpoint: String,
2047 machine_token: String,
2048 cache: super::TokenCache,
2049}
2050
2051impl RemoteBroker {
2052 pub fn new(
2053 endpoint: impl Into<String>,
2054 machine_token: impl Into<String>,
2055 http: reqwest::Client,
2056 cache: super::TokenCache,
2057 ) -> Self {
2058 let endpoint = endpoint.into().trim_end_matches('/').to_string();
2059 let http = reqwest::Client::builder()
2062 .redirect(reqwest::redirect::Policy::none())
2063 .connect_timeout(Duration::from_secs(10))
2064 .timeout(Duration::from_secs(90))
2065 .build()
2066 .unwrap_or(http);
2067 Self {
2068 http,
2069 endpoint,
2070 machine_token: machine_token.into(),
2071 cache,
2072 }
2073 }
2074
2075 async fn post_proxy(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
2076 request.validate()?;
2077 let resp = self
2078 .http
2079 .post(format!("{}/proxy", self.endpoint))
2080 .bearer_auth(&self.machine_token)
2081 .json(request)
2082 .send()
2083 .await
2084 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
2085 match resp.status().as_u16() {
2086 401 => Err(BrokerError::Unauthorized),
2087 s @ (400 | 403) => {
2088 drop(resp);
2092 Err(BrokerError::Denied(format!(
2093 "broker rejected the request (HTTP {s})"
2094 )))
2095 }
2096 _ => Ok(resp),
2097 }
2098 }
2099}
2100
2101#[async_trait]
2102impl CredentialBroker for RemoteBroker {
2103 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
2104 let fetcher = super::BrokerClient::with_client(
2105 self.endpoint.clone(),
2106 self.machine_token.clone(),
2107 self.http.clone(),
2108 );
2109 let tok = super::resolve_remote(
2110 &fetcher,
2111 &self.cache,
2112 provider.as_str(),
2113 super::DEFAULT_MARGIN_MS,
2114 )
2115 .await
2116 .map_err(BrokerError::Transport)?;
2117 Ok(AccessToken {
2118 token: tok.access_token,
2119 expires: tok.expires,
2120 })
2121 }
2122
2123 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
2124 let mut request = request;
2125 request.stream = false;
2126 let resp = self.post_proxy(&request).await?;
2127 if !resp.status().is_success() {
2128 return Err(BrokerError::Transport(format!(
2129 "broker proxy returned HTTP {}",
2130 resp.status()
2131 )));
2132 }
2133 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
2134 serde_json::from_str::<ProxyResponse>(&body)
2135 .map_err(|e| BrokerError::Transport(format!("invalid broker proxy response: {e}")))
2136 }
2137
2138 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
2139 let mut request = request;
2140 request.stream = true;
2141 let resp = self.post_proxy(&request).await?;
2142 let status = resp.status();
2143 if !status.is_success() {
2144 drop(resp);
2146 return Err(BrokerError::Transport(format!(
2147 "broker proxy stream failed: {status}"
2148 )));
2149 }
2150 use futures::StreamExt;
2151 let stream = resp
2152 .bytes_stream()
2153 .map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
2154 Ok(Box::pin(stream))
2155 }
2156
2157 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
2158 let resp = self
2161 .http
2162 .get(format!("{}/usage", self.endpoint))
2163 .bearer_auth(&self.machine_token)
2164 .send()
2165 .await
2166 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
2167 match resp.status().as_u16() {
2168 401 => Err(BrokerError::Unauthorized),
2169 s if !(200..300).contains(&s) => {
2170 drop(resp);
2172 Err(BrokerError::Transport(format!(
2173 "broker usage returned HTTP {s}"
2174 )))
2175 }
2176 _ => {
2177 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
2178 serde_json::from_str(&body)
2179 .map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
2180 }
2181 }
2182 }
2183
2184 async fn cloud_catalog(
2185 &self,
2186 provider: CloudProviderId,
2187 context_ref: &str,
2188 allow_stale: bool,
2189 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
2190 let resp = self.http.post(format!("{}/cloud/catalog", self.endpoint)).bearer_auth(&self.machine_token).json(&serde_json::json!({"provider":provider,"context_ref":context_ref,"allow_stale":allow_stale})).send().await.map_err(|e| BrokerError::Transport(e.to_string()))?;
2191 if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
2192 return Err(BrokerError::Unauthorized);
2193 }
2194 if !resp.status().is_success() {
2195 return Err(BrokerError::Transport(format!(
2196 "cloud catalog returned HTTP {}",
2197 resp.status()
2198 )));
2199 }
2200 let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
2201 let entries: Vec<CloudCatalogEntry> = serde_json::from_str(&body)
2202 .map_err(|e| BrokerError::Transport(format!("invalid cloud catalog: {e}")))?;
2203 if entries.len() > MAX_CLOUD_CATALOG_ENTRIES
2204 || entries.iter().any(|e| {
2205 e.provider != provider || !e.context_ref.starts_with("ctx-") || e.id.len() > 512
2206 })
2207 {
2208 return Err(BrokerError::Transport(
2209 "broker returned an invalid cloud catalog".into(),
2210 ));
2211 }
2212 Ok(entries)
2213 }
2214
2215 async fn cloud_invoke(
2216 &self,
2217 provider: CloudProviderId,
2218 context_ref: &str,
2219 model_id: &str,
2220 request: InvokeRequest,
2221 ) -> Result<CloudEventStream, BrokerError> {
2222 let resp = self.http.post(format!("{}/cloud/invoke", self.endpoint)).bearer_auth(&self.machine_token).json(&serde_json::json!({"provider":provider,"context_ref":context_ref,"model_id":model_id,"request":request})).send().await.map_err(|e| BrokerError::Transport(e.to_string()))?;
2223 if !resp.status().is_success() {
2224 return Err(if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
2225 BrokerError::Unauthorized
2226 } else {
2227 BrokerError::Transport(format!("cloud invoke returned HTTP {}", resp.status()))
2228 });
2229 }
2230 use futures::StreamExt;
2231 let chunks = Box::pin(resp.bytes_stream());
2232 let stream = futures::stream::unfold(
2233 (chunks, Vec::<u8>::new(), false, false),
2234 |(mut chunks, mut buffer, done, failed)| async move {
2235 if failed {
2236 return None;
2237 }
2238 loop {
2239 if let Some(end) = buffer.iter().position(|b| *b == b'\n') {
2240 let line: Vec<u8> = buffer.drain(..=end).collect();
2241 let parsed = serde_json::from_slice::<CloudEvent>(&line[..end])
2242 .map_err(|_| BrokerError::Transport("invalid cloud event".into()));
2243 let event = match parsed {
2244 Ok(CloudEvent::Done) if done => Err(BrokerError::Transport(
2245 "duplicate cloud terminal event".into(),
2246 )),
2247 Ok(CloudEvent::Done) => Ok(CloudEvent::Done),
2248 Ok(_) if done => Err(BrokerError::Transport(
2249 "cloud data followed terminal event".into(),
2250 )),
2251 other => other,
2252 };
2253 let next_done = done || matches!(event, Ok(CloudEvent::Done));
2254 let next_failed = event.is_err();
2255 return Some((event, (chunks, buffer, next_done, next_failed)));
2256 }
2257 match chunks.next().await {
2258 Some(Ok(chunk))
2259 if buffer.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES =>
2260 {
2261 buffer.extend_from_slice(&chunk)
2262 }
2263 Some(Ok(_)) => {
2264 return Some((
2265 Err(BrokerError::Transport("cloud event exceeded limit".into())),
2266 (chunks, Vec::new(), done, true),
2267 ))
2268 }
2269 Some(Err(_)) => {
2270 return Some((
2271 Err(BrokerError::Transport("broker stream failed".into())),
2272 (chunks, Vec::new(), done, true),
2273 ))
2274 }
2275 None if buffer.is_empty() && done => return None,
2276 None if buffer.is_empty() => {
2277 return Some((
2278 Err(BrokerError::Transport(
2279 "cloud stream ended without terminal event".into(),
2280 )),
2281 (chunks, Vec::new(), done, true),
2282 ))
2283 }
2284 None => {
2285 return Some((
2286 Err(BrokerError::Transport("truncated cloud event".into())),
2287 (chunks, Vec::new(), done, true),
2288 ))
2289 }
2290 }
2291 }
2292 },
2293 );
2294 Ok(Box::pin(stream))
2295 }
2296
2297 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
2298 let resp = self
2299 .http
2300 .get(format!("{}/capabilities", self.endpoint))
2301 .bearer_auth(&self.machine_token)
2302 .send()
2303 .await
2304 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
2305 match resp.status().as_u16() {
2306 401 => Err(BrokerError::Unauthorized),
2307 s if !(200..300).contains(&s) => {
2308 Err(BrokerError::Transport(format!("broker returned HTTP {s}")))
2309 }
2310 _ => {
2311 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
2312 serde_json::from_str::<Vec<ProviderStatus>>(&body).map_err(|e| {
2313 BrokerError::Transport(format!("invalid capabilities response: {e}"))
2314 })
2315 }
2316 }
2317 }
2318}
2319
2320pub fn preflight_cloud_capability(
2329 provider: CloudProviderId,
2330 needs_tools: bool,
2331) -> Result<(), BrokerError> {
2332 if needs_tools && !provider.supports_tools() {
2333 return Err(BrokerError::UnsupportedCapability {
2334 provider: provider.to_string(),
2335 capability: "tools".into(),
2336 });
2337 }
2338 Ok(())
2339}
2340
2341pub fn broker_from_source(
2343 source: &super::CredentialSource,
2344 cache: &super::TokenCache,
2345 http: reqwest::Client,
2346) -> Arc<dyn CredentialBroker> {
2347 match source {
2348 super::CredentialSource::Local => Arc::new(LocalBroker::new(http)),
2349 super::CredentialSource::Remote {
2350 endpoint,
2351 machine_token,
2352 } => Arc::new(RemoteBroker::new(
2353 endpoint.clone(),
2354 machine_token.clone(),
2355 http,
2356 cache.clone(),
2357 )),
2358 }
2359}
2360
2361static GLOBAL_BROKER: std::sync::RwLock<Option<Arc<dyn CredentialBroker>>> =
2362 std::sync::RwLock::new(None);
2363
2364static GLOBAL_BROKER_INSTALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
2365
2366pub fn set_global_broker(broker: Arc<dyn CredentialBroker>) {
2371 *GLOBAL_BROKER.write().expect("broker registry poisoned") = Some(broker);
2372 let n = GLOBAL_BROKER_INSTALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
2373 if crate::core::memstat::mem_trace_enabled() {
2374 tracing::info!(target: "agent_core::memstat", installs = n, "global broker installed");
2375 }
2376}
2377
2378pub fn global_broker_install_count() -> u64 {
2380 GLOBAL_BROKER_INSTALLS.load(std::sync::atomic::Ordering::Relaxed)
2381}
2382
2383pub fn global_broker() -> Arc<dyn CredentialBroker> {
2386 if let Some(b) = GLOBAL_BROKER
2387 .read()
2388 .expect("broker registry poisoned")
2389 .clone()
2390 {
2391 return b;
2392 }
2393 let default: Arc<dyn CredentialBroker> = Arc::new(LocalBroker::new(
2394 reqwest::Client::builder()
2395 .connect_timeout(std::time::Duration::from_secs(10))
2396 .timeout(std::time::Duration::from_secs(300))
2397 .build()
2398 .unwrap_or_default(),
2399 ));
2400 let mut guard = GLOBAL_BROKER.write().expect("broker registry poisoned");
2401 if let Some(b) = guard.clone() {
2402 return b;
2403 }
2404 *guard = Some(default.clone());
2405 default
2406}
2407
2408pub async fn broker_access_token(
2411 provider: OAuthProviderId,
2412 source: &super::CredentialSource,
2413 cache: &super::TokenCache,
2414 http: &reqwest::Client,
2415) -> Result<AccessToken, BrokerError> {
2416 broker_from_source(source, cache, http.clone())
2417 .access_token(provider)
2418 .await
2419}
2420
2421pub fn static_key_status_map() -> BTreeMap<String, StaticKeyStatus> {
2424 STATIC_PROVIDERS
2425 .iter()
2426 .map(|s| (s.key.to_string(), static_key_status(s.key)))
2427 .collect()
2428}
2429
2430#[cfg(test)]
2431mod tests {
2432 use super::*;
2433
2434 #[test]
2439 fn preflight_rejects_tool_requiring_cloud_routes_with_typed_error() {
2440 for provider in [
2441 CloudProviderId::AzureOpenAi,
2442 CloudProviderId::AwsBedrock,
2443 CloudProviderId::GoogleVertex,
2444 ] {
2445 let err = preflight_cloud_capability(provider, true)
2446 .expect_err("tool-requiring cloud route must fail pre-flight");
2447 assert_eq!(
2448 err,
2449 BrokerError::UnsupportedCapability {
2450 provider: provider.to_string(),
2451 capability: "tools".into(),
2452 }
2453 );
2454 assert!(err.to_string().contains("text-only"));
2455 assert_eq!(preflight_cloud_capability(provider, false), Ok(()));
2456 }
2457 }
2458
2459 async fn spawn_ndjson(body: &'static str) -> String {
2460 use axum::{body::Body, routing::post, Router};
2461 use bytes::Bytes;
2462 let app = Router::new().route(
2463 "/cloud/invoke",
2464 post(move || async move {
2465 let chunks = body
2466 .as_bytes()
2467 .chunks(3)
2468 .map(|chunk| Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(chunk)));
2469 axum::response::Response::builder()
2470 .header("content-type", "application/x-ndjson")
2471 .body(Body::from_stream(futures::stream::iter(chunks)))
2472 .unwrap()
2473 }),
2474 );
2475 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2476 let addr = listener.local_addr().unwrap();
2477 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
2478 format!("http://{addr}")
2479 }
2480
2481 async fn remote_events(body: &'static str) -> Vec<Result<CloudEvent, BrokerError>> {
2482 use futures::StreamExt;
2483 let broker = RemoteBroker::new(
2484 spawn_ndjson(body).await,
2485 "opaque-machine-token",
2486 reqwest::Client::new(),
2487 super::super::TokenCache::new(),
2488 );
2489 broker
2490 .cloud_invoke(
2491 CloudProviderId::GoogleVertex,
2492 "ctx-opaque",
2493 "google-vertex/publishers/google/models/test",
2494 InvokeRequest {
2495 messages: vec![],
2496 tools: vec![],
2497 stream: true,
2498 options: Default::default(),
2499 },
2500 )
2501 .await
2502 .unwrap()
2503 .collect()
2504 .await
2505 }
2506
2507 #[tokio::test]
2508 async fn remote_ndjson_requires_exactly_one_terminal_and_stops_after_error() {
2509 let valid =
2510 remote_events("{\"type\":\"text_delta\",\"delta\":\"hi\"}\n{\"type\":\"done\"}\n")
2511 .await;
2512 assert_eq!(valid.len(), 2);
2513 assert!(matches!(valid[1], Ok(CloudEvent::Done)));
2514
2515 for body in [
2516 "{\"type\":\"text_delta\",\"delta\":\"hi\"}\n",
2517 "{\"type\":\"done\"}\n{\"type\":\"done\"}\n",
2518 "{\"type\":\"done\"}\n{\"type\":\"text_delta\",\"delta\":\"late\"}\n",
2519 ] {
2520 let events = remote_events(body).await;
2521 assert_eq!(events.iter().filter(|event| event.is_err()).count(), 1);
2522 assert!(events.last().unwrap().is_err());
2523 }
2524 }
2525
2526 #[test]
2529 fn proxy_rejects_unlisted_same_host_paths() {
2530 for (provider, path) in [
2531 ("groq", "/v1/keys"),
2532 ("groq", "/admin"),
2533 ("openrouter", "/api/v1/auth/keys"),
2534 ("local", "/audio/speech"),
2535 ] {
2536 let req = ProxyRequest {
2537 provider: provider.into(),
2538 method: ProxyMethod::Get,
2539 path: path.into(),
2540 body: None,
2541 stream: false,
2542 body_bytes: None,
2543 };
2544 match req.validate() {
2545 Err(BrokerError::Denied(msg)) => {
2546 assert!(msg.contains("allowlist"), "got: {msg}")
2547 }
2548 other => panic!("{provider} {path} must be denied, got {other:?}"),
2549 }
2550 }
2551 for path in ["/models", "/chat/completions"] {
2553 let req = ProxyRequest {
2554 provider: "groq".into(),
2555 method: ProxyMethod::Get,
2556 path: path.into(),
2557 body: None,
2558 stream: false,
2559 body_bytes: None,
2560 };
2561 assert!(req.validate().is_ok(), "{path} must be allowed");
2562 }
2563 }
2564
2565 #[test]
2568 fn proxy_rejects_oversize_request_body() {
2569 let req = ProxyRequest {
2570 provider: "groq".into(),
2571 method: ProxyMethod::Post,
2572 path: "/chat/completions".into(),
2573 body: Some(serde_json::json!({
2574 "blob": "x".repeat(MAX_PROXY_REQUEST_BYTES + 1)
2575 })),
2576 stream: false,
2577 body_bytes: None,
2578 };
2579 match req.validate() {
2580 Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
2581 other => panic!("oversize body must be denied, got {other:?}"),
2582 }
2583 }
2584
2585 #[test]
2589 fn proxy_rejects_oversize_body_bytes_without_json_body() {
2590 let req = ProxyRequest {
2591 provider: "groq".into(),
2592 method: ProxyMethod::Post,
2593 path: "/chat/completions".into(),
2594 body: None,
2595 stream: false,
2596 body_bytes: Some(bytes::Bytes::from(vec![b'x'; MAX_PROXY_REQUEST_BYTES + 1])),
2597 };
2598 match req.validate() {
2599 Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
2600 other => panic!("oversize body_bytes must be denied, got {other:?}"),
2601 }
2602 }
2603
2604 #[test]
2608 fn post_json_exact_sets_coherent_body_and_bytes() {
2609 let value = serde_json::json!({"model": "m", "stream": true});
2610 let (req, digest_bytes) =
2611 ProxyRequest::post_json_exact("groq", "/chat/completions", value.clone(), true)
2612 .expect("serializable body must construct");
2613 assert_eq!(req.method, ProxyMethod::Post);
2614 assert!(req.stream);
2615 assert_eq!(req.body.as_ref(), Some(&value));
2616 let stored = req.body_bytes.as_ref().expect("bytes must be set");
2617 assert_eq!(stored, &digest_bytes, "caller digest bytes == wire bytes");
2618 assert_eq!(
2619 serde_json::from_slice::<serde_json::Value>(stored).unwrap(),
2620 value
2621 );
2622 req.validate().expect("coherent request must validate");
2623 }
2624
2625 #[test]
2629 fn mismatched_body_bytes_rejected_by_debug_validation() {
2630 let req = ProxyRequest {
2631 provider: "groq".into(),
2632 method: ProxyMethod::Post,
2633 path: "/chat/completions".into(),
2634 body: Some(serde_json::json!({"model": "claimed"})),
2635 stream: true,
2636 body_bytes: Some(bytes::Bytes::from_static(b"{\"model\":\"actually-sent\"}")),
2637 };
2638 match req.validate() {
2639 Err(BrokerError::Denied(msg)) => {
2640 assert!(msg.contains("body_bytes"), "got: {msg}");
2641 }
2642 other => panic!("incoherent handoff must be denied, got {other:?}"),
2643 }
2644 }
2645
2646 #[test]
2647 fn proxy_request_validation_fails_closed() {
2648 let ok = ProxyRequest {
2649 provider: "groq".into(),
2650 method: ProxyMethod::Post,
2651 path: "/chat/completions".into(),
2652 body: None,
2653 stream: false,
2654 body_bytes: None,
2655 };
2656 assert!(ok.validate().is_ok());
2657
2658 let unknown = ProxyRequest {
2659 provider: "evil".into(),
2660 ..ok.clone()
2661 };
2662 assert!(matches!(
2663 unknown.validate(),
2664 Err(BrokerError::UnknownProvider(_))
2665 ));
2666
2667 let absolute = ProxyRequest {
2668 path: "https://evil.example/x".into(),
2669 ..ok.clone()
2670 };
2671 assert!(matches!(absolute.validate(), Err(BrokerError::Denied(_))));
2672
2673 let traversal = ProxyRequest {
2674 path: "/../secrets".into(),
2675 ..ok.clone()
2676 };
2677 assert!(matches!(traversal.validate(), Err(BrokerError::Denied(_))));
2678
2679 let relative = ProxyRequest {
2680 path: "chat/completions".into(),
2681 ..ok
2682 };
2683 assert!(matches!(relative.validate(), Err(BrokerError::Denied(_))));
2684 }
2685
2686 #[test]
2689 fn proxy_allows_only_pinned_kimi_code_paths() {
2690 let chat = ProxyRequest {
2691 provider: "kimi-code".into(),
2692 method: ProxyMethod::Post,
2693 path: "/chat/completions".into(),
2694 body: None,
2695 stream: true,
2696 body_bytes: None,
2697 };
2698 assert!(chat.validate().is_ok(), "kimi-code chat must be allowed");
2699
2700 for path in ["/models", "/me", "/usages"] {
2701 let get = ProxyRequest {
2702 provider: "kimi-code".into(),
2703 method: ProxyMethod::Get,
2704 path: path.into(),
2705 body: None,
2706 stream: false,
2707 body_bytes: None,
2708 };
2709 assert!(get.validate().is_ok(), "kimi-code GET {path} must pass");
2710 let post = ProxyRequest {
2712 method: ProxyMethod::Post,
2713 ..get
2714 };
2715 assert!(
2716 matches!(post.validate(), Err(BrokerError::Denied(_))),
2717 "kimi-code POST {path} must be denied"
2718 );
2719 }
2720
2721 for path in ["/feedback", "/search", "/fetch", "/completions"] {
2723 let denied = ProxyRequest {
2724 provider: "kimi-code".into(),
2725 method: ProxyMethod::Post,
2726 path: path.into(),
2727 body: None,
2728 stream: false,
2729 body_bytes: None,
2730 };
2731 assert!(
2732 matches!(denied.validate(), Err(BrokerError::Denied(_))),
2733 "kimi-code {path} must be denied"
2734 );
2735 }
2736
2737 let static_kimi = ProxyRequest {
2740 provider: "kimi".into(),
2741 method: ProxyMethod::Post,
2742 path: "/chat/completions".into(),
2743 body: None,
2744 stream: false,
2745 body_bytes: None,
2746 };
2747 assert!(static_kimi.validate().is_ok());
2748 }
2749
2750 #[test]
2754 fn proxy_rejects_oauth_providers_except_codex_catalog() {
2755 let claude = ProxyRequest {
2757 provider: "claude".into(),
2758 method: ProxyMethod::Get,
2759 path: "/models".into(),
2760 body: None,
2761 stream: false,
2762 body_bytes: None,
2763 };
2764 assert!(
2765 matches!(claude.validate(), Err(BrokerError::UnknownProvider(_))),
2766 "claude alias must not be proxyable"
2767 );
2768 let anth = ProxyRequest {
2770 provider: "anthropic".into(),
2771 method: ProxyMethod::Get,
2772 path: "/models".into(),
2773 body: None,
2774 stream: false,
2775 body_bytes: None,
2776 };
2777 assert!(
2778 matches!(anth.validate(), Err(BrokerError::Denied(_))),
2779 "anthropic /models must be denied"
2780 );
2781 let bare = ProxyRequest {
2783 provider: "openai-codex".into(),
2784 method: ProxyMethod::Get,
2785 path: "/models".into(),
2786 body: None,
2787 stream: false,
2788 body_bytes: None,
2789 };
2790 assert!(
2791 matches!(bare.validate(), Err(BrokerError::Denied(_))),
2792 "openai-codex /models without client_version must be denied"
2793 );
2794 }
2795
2796 #[test]
2800 fn proxy_allows_only_pinned_openai_codex_models_path() {
2801 let allowed = ProxyRequest {
2802 provider: "openai-codex".into(),
2803 method: ProxyMethod::Get,
2804 path: "/codex/models?client_version=0.153.3".into(),
2805 body: None,
2806 stream: false,
2807 body_bytes: None,
2808 };
2809 assert!(
2810 allowed.validate().is_ok(),
2811 "exact codex models path must be allowed"
2812 );
2813
2814 for path in [
2815 "/codex/models",
2816 "/codex/models?client_version=",
2817 "/codex/models?foo=1",
2818 "/codex/models?client_version=0.153.3&extra=1",
2819 "/codex/responses",
2820 "/models",
2821 "/models?client_version=0.153.3",
2822 "/backend-api/codex/models?client_version=0.153.3",
2823 "/v1/models?client_version=0.153.3",
2824 "/chat/completions",
2825 ] {
2826 let req = ProxyRequest {
2827 provider: "openai-codex".into(),
2828 method: ProxyMethod::Get,
2829 path: path.into(),
2830 body: None,
2831 stream: false,
2832 body_bytes: None,
2833 };
2834 assert!(
2835 matches!(req.validate(), Err(BrokerError::Denied(_))),
2836 "openai-codex path {path} must be denied"
2837 );
2838 }
2839
2840 let post = ProxyRequest {
2842 provider: "openai-codex".into(),
2843 method: ProxyMethod::Post,
2844 path: "/codex/models?client_version=0.153.3".into(),
2845 body: None,
2846 stream: false,
2847 body_bytes: None,
2848 };
2849 assert!(
2850 matches!(post.validate(), Err(BrokerError::Denied(_))),
2851 "POST openai-codex models must be denied"
2852 );
2853 }
2854
2855 #[test]
2857 fn proxy_allows_only_pinned_github_copilot_paths() {
2858 let models = ProxyRequest {
2859 provider: "github-copilot".into(),
2860 method: ProxyMethod::Get,
2861 path: "/models".into(),
2862 body: None,
2863 stream: false,
2864 body_bytes: None,
2865 };
2866 assert!(models.validate().is_ok());
2867
2868 let chat = ProxyRequest {
2869 provider: "github-copilot".into(),
2870 method: ProxyMethod::Post,
2871 path: "/chat/completions".into(),
2872 body: None,
2873 stream: false,
2874 body_bytes: None,
2875 };
2876 assert!(chat.validate().is_ok());
2877 let responses = ProxyRequest {
2878 provider: "github-copilot".into(),
2879 method: ProxyMethod::Post,
2880 path: "/responses".into(),
2881 body: None,
2882 stream: false,
2883 body_bytes: None,
2884 };
2885 assert!(responses.validate().is_ok());
2886 for path in ["/v1/messages", "/models?x=1", "/embeddings"] {
2887 let request = ProxyRequest {
2888 provider: "github-copilot".into(),
2889 method: ProxyMethod::Post,
2890 path: path.into(),
2891 body: None,
2892 stream: false,
2893 body_bytes: None,
2894 };
2895 assert!(matches!(request.validate(), Err(BrokerError::Denied(_))));
2896 }
2897 }
2898
2899 #[test]
2901 fn proxy_allows_only_pinned_anthropic_models_path() {
2902 for path in [
2903 "/v1/models",
2904 "/v1/models?limit=100",
2905 "/v1/models?limit=100&after_id=claude-opus-4-7",
2906 ] {
2907 let req = ProxyRequest {
2908 provider: "anthropic".into(),
2909 method: ProxyMethod::Get,
2910 path: path.into(),
2911 body: None,
2912 stream: false,
2913 body_bytes: None,
2914 };
2915 assert!(req.validate().is_ok(), "{path} must be allowed");
2916 }
2917 for path in [
2918 "/v1/messages",
2919 "/models",
2920 "/v1/models?foo=1",
2921 "/v1/models?limit=",
2922 "/api/oauth/usage",
2923 "/v1/models?limit=100&evil=1",
2924 "/v1/models?limit=100&limit=50",
2925 "/v1/models?after_id=a&after_id=b",
2926 "/v1/models?limit=100&after_id=a&limit=50",
2927 ] {
2928 let req = ProxyRequest {
2929 provider: "anthropic".into(),
2930 method: ProxyMethod::Get,
2931 path: path.into(),
2932 body: None,
2933 stream: false,
2934 body_bytes: None,
2935 };
2936 assert!(
2937 matches!(req.validate(), Err(BrokerError::Denied(_))),
2938 "anthropic path {path} must be denied"
2939 );
2940 }
2941 }
2942
2943 #[test]
2944 fn anthropic_oauth_catalog_headers_are_bearer_version_beta_not_x_api_key() {
2945 let headers = anthropic_oauth_catalog_request_headers();
2946 let names: Vec<_> = headers.iter().map(|(n, _)| *n).collect();
2947 assert!(names.contains(&"anthropic-version"));
2948 assert!(names.contains(&"anthropic-beta"));
2949 assert!(names.contains(&"accept"));
2950 assert!(!names.iter().any(|n| n.eq_ignore_ascii_case("x-api-key")));
2951 assert!(!names
2952 .iter()
2953 .any(|n| n.eq_ignore_ascii_case("authorization")));
2954 let beta = headers
2955 .iter()
2956 .find(|(n, _)| *n == "anthropic-beta")
2957 .map(|(_, v)| *v)
2958 .expect("beta");
2959 assert_eq!(beta, "oauth-2025-04-20");
2960 }
2961
2962 #[test]
2963 fn anthropic_models_path_round_trip_and_rejects_duplicates() {
2964 assert!(is_allowed_anthropic_path("/v1/models"));
2965 assert!(is_allowed_anthropic_path("/v1/models?limit=100"));
2966 assert!(is_allowed_anthropic_path(
2967 "/v1/models?limit=100&after_id=claude-opus-4-7"
2968 ));
2969 assert!(!is_allowed_anthropic_path("/v1/models?limit=100&limit=50"));
2970 assert!(!is_allowed_anthropic_path(
2971 "/v1/models?after_id=a&after_id=b"
2972 ));
2973 let page0 = "/v1/models?limit=100";
2975 let page1 = "/v1/models?limit=100&after_id=model-1";
2976 assert!(is_allowed_anthropic_path(page0));
2977 assert!(is_allowed_anthropic_path(page1));
2978 }
2979
2980 #[test]
2982 fn proxy_allows_only_pinned_google_gemini_paths() {
2983 for path in [
2984 "/v1internal:loadCodeAssist",
2985 "/v1internal:onboardUser",
2986 "/v1internal:streamGenerateContent",
2987 "/v1internal:countTokens",
2988 "/v1internal/operations/op-12345",
2989 ] {
2990 let req = ProxyRequest {
2991 provider: "google-gemini".into(),
2992 method: ProxyMethod::Post,
2993 path: path.into(),
2994 body: None,
2995 stream: false,
2996 body_bytes: None,
2997 };
2998 assert!(req.validate().is_ok(), "{path} must be allowed");
2999 }
3000 for path in [
3001 "/v1internal:listExperiments",
3003 "/v1internal:fetchAdminControls",
3004 "/v1internal:setCodeAssistGlobalUserSetting",
3005 "/v1internal:generateContent",
3006 "/v2/models",
3008 "/v1beta/models",
3009 "/",
3011 "/v1internal:",
3012 ] {
3013 let req = ProxyRequest {
3014 provider: "google-gemini".into(),
3015 method: ProxyMethod::Post,
3016 path: path.into(),
3017 body: None,
3018 stream: false,
3019 body_bytes: None,
3020 };
3021 assert!(
3022 matches!(req.validate(), Err(BrokerError::Denied(_))),
3023 "{path} must be denied"
3024 );
3025 }
3026 }
3027
3028 #[test]
3029 fn access_token_type_has_no_refresh_field() {
3030 let t: AccessToken =
3033 serde_json::from_str(r#"{"token":"sk-x","expires":123,"refresh":"MUST-NOT-EXIST"}"#)
3034 .unwrap();
3035 let round = serde_json::to_value(&t).unwrap();
3036 assert_eq!(round.get("refresh"), None);
3037 assert_eq!(round["token"], "sk-x");
3038 }
3039
3040 #[test]
3041 fn broker_error_display_never_echoes_values() {
3042 let e = BrokerError::NotConfigured("groq".into());
3044 assert!(!format!("{e}").contains("sk-"));
3045 let e = BrokerError::Unauthorized;
3046 assert_eq!(format!("{e}"), "broker rejected machine auth");
3047 }
3048
3049 #[test]
3050 fn mask_key_is_short_and_lossy() {
3051 assert_eq!(mask_key("gsk-live-1234567890abcdef"), "gsk-…cdef");
3052 assert_eq!(mask_key("short"), "…");
3053 }
3054
3055 #[test]
3056 fn proxy_request_serde_roundtrip() {
3057 let req = ProxyRequest {
3058 provider: "openrouter".into(),
3059 method: ProxyMethod::Get,
3060 path: "/models".into(),
3061 body: Some(serde_json::json!({"a": 1})),
3062 stream: true,
3063 body_bytes: None,
3064 };
3065 let json = serde_json::to_string(&req).unwrap();
3066 let back: ProxyRequest = serde_json::from_str(&json).unwrap();
3067 assert_eq!(back.provider, "openrouter");
3068 assert!(back.stream);
3069 assert!(matches!(back.method, ProxyMethod::Get));
3070 }
3071
3072 async fn spawn_upstream(app: axum::Router) -> String {
3075 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3076 let addr = listener.local_addr().unwrap();
3077 tokio::spawn(async move {
3078 axum::serve(listener, app).await.unwrap();
3079 });
3080 format!("http://{addr}")
3081 }
3082
3083 #[tokio::test]
3087 async fn local_broker_streams_local_endpoint_and_applies_key() {
3088 use axum::routing::post;
3089 let seen_auth = Arc::new(std::sync::Mutex::new(String::new()));
3090 let seen = seen_auth.clone();
3091 let app = axum::Router::new().route(
3092 "/chat/completions",
3093 post(move |headers: axum::http::HeaderMap| {
3094 let seen = seen.clone();
3095 async move {
3096 *seen.lock().unwrap() = headers
3097 .get("authorization")
3098 .and_then(|v| v.to_str().ok())
3099 .unwrap_or("")
3100 .to_string();
3101 (
3102 [("content-type", "text/event-stream")],
3103 "data: {\"x\":1}\n\ndata: [DONE]\n\n",
3104 )
3105 }
3106 }),
3107 );
3108 let url = spawn_upstream(app).await;
3109 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3110 let mut stream = broker
3111 .proxy_stream(ProxyRequest {
3112 provider: LOCAL_PROVIDER_KEY.into(),
3113 method: ProxyMethod::Post,
3114 path: "/chat/completions".into(),
3115 body: Some(serde_json::json!({"model": "m"})),
3116 stream: true,
3117 body_bytes: None,
3118 })
3119 .await
3120 .expect("stream must open");
3121 use futures::StreamExt;
3122 let mut collected = Vec::new();
3123 while let Some(chunk) = stream.next().await {
3124 collected.extend_from_slice(&chunk.unwrap());
3125 }
3126 let text = String::from_utf8(collected).unwrap();
3127 assert!(text.contains("data: {\"x\":1}"));
3128 assert!(text.contains("[DONE]"));
3129 assert_eq!(&*seen_auth.lock().unwrap(), "Bearer local");
3131 }
3132
3133 #[tokio::test]
3136 async fn local_broker_stream_error_is_typed_and_keyless() {
3137 use axum::routing::post;
3138 let app = axum::Router::new().route(
3139 "/chat/completions",
3140 post(|| async {
3141 (
3142 axum::http::StatusCode::UNAUTHORIZED,
3143 "{\"error\":\"bad key\"}",
3144 )
3145 }),
3146 );
3147 let url = spawn_upstream(app).await;
3148 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3149 let err = match broker
3150 .proxy_stream(ProxyRequest {
3151 provider: LOCAL_PROVIDER_KEY.into(),
3152 method: ProxyMethod::Post,
3153 path: "/chat/completions".into(),
3154 body: None,
3155 stream: true,
3156 body_bytes: None,
3157 })
3158 .await
3159 {
3160 Err(e) => e,
3161 Ok(_) => panic!("non-2xx upstream must not yield a stream"),
3162 };
3163 let msg = format!("{err}");
3164 assert!(msg.contains("401"), "got: {msg}");
3165 assert!(
3166 !msg.to_lowercase().contains("bearer"),
3167 "no auth material in errors"
3168 );
3169 }
3170
3171 #[tokio::test]
3173 async fn local_broker_proxy_returns_status_and_body() {
3174 use axum::routing::get;
3175 let app =
3176 axum::Router::new().route("/models", get(|| async { "{\"data\":[{\"id\":\"m1\"}]}" }));
3177 let url = spawn_upstream(app).await;
3178 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3179 let resp = broker
3180 .proxy(ProxyRequest {
3181 provider: LOCAL_PROVIDER_KEY.into(),
3182 method: ProxyMethod::Get,
3183 path: "/models".into(),
3184 body: None,
3185 stream: false,
3186 body_bytes: None,
3187 })
3188 .await
3189 .unwrap();
3190 assert_eq!(resp.status, 200);
3191 assert!(resp.body.contains("m1"));
3192 }
3193
3194 #[tokio::test]
3197 async fn local_broker_proxy_rejects_oversize_response() {
3198 use axum::routing::get;
3199 let app = axum::Router::new().route("/models", get(|| async { "x".repeat(4096) }));
3200 let url = spawn_upstream(app).await;
3201 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
3202 .with_max_response_bytes(64);
3203 let err = broker
3204 .proxy(ProxyRequest {
3205 provider: LOCAL_PROVIDER_KEY.into(),
3206 method: ProxyMethod::Get,
3207 path: "/models".into(),
3208 body: None,
3209 stream: false,
3210 body_bytes: None,
3211 })
3212 .await
3213 .expect_err("oversize body must be rejected");
3214 let msg = format!("{err}");
3215 assert!(msg.contains("limit"), "got: {msg}");
3216 }
3217
3218 #[tokio::test]
3221 async fn local_broker_stream_error_is_status_only() {
3222 use axum::routing::post;
3223 let app = axum::Router::new().route(
3224 "/chat/completions",
3225 post(|| async {
3226 (
3227 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
3228 format!("\x1b[2J{}", "E".repeat(64 * 1024)),
3229 )
3230 }),
3231 );
3232 let url = spawn_upstream(app).await;
3233 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3234 let err = broker
3235 .proxy_stream(ProxyRequest {
3236 provider: LOCAL_PROVIDER_KEY.into(),
3237 method: ProxyMethod::Post,
3238 path: "/chat/completions".into(),
3239 body: None,
3240 stream: true,
3241 body_bytes: None,
3242 })
3243 .await
3244 .err()
3245 .expect("non-2xx upstream must not yield a stream");
3246 let msg = format!("{err}");
3247 assert!(
3248 msg.len() < 700,
3249 "error must be bounded, got {} bytes",
3250 msg.len()
3251 );
3252 assert!(!msg.contains('\x1b'), "control chars must be stripped");
3253 assert!(msg.contains("500"));
3254 }
3255
3256 #[tokio::test]
3259 async fn local_broker_buffered_request_times_out() {
3260 use axum::routing::get;
3261 let app = axum::Router::new().route(
3262 "/models",
3263 get(|| async {
3264 tokio::time::sleep(Duration::from_secs(30)).await;
3265 "too late"
3266 }),
3267 );
3268 let url = spawn_upstream(app).await;
3269 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
3270 .with_request_timeout(Duration::from_millis(200));
3271 let started = std::time::Instant::now();
3272 let err = broker
3273 .proxy(ProxyRequest {
3274 provider: LOCAL_PROVIDER_KEY.into(),
3275 method: ProxyMethod::Get,
3276 path: "/models".into(),
3277 body: None,
3278 stream: false,
3279 body_bytes: None,
3280 })
3281 .await
3282 .expect_err("hung upstream must time out");
3283 assert!(matches!(err, BrokerError::Transport(_)));
3284 assert!(
3285 started.elapsed() < Duration::from_secs(5),
3286 "timeout must be enforced promptly"
3287 );
3288 }
3289
3290 #[tokio::test]
3293 async fn global_broker_defaults_to_in_process_local() {
3294 let broker = global_broker();
3295 let caps = broker
3296 .capabilities()
3297 .await
3298 .expect("in-process broker needs no daemon");
3299 assert!(caps.iter().any(|c| c.key == LOCAL_PROVIDER_KEY));
3300 let again = global_broker();
3302 assert!(again.capabilities().await.is_ok());
3303 }
3304
3305 const HOSTILE_SENTINEL: &str = "ZX9-HOSTILE-SENTINEL-7Q";
3315
3316 fn hostile_body() -> String {
3319 format!(
3320 "{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} \
3321 {{\\\"messages\\\":[{{\\\"content\\\":\\\"my secret prompt\\\"}}],\
3322 \\\"tools\\\":[{{\\\"input_schema\\\":{{}}}}]}}\"}}}}"
3323 )
3324 }
3325
3326 fn hostile_broker_json() -> String {
3329 format!(
3330 "{{\"error\":\"ECHOED {HOSTILE_SENTINEL} input_schema \
3331 my secret prompt\"}}"
3332 )
3333 }
3334
3335 fn assert_no_hostile_bytes(msg: &str) {
3336 assert!(
3337 !msg.contains(HOSTILE_SENTINEL),
3338 "sentinel leaked into error: {msg}"
3339 );
3340 assert!(!msg.contains("ECHOED"), "echoed body leaked: {msg}");
3341 assert!(
3342 !msg.contains("input_schema"),
3343 "request-shaped data leaked: {msg}"
3344 );
3345 assert!(
3346 !msg.contains("secret prompt"),
3347 "prompt content leaked: {msg}"
3348 );
3349 }
3350
3351 #[tokio::test]
3356 async fn local_broker_stream_error_omits_upstream_body() {
3357 use axum::routing::post;
3358 let app = axum::Router::new().route(
3359 "/chat/completions",
3360 post(|| async {
3361 (
3362 axum::http::StatusCode::TOO_MANY_REQUESTS,
3363 format!("\x1b[2J{}", hostile_body()),
3364 )
3365 }),
3366 );
3367 let url = spawn_upstream(app).await;
3368 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3369 let err = broker
3370 .proxy_stream(ProxyRequest {
3371 provider: LOCAL_PROVIDER_KEY.into(),
3372 method: ProxyMethod::Post,
3373 path: "/chat/completions".into(),
3374 body: None,
3375 stream: true,
3376 body_bytes: None,
3377 })
3378 .await
3379 .err()
3380 .expect("non-2xx upstream must not yield a stream");
3381 assert!(matches!(err, BrokerError::Transport(_)));
3382 let msg = format!("{err}");
3383 assert!(
3384 msg.contains("provider request failed: 429 Too Many Requests"),
3385 "status prefix must survive for engine classification: {msg}"
3386 );
3387 assert!(
3388 !msg.contains('\x1b'),
3389 "escapes must not reach errors: {msg}"
3390 );
3391 assert_no_hostile_bytes(&msg);
3392 }
3393
3394 async fn stream_error_message(status: axum::http::StatusCode, body: String) -> String {
3397 use axum::routing::post;
3398 let app = axum::Router::new().route(
3399 "/chat/completions",
3400 post(move || async move { (status, body) }),
3401 );
3402 let url = spawn_upstream(app).await;
3403 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3404 let err = broker
3405 .proxy_stream(ProxyRequest {
3406 provider: LOCAL_PROVIDER_KEY.into(),
3407 method: ProxyMethod::Post,
3408 path: "/chat/completions".into(),
3409 body: None,
3410 stream: true,
3411 body_bytes: None,
3412 })
3413 .await
3414 .err()
3415 .expect("non-2xx upstream must not yield a stream");
3416 assert!(matches!(err, BrokerError::Transport(_)));
3417 format!("{err}")
3418 }
3419
3420 fn after_marker(msg: &str) -> &str {
3422 let marker = "provider request failed: ";
3423 let idx = msg.find(marker).expect("marker present");
3424 &msg[idx + marker.len()..]
3425 }
3426
3427 #[tokio::test]
3433 async fn local_broker_stream_403_kimi_quota_surfaces_vetted_label_only() {
3434 let kimi_body = "{\"error\":{\"type\":\"access_terminated_error\",\"message\":\
3435 \"You've reached your weekly (7-day) usage limit. Your quota will \
3436 reset when the current 7-day window ends. \"}}";
3437 let msg = stream_error_message(axum::http::StatusCode::FORBIDDEN, kimi_body.into()).await;
3438 assert!(
3439 msg.ends_with("provider request failed: 403 Forbidden [access_terminated_error]"),
3440 "expected exact labelled prefix: {msg}"
3441 );
3442 assert!(!msg.contains("weekly"), "body text leaked: {msg}");
3443 assert!(!msg.contains("quota"), "body text leaked: {msg}");
3444 assert!(!msg.contains('{'), "raw JSON leaked: {msg}");
3445 assert!(
3446 !after_marker(&msg).contains(':'),
3447 "suffix must not contain ':' (engine truncates there): {msg}"
3448 );
3449 }
3450
3451 #[tokio::test]
3454 async fn local_broker_stream_429_unvetted_type_yields_no_label() {
3455 let body = format!(
3456 "{{\"error\":{{\"type\":\"<script>ECHO:secret\",\"message\":\"ECHOED {HOSTILE_SENTINEL}\"}}}}"
3457 );
3458 let msg = stream_error_message(axum::http::StatusCode::TOO_MANY_REQUESTS, body).await;
3459 assert!(
3460 msg.ends_with("provider request failed: 429 Too Many Requests"),
3461 "no label expected: {msg}"
3462 );
3463 assert!(!msg.contains('['), "unexpected label: {msg}");
3464 assert!(!msg.contains("script"), "hostile type leaked: {msg}");
3465 assert!(!msg.contains("secret"), "hostile type leaked: {msg}");
3466 assert!(!after_marker(&msg).contains(':'), "got: {msg}");
3467 assert_no_hostile_bytes(&msg);
3468 }
3469
3470 #[tokio::test]
3472 async fn local_broker_stream_html_body_yields_no_label() {
3473 let body =
3474 format!("<html><body>ECHOED {HOSTILE_SENTINEL} access_terminated_error</body></html>");
3475 let msg = stream_error_message(axum::http::StatusCode::BAD_GATEWAY, body).await;
3476 assert!(
3477 msg.ends_with("provider request failed: 502 Bad Gateway"),
3478 "no label expected: {msg}"
3479 );
3480 assert!(!msg.contains("html"), "body leaked: {msg}");
3481 assert!(!msg.contains('['), "unexpected label: {msg}");
3482 assert_no_hostile_bytes(&msg);
3483 }
3484
3485 #[tokio::test]
3488 async fn local_broker_stream_oversized_body_yields_no_label_without_panic() {
3489 let pad = "x".repeat(MAX_PROXY_ERROR_CLASSIFY_BYTES * 4);
3492 let body = format!(
3493 "{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} {pad}\",\"type\":\"access_terminated_error\"}}}}"
3494 );
3495 let msg = stream_error_message(axum::http::StatusCode::FORBIDDEN, body).await;
3496 assert!(
3497 msg.ends_with("provider request failed: 403 Forbidden"),
3498 "no label expected for oversized body: {msg}"
3499 );
3500 assert!(!msg.contains('['), "unexpected label: {msg}");
3501 assert!(
3502 msg.len() < 200,
3503 "oversized body leaked: {} bytes",
3504 msg.len()
3505 );
3506 assert_no_hostile_bytes(&msg);
3507 }
3508
3509 #[tokio::test]
3513 async fn read_error_body_for_classification_truncates_and_never_fails() {
3514 use axum::routing::get;
3515 let app = axum::Router::new()
3516 .route("/big", get(|| async { "a".repeat(100) }))
3517 .route("/utf8", get(|| async { format!("{}é", "a".repeat(9)) }));
3518 let url = spawn_upstream(app).await;
3519 let client = reqwest::Client::new();
3520
3521 let resp = client.get(format!("{url}/big")).send().await.unwrap();
3522 let got = read_error_body_for_classification(resp, 10).await;
3523 assert_eq!(got.as_deref(), Some("aaaaaaaaaa"));
3524
3525 let resp = client.get(format!("{url}/utf8")).send().await.unwrap();
3527 let got = read_error_body_for_classification(resp, 10).await;
3528 assert_eq!(got, None);
3529
3530 let resp = client.get(format!("{url}/utf8")).send().await.unwrap();
3532 let got = read_error_body_for_classification(resp, 1024).await;
3533 assert_eq!(got.as_deref(), Some("aaaaaaaaaé"));
3534 }
3535
3536 #[tokio::test]
3540 #[serial_test::serial]
3541 async fn local_broker_usage_error_omits_upstream_body() {
3542 use axum::routing::get;
3543 let app = axum::Router::new().route(
3544 "/usage",
3545 get(|| async {
3546 (
3547 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
3548 hostile_body(),
3549 )
3550 }),
3551 );
3552 let url = spawn_upstream(app).await;
3553
3554 let dir = tempfile::tempdir().unwrap();
3560 let base = dir.path().join(".synaps-cli");
3561 std::fs::create_dir_all(&base).unwrap();
3562 let creds = crate::auth::OAuthCredentials {
3563 auth_type: "oauth".into(),
3564 refresh: String::new(),
3565 access: "test-access-token".into(),
3566 expires: crate::epoch_millis() + 3_600_000,
3567 account_id: None,
3568 };
3569 super::super::storage::save_provider_auth_at_test_hook(
3570 &base.join("auth.json"),
3571 "anthropic",
3572 &creds,
3573 )
3574 .unwrap();
3575 let old_base = std::env::var("SYNAPS_BASE_DIR").ok();
3576 let old_home = std::env::var("HOME").ok();
3577 std::env::remove_var("SYNAPS_BASE_DIR");
3578 std::env::set_var("HOME", dir.path());
3579
3580 let broker = LocalBroker::new(reqwest::Client::new())
3581 .with_anthropic_usage_url(format!("{url}/usage"));
3582 let result = broker.anthropic_usage().await;
3583
3584 match old_home {
3585 Some(v) => std::env::set_var("HOME", v),
3586 None => std::env::remove_var("HOME"),
3587 }
3588 if let Some(v) = old_base {
3589 std::env::set_var("SYNAPS_BASE_DIR", v);
3590 }
3591
3592 let err = result.expect_err("non-2xx usage response must fail");
3593 assert!(matches!(err, BrokerError::Transport(_)));
3594 let msg = format!("{err}");
3595 assert!(
3596 msg.contains("usage request failed") && msg.contains("500"),
3597 "operation + status must survive: {msg}"
3598 );
3599 assert_no_hostile_bytes(&msg);
3600 }
3601
3602 async fn spawn_remote_broker_error(status: u16) -> RemoteBroker {
3603 use axum::routing::{get, post};
3604 let code = axum::http::StatusCode::from_u16(status).unwrap();
3605 let app = axum::Router::new()
3606 .route(
3607 "/proxy",
3608 post(move || async move { (code, hostile_broker_json()) }),
3609 )
3610 .route(
3611 "/usage",
3612 get(move || async move { (code, hostile_broker_json()) }),
3613 );
3614 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3615 let addr = listener.local_addr().unwrap();
3616 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
3617 RemoteBroker::new(
3618 format!("http://{addr}"),
3619 "opaque-machine-token",
3620 reqwest::Client::new(),
3621 super::super::TokenCache::new(),
3622 )
3623 }
3624
3625 fn groq_request(stream: bool) -> ProxyRequest {
3626 ProxyRequest {
3627 provider: "groq".into(),
3628 method: ProxyMethod::Post,
3629 path: "/chat/completions".into(),
3630 body: None,
3631 stream,
3632 body_bytes: None,
3633 }
3634 }
3635
3636 #[tokio::test]
3639 async fn remote_broker_denied_omits_broker_error_body() {
3640 for status in [400u16, 403] {
3641 let broker = spawn_remote_broker_error(status).await;
3642 let err = broker
3643 .proxy(groq_request(false))
3644 .await
3645 .expect_err("4xx broker response must fail");
3646 assert!(
3647 matches!(err, BrokerError::Denied(_)),
3648 "typed category must survive, got {err:?}"
3649 );
3650 let msg = format!("{err}");
3651 assert!(msg.contains(&status.to_string()), "status lost: {msg}");
3652 assert_no_hostile_bytes(&msg);
3653 }
3654 }
3655
3656 #[tokio::test]
3659 async fn remote_broker_stream_error_omits_broker_body() {
3660 let broker = spawn_remote_broker_error(500).await;
3661 let err = broker
3662 .proxy_stream(groq_request(true))
3663 .await
3664 .err()
3665 .expect("non-2xx broker response must not yield a stream");
3666 assert!(matches!(err, BrokerError::Transport(_)));
3667 let msg = format!("{err}");
3668 assert!(
3669 msg.contains("broker proxy stream failed") && msg.contains("500"),
3670 "operation + status must survive: {msg}"
3671 );
3672 assert_no_hostile_bytes(&msg);
3673 }
3674
3675 #[tokio::test]
3678 async fn remote_broker_usage_error_omits_broker_body() {
3679 let broker = spawn_remote_broker_error(502).await;
3680 let err = broker
3681 .anthropic_usage()
3682 .await
3683 .expect_err("non-2xx broker response must fail");
3684 assert!(matches!(err, BrokerError::Transport(_)));
3685 let msg = format!("{err}");
3686 assert!(
3687 msg.contains("broker usage returned HTTP 502"),
3688 "operation + status must survive: {msg}"
3689 );
3690 assert_no_hostile_bytes(&msg);
3691 }
3692
3693 #[tokio::test]
3696 async fn capabilities_expose_no_secret_fields() {
3697 let broker = LocalBroker::new(reqwest::Client::new());
3698 let caps = broker.capabilities().await.unwrap();
3699 assert!(caps
3700 .iter()
3701 .any(|c| c.key == "anthropic" && c.kind == CredentialKind::OAuth));
3702 assert!(caps
3703 .iter()
3704 .any(|c| c.key == "groq" && c.kind == CredentialKind::StaticKey));
3705 let json = serde_json::to_string(&caps).unwrap();
3706 for field in ["key\":", "name\":", "kind\":", "configured\":"] {
3707 assert!(json.contains(field));
3708 }
3709 assert!(!json.contains("refresh"));
3710 assert!(!json.contains("access"));
3711 }
3712}