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