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
1727fn discover_legacy_static_key(spec: &StaticProviderSpec) -> Option<String> {
1730 if let Some(v) = crate::config::get_provider_keys().get(spec.key) {
1731 if !v.is_empty() {
1732 return Some(v.clone());
1733 }
1734 }
1735 spec.env_vars
1736 .iter()
1737 .find_map(|var| std::env::var(var).ok().filter(|v| !v.is_empty()))
1738}
1739
1740pub fn static_key_configured(provider: &str) -> bool {
1743 if provider == LOCAL_PROVIDER_KEY {
1744 return true; }
1746 let Some(spec) = static_provider(provider) else {
1747 return false;
1748 };
1749 storage::load_static_key(provider).ok().flatten().is_some()
1750 || discover_legacy_static_key(spec).is_some()
1751}
1752
1753pub fn static_key_status(provider: &str) -> StaticKeyStatus {
1757 let stored = storage::load_static_key(provider)
1758 .ok()
1759 .flatten()
1760 .or_else(|| {
1761 crate::config::get_provider_keys()
1762 .get(provider)
1763 .filter(|s| !s.is_empty())
1764 .cloned()
1765 });
1766 if let Some(key) = stored {
1767 return StaticKeyStatus::Configured {
1768 masked: mask_key(&key),
1769 };
1770 }
1771 if let Some(spec) = static_provider(provider) {
1772 if spec
1773 .env_vars
1774 .iter()
1775 .any(|var| std::env::var(var).is_ok_and(|v| !v.is_empty()))
1776 {
1777 return StaticKeyStatus::FromEnv;
1778 }
1779 }
1780 StaticKeyStatus::NotSet
1781}
1782
1783fn mask_key(key: &str) -> String {
1785 let chars: Vec<char> = key.chars().collect();
1786 if chars.len() <= 8 {
1787 return "…".to_string();
1788 }
1789 let head: String = chars[..4].iter().collect();
1790 let tail: String = chars[chars.len() - 4..].iter().collect();
1791 format!("{head}…{tail}")
1792}
1793
1794pub fn configured_static_provider_keys() -> std::collections::BTreeSet<String> {
1796 STATIC_PROVIDERS
1797 .iter()
1798 .filter(|s| static_key_configured(s.key))
1799 .map(|s| s.key.to_string())
1800 .collect()
1801}
1802
1803pub fn local_endpoint_url() -> String {
1806 local_endpoint_config()
1807 .unwrap_or_else(|| LOCAL_DEFAULT_BASE_URL.to_string())
1808 .trim_end_matches('/')
1809 .to_string()
1810}
1811
1812pub fn local_endpoint_config() -> Option<String> {
1814 crate::config::get_provider_keys()
1815 .get("local.url")
1816 .filter(|s| !s.is_empty())
1817 .cloned()
1818 .or_else(|| {
1819 std::env::var("LOCAL_ENDPOINT")
1820 .ok()
1821 .filter(|s| !s.is_empty())
1822 })
1823}
1824
1825pub fn local_model_ids() -> Vec<String> {
1827 crate::config::get_provider_keys()
1828 .get("local.models")
1829 .map(|value| {
1830 value
1831 .split(',')
1832 .map(|s| s.trim().to_string())
1833 .filter(|s| !s.is_empty())
1834 .collect()
1835 })
1836 .unwrap_or_default()
1837}
1838
1839pub fn anthropic_credential_available() -> bool {
1842 let oauth = load_provider_auth(OAuthProviderId::Anthropic.as_str())
1843 .ok()
1844 .flatten()
1845 .map(|c| c.auth_type == "oauth" && !c.access.is_empty())
1846 .unwrap_or(false);
1847 oauth || std::env::var("ANTHROPIC_API_KEY").is_ok_and(|v| !v.is_empty())
1848}
1849
1850pub fn oauth_provider_logged_in(provider: OAuthProviderId) -> bool {
1854 load_provider_auth(provider.as_str())
1855 .ok()
1856 .flatten()
1857 .is_some_and(|creds| {
1858 creds.auth_type == "oauth" && (!creds.refresh.is_empty() || !creds.access.is_empty())
1859 })
1860}
1861
1862#[async_trait]
1863impl CredentialBroker for LocalBroker {
1864 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
1865 let creds = super::ensure_fresh_provider_token(&self.http, provider)
1866 .await
1867 .map_err(BrokerError::Credential)?;
1868 Ok(AccessToken {
1870 token: creds.access,
1871 expires: creds.expires,
1872 })
1873 }
1874
1875 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
1876 let resp = self.send(&request).await?;
1877 let status = resp.status().as_u16();
1878 let body = read_body_capped(resp, self.max_response_bytes).await?;
1879 Ok(ProxyResponse { status, body })
1880 }
1881
1882 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
1883 let mut request = request;
1884 request.stream = true;
1885 let resp = self.send(&request).await?;
1886 let status = resp.status();
1887 if !status.is_success() {
1888 drop(resp);
1895 return Err(BrokerError::Transport(format!(
1896 "provider request failed: {status}"
1897 )));
1898 }
1899 use futures::StreamExt;
1900 let stream = resp
1901 .bytes_stream()
1902 .map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
1903 Ok(Box::pin(stream))
1904 }
1905
1906 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
1907 let token = self.access_token(OAuthProviderId::Anthropic).await?;
1910 let url = self
1911 .anthropic_usage_url
1912 .as_deref()
1913 .unwrap_or(ANTHROPIC_USAGE_URL);
1914 let resp = self
1915 .http
1916 .get(url)
1917 .timeout(self.request_timeout)
1918 .bearer_auth(&token.token)
1919 .header("anthropic-beta", ANTHROPIC_OAUTH_BETA)
1920 .send()
1921 .await
1922 .map_err(|e| BrokerError::Transport(format!("usage request failed: {e}")))?;
1923 let status = resp.status();
1924 if !status.is_success() {
1925 drop(resp);
1928 return Err(BrokerError::Transport(format!(
1929 "usage request failed: {status}"
1930 )));
1931 }
1932 let body = read_body_capped(resp, self.max_response_bytes).await?;
1933 serde_json::from_str(&body)
1934 .map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
1935 }
1936
1937 async fn cloud_catalog(
1938 &self,
1939 provider: CloudProviderId,
1940 context_ref: &str,
1941 allow_stale: bool,
1942 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
1943 self.cloud_backend
1944 .as_ref()
1945 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
1946 .catalog(provider, context_ref, allow_stale)
1947 .await
1948 }
1949
1950 async fn cloud_invoke(
1951 &self,
1952 provider: CloudProviderId,
1953 context_ref: &str,
1954 model_id: &str,
1955 request: InvokeRequest,
1956 ) -> Result<CloudEventStream, BrokerError> {
1957 self.cloud_backend
1958 .as_ref()
1959 .ok_or_else(|| BrokerError::NotConfigured(provider.to_string()))?
1960 .invoke(provider, context_ref, model_id, request)
1961 .await
1962 }
1963
1964 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
1965 let mut out = Vec::new();
1966 for descriptor in super::provider::registry().iter() {
1967 out.push(ProviderStatus {
1968 key: descriptor.id.as_str().to_string(),
1969 name: descriptor.display_name.to_string(),
1970 kind: CredentialKind::OAuth,
1971 configured: load_provider_auth(descriptor.id.as_str())
1972 .ok()
1973 .flatten()
1974 .is_some(),
1975 });
1976 }
1977 for spec in STATIC_PROVIDERS {
1978 out.push(ProviderStatus {
1979 key: spec.key.to_string(),
1980 name: spec.name.to_string(),
1981 kind: CredentialKind::StaticKey,
1982 configured: static_key_configured(spec.key),
1983 });
1984 }
1985 out.push(ProviderStatus {
1986 key: LOCAL_PROVIDER_KEY.to_string(),
1987 name: "Local endpoint".to_string(),
1988 kind: CredentialKind::LocalEndpoint,
1989 configured: true,
1990 });
1991 Ok(out)
1992 }
1993}
1994
1995pub struct RemoteBroker {
2001 http: reqwest::Client,
2002 endpoint: String,
2003 machine_token: String,
2004 cache: super::TokenCache,
2005}
2006
2007impl RemoteBroker {
2008 pub fn new(
2009 endpoint: impl Into<String>,
2010 machine_token: impl Into<String>,
2011 http: reqwest::Client,
2012 cache: super::TokenCache,
2013 ) -> Self {
2014 let endpoint = endpoint.into().trim_end_matches('/').to_string();
2015 let http = reqwest::Client::builder()
2018 .redirect(reqwest::redirect::Policy::none())
2019 .connect_timeout(Duration::from_secs(10))
2020 .timeout(Duration::from_secs(90))
2021 .build()
2022 .unwrap_or(http);
2023 Self {
2024 http,
2025 endpoint,
2026 machine_token: machine_token.into(),
2027 cache,
2028 }
2029 }
2030
2031 async fn post_proxy(&self, request: &ProxyRequest) -> Result<reqwest::Response, BrokerError> {
2032 request.validate()?;
2033 let resp = self
2034 .http
2035 .post(format!("{}/proxy", self.endpoint))
2036 .bearer_auth(&self.machine_token)
2037 .json(request)
2038 .send()
2039 .await
2040 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
2041 match resp.status().as_u16() {
2042 401 => Err(BrokerError::Unauthorized),
2043 s @ (400 | 403) => {
2044 drop(resp);
2048 Err(BrokerError::Denied(format!(
2049 "broker rejected the request (HTTP {s})"
2050 )))
2051 }
2052 _ => Ok(resp),
2053 }
2054 }
2055}
2056
2057#[async_trait]
2058impl CredentialBroker for RemoteBroker {
2059 async fn access_token(&self, provider: OAuthProviderId) -> Result<AccessToken, BrokerError> {
2060 let fetcher = super::BrokerClient::with_client(
2061 self.endpoint.clone(),
2062 self.machine_token.clone(),
2063 self.http.clone(),
2064 );
2065 let tok = super::resolve_remote(
2066 &fetcher,
2067 &self.cache,
2068 provider.as_str(),
2069 super::DEFAULT_MARGIN_MS,
2070 )
2071 .await
2072 .map_err(BrokerError::Transport)?;
2073 Ok(AccessToken {
2074 token: tok.access_token,
2075 expires: tok.expires,
2076 })
2077 }
2078
2079 async fn proxy(&self, request: ProxyRequest) -> Result<ProxyResponse, BrokerError> {
2080 let mut request = request;
2081 request.stream = false;
2082 let resp = self.post_proxy(&request).await?;
2083 if !resp.status().is_success() {
2084 return Err(BrokerError::Transport(format!(
2085 "broker proxy returned HTTP {}",
2086 resp.status()
2087 )));
2088 }
2089 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
2090 serde_json::from_str::<ProxyResponse>(&body)
2091 .map_err(|e| BrokerError::Transport(format!("invalid broker proxy response: {e}")))
2092 }
2093
2094 async fn proxy_stream(&self, request: ProxyRequest) -> Result<ProxyByteStream, BrokerError> {
2095 let mut request = request;
2096 request.stream = true;
2097 let resp = self.post_proxy(&request).await?;
2098 let status = resp.status();
2099 if !status.is_success() {
2100 drop(resp);
2102 return Err(BrokerError::Transport(format!(
2103 "broker proxy stream failed: {status}"
2104 )));
2105 }
2106 use futures::StreamExt;
2107 let stream = resp
2108 .bytes_stream()
2109 .map(|chunk| chunk.map_err(|e| BrokerError::Transport(format!("stream error: {e}"))));
2110 Ok(Box::pin(stream))
2111 }
2112
2113 async fn anthropic_usage(&self) -> Result<serde_json::Value, BrokerError> {
2114 let resp = self
2117 .http
2118 .get(format!("{}/usage", self.endpoint))
2119 .bearer_auth(&self.machine_token)
2120 .send()
2121 .await
2122 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
2123 match resp.status().as_u16() {
2124 401 => Err(BrokerError::Unauthorized),
2125 s if !(200..300).contains(&s) => {
2126 drop(resp);
2128 Err(BrokerError::Transport(format!(
2129 "broker usage returned HTTP {s}"
2130 )))
2131 }
2132 _ => {
2133 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
2134 serde_json::from_str(&body)
2135 .map_err(|e| BrokerError::Transport(format!("invalid usage response: {e}")))
2136 }
2137 }
2138 }
2139
2140 async fn cloud_catalog(
2141 &self,
2142 provider: CloudProviderId,
2143 context_ref: &str,
2144 allow_stale: bool,
2145 ) -> Result<Vec<CloudCatalogEntry>, BrokerError> {
2146 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()))?;
2147 if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
2148 return Err(BrokerError::Unauthorized);
2149 }
2150 if !resp.status().is_success() {
2151 return Err(BrokerError::Transport(format!(
2152 "cloud catalog returned HTTP {}",
2153 resp.status()
2154 )));
2155 }
2156 let body = read_body_capped(resp, MAX_CLOUD_CATALOG_BODY_BYTES).await?;
2157 let entries: Vec<CloudCatalogEntry> = serde_json::from_str(&body)
2158 .map_err(|e| BrokerError::Transport(format!("invalid cloud catalog: {e}")))?;
2159 if entries.len() > MAX_CLOUD_CATALOG_ENTRIES
2160 || entries.iter().any(|e| {
2161 e.provider != provider || !e.context_ref.starts_with("ctx-") || e.id.len() > 512
2162 })
2163 {
2164 return Err(BrokerError::Transport(
2165 "broker returned an invalid cloud catalog".into(),
2166 ));
2167 }
2168 Ok(entries)
2169 }
2170
2171 async fn cloud_invoke(
2172 &self,
2173 provider: CloudProviderId,
2174 context_ref: &str,
2175 model_id: &str,
2176 request: InvokeRequest,
2177 ) -> Result<CloudEventStream, BrokerError> {
2178 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()))?;
2179 if !resp.status().is_success() {
2180 return Err(if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
2181 BrokerError::Unauthorized
2182 } else {
2183 BrokerError::Transport(format!("cloud invoke returned HTTP {}", resp.status()))
2184 });
2185 }
2186 use futures::StreamExt;
2187 let chunks = Box::pin(resp.bytes_stream());
2188 let stream = futures::stream::unfold(
2189 (chunks, Vec::<u8>::new(), false, false),
2190 |(mut chunks, mut buffer, done, failed)| async move {
2191 if failed {
2192 return None;
2193 }
2194 loop {
2195 if let Some(end) = buffer.iter().position(|b| *b == b'\n') {
2196 let line: Vec<u8> = buffer.drain(..=end).collect();
2197 let parsed = serde_json::from_slice::<CloudEvent>(&line[..end])
2198 .map_err(|_| BrokerError::Transport("invalid cloud event".into()));
2199 let event = match parsed {
2200 Ok(CloudEvent::Done) if done => Err(BrokerError::Transport(
2201 "duplicate cloud terminal event".into(),
2202 )),
2203 Ok(CloudEvent::Done) => Ok(CloudEvent::Done),
2204 Ok(_) if done => Err(BrokerError::Transport(
2205 "cloud data followed terminal event".into(),
2206 )),
2207 other => other,
2208 };
2209 let next_done = done || matches!(event, Ok(CloudEvent::Done));
2210 let next_failed = event.is_err();
2211 return Some((event, (chunks, buffer, next_done, next_failed)));
2212 }
2213 match chunks.next().await {
2214 Some(Ok(chunk))
2215 if buffer.len() + chunk.len() <= MAX_CLOUD_STREAM_EVENT_BYTES =>
2216 {
2217 buffer.extend_from_slice(&chunk)
2218 }
2219 Some(Ok(_)) => {
2220 return Some((
2221 Err(BrokerError::Transport("cloud event exceeded limit".into())),
2222 (chunks, Vec::new(), done, true),
2223 ))
2224 }
2225 Some(Err(_)) => {
2226 return Some((
2227 Err(BrokerError::Transport("broker stream failed".into())),
2228 (chunks, Vec::new(), done, true),
2229 ))
2230 }
2231 None if buffer.is_empty() && done => return None,
2232 None if buffer.is_empty() => {
2233 return Some((
2234 Err(BrokerError::Transport(
2235 "cloud stream ended without terminal event".into(),
2236 )),
2237 (chunks, Vec::new(), done, true),
2238 ))
2239 }
2240 None => {
2241 return Some((
2242 Err(BrokerError::Transport("truncated cloud event".into())),
2243 (chunks, Vec::new(), done, true),
2244 ))
2245 }
2246 }
2247 }
2248 },
2249 );
2250 Ok(Box::pin(stream))
2251 }
2252
2253 async fn capabilities(&self) -> Result<Vec<ProviderStatus>, BrokerError> {
2254 let resp = self
2255 .http
2256 .get(format!("{}/capabilities", self.endpoint))
2257 .bearer_auth(&self.machine_token)
2258 .send()
2259 .await
2260 .map_err(|e| BrokerError::Transport(format!("broker request failed: {e}")))?;
2261 match resp.status().as_u16() {
2262 401 => Err(BrokerError::Unauthorized),
2263 s if !(200..300).contains(&s) => {
2264 Err(BrokerError::Transport(format!("broker returned HTTP {s}")))
2265 }
2266 _ => {
2267 let body = read_body_capped(resp, MAX_PROXY_RESPONSE_BYTES).await?;
2268 serde_json::from_str::<Vec<ProviderStatus>>(&body).map_err(|e| {
2269 BrokerError::Transport(format!("invalid capabilities response: {e}"))
2270 })
2271 }
2272 }
2273 }
2274}
2275
2276pub fn preflight_cloud_capability(
2285 provider: CloudProviderId,
2286 needs_tools: bool,
2287) -> Result<(), BrokerError> {
2288 if needs_tools && !provider.supports_tools() {
2289 return Err(BrokerError::UnsupportedCapability {
2290 provider: provider.to_string(),
2291 capability: "tools".into(),
2292 });
2293 }
2294 Ok(())
2295}
2296
2297pub fn broker_from_source(
2299 source: &super::CredentialSource,
2300 cache: &super::TokenCache,
2301 http: reqwest::Client,
2302) -> Arc<dyn CredentialBroker> {
2303 match source {
2304 super::CredentialSource::Local => Arc::new(LocalBroker::new(http)),
2305 super::CredentialSource::Remote {
2306 endpoint,
2307 machine_token,
2308 } => Arc::new(RemoteBroker::new(
2309 endpoint.clone(),
2310 machine_token.clone(),
2311 http,
2312 cache.clone(),
2313 )),
2314 }
2315}
2316
2317static GLOBAL_BROKER: std::sync::RwLock<Option<Arc<dyn CredentialBroker>>> =
2318 std::sync::RwLock::new(None);
2319
2320pub fn set_global_broker(broker: Arc<dyn CredentialBroker>) {
2322 *GLOBAL_BROKER.write().expect("broker registry poisoned") = Some(broker);
2323}
2324
2325pub fn global_broker() -> Arc<dyn CredentialBroker> {
2328 if let Some(b) = GLOBAL_BROKER
2329 .read()
2330 .expect("broker registry poisoned")
2331 .clone()
2332 {
2333 return b;
2334 }
2335 let default: Arc<dyn CredentialBroker> = Arc::new(LocalBroker::new(
2336 reqwest::Client::builder()
2337 .connect_timeout(std::time::Duration::from_secs(10))
2338 .timeout(std::time::Duration::from_secs(300))
2339 .build()
2340 .unwrap_or_default(),
2341 ));
2342 let mut guard = GLOBAL_BROKER.write().expect("broker registry poisoned");
2343 if let Some(b) = guard.clone() {
2344 return b;
2345 }
2346 *guard = Some(default.clone());
2347 default
2348}
2349
2350pub async fn broker_access_token(
2353 provider: OAuthProviderId,
2354 source: &super::CredentialSource,
2355 cache: &super::TokenCache,
2356 http: &reqwest::Client,
2357) -> Result<AccessToken, BrokerError> {
2358 broker_from_source(source, cache, http.clone())
2359 .access_token(provider)
2360 .await
2361}
2362
2363pub fn static_key_status_map() -> BTreeMap<String, StaticKeyStatus> {
2366 STATIC_PROVIDERS
2367 .iter()
2368 .map(|s| (s.key.to_string(), static_key_status(s.key)))
2369 .collect()
2370}
2371
2372#[cfg(test)]
2373mod tests {
2374 use super::*;
2375
2376 #[test]
2381 fn preflight_rejects_tool_requiring_cloud_routes_with_typed_error() {
2382 for provider in [
2383 CloudProviderId::AzureOpenAi,
2384 CloudProviderId::AwsBedrock,
2385 CloudProviderId::GoogleVertex,
2386 ] {
2387 let err = preflight_cloud_capability(provider, true)
2388 .expect_err("tool-requiring cloud route must fail pre-flight");
2389 assert_eq!(
2390 err,
2391 BrokerError::UnsupportedCapability {
2392 provider: provider.to_string(),
2393 capability: "tools".into(),
2394 }
2395 );
2396 assert!(err.to_string().contains("text-only"));
2397 assert_eq!(preflight_cloud_capability(provider, false), Ok(()));
2398 }
2399 }
2400
2401 async fn spawn_ndjson(body: &'static str) -> String {
2402 use axum::{body::Body, routing::post, Router};
2403 use bytes::Bytes;
2404 let app = Router::new().route(
2405 "/cloud/invoke",
2406 post(move || async move {
2407 let chunks = body
2408 .as_bytes()
2409 .chunks(3)
2410 .map(|chunk| Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(chunk)));
2411 axum::response::Response::builder()
2412 .header("content-type", "application/x-ndjson")
2413 .body(Body::from_stream(futures::stream::iter(chunks)))
2414 .unwrap()
2415 }),
2416 );
2417 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2418 let addr = listener.local_addr().unwrap();
2419 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
2420 format!("http://{addr}")
2421 }
2422
2423 async fn remote_events(body: &'static str) -> Vec<Result<CloudEvent, BrokerError>> {
2424 use futures::StreamExt;
2425 let broker = RemoteBroker::new(
2426 spawn_ndjson(body).await,
2427 "opaque-machine-token",
2428 reqwest::Client::new(),
2429 super::super::TokenCache::new(),
2430 );
2431 broker
2432 .cloud_invoke(
2433 CloudProviderId::GoogleVertex,
2434 "ctx-opaque",
2435 "google-vertex/publishers/google/models/test",
2436 InvokeRequest {
2437 messages: vec![],
2438 tools: vec![],
2439 stream: true,
2440 options: Default::default(),
2441 },
2442 )
2443 .await
2444 .unwrap()
2445 .collect()
2446 .await
2447 }
2448
2449 #[tokio::test]
2450 async fn remote_ndjson_requires_exactly_one_terminal_and_stops_after_error() {
2451 let valid =
2452 remote_events("{\"type\":\"text_delta\",\"delta\":\"hi\"}\n{\"type\":\"done\"}\n")
2453 .await;
2454 assert_eq!(valid.len(), 2);
2455 assert!(matches!(valid[1], Ok(CloudEvent::Done)));
2456
2457 for body in [
2458 "{\"type\":\"text_delta\",\"delta\":\"hi\"}\n",
2459 "{\"type\":\"done\"}\n{\"type\":\"done\"}\n",
2460 "{\"type\":\"done\"}\n{\"type\":\"text_delta\",\"delta\":\"late\"}\n",
2461 ] {
2462 let events = remote_events(body).await;
2463 assert_eq!(events.iter().filter(|event| event.is_err()).count(), 1);
2464 assert!(events.last().unwrap().is_err());
2465 }
2466 }
2467
2468 #[test]
2471 fn proxy_rejects_unlisted_same_host_paths() {
2472 for (provider, path) in [
2473 ("groq", "/v1/keys"),
2474 ("groq", "/admin"),
2475 ("openrouter", "/api/v1/auth/keys"),
2476 ("local", "/audio/speech"),
2477 ] {
2478 let req = ProxyRequest {
2479 provider: provider.into(),
2480 method: ProxyMethod::Get,
2481 path: path.into(),
2482 body: None,
2483 stream: false,
2484 body_bytes: None,
2485 };
2486 match req.validate() {
2487 Err(BrokerError::Denied(msg)) => {
2488 assert!(msg.contains("allowlist"), "got: {msg}")
2489 }
2490 other => panic!("{provider} {path} must be denied, got {other:?}"),
2491 }
2492 }
2493 for path in ["/models", "/chat/completions"] {
2495 let req = ProxyRequest {
2496 provider: "groq".into(),
2497 method: ProxyMethod::Get,
2498 path: path.into(),
2499 body: None,
2500 stream: false,
2501 body_bytes: None,
2502 };
2503 assert!(req.validate().is_ok(), "{path} must be allowed");
2504 }
2505 }
2506
2507 #[test]
2510 fn proxy_rejects_oversize_request_body() {
2511 let req = ProxyRequest {
2512 provider: "groq".into(),
2513 method: ProxyMethod::Post,
2514 path: "/chat/completions".into(),
2515 body: Some(serde_json::json!({
2516 "blob": "x".repeat(MAX_PROXY_REQUEST_BYTES + 1)
2517 })),
2518 stream: false,
2519 body_bytes: None,
2520 };
2521 match req.validate() {
2522 Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
2523 other => panic!("oversize body must be denied, got {other:?}"),
2524 }
2525 }
2526
2527 #[test]
2531 fn proxy_rejects_oversize_body_bytes_without_json_body() {
2532 let req = ProxyRequest {
2533 provider: "groq".into(),
2534 method: ProxyMethod::Post,
2535 path: "/chat/completions".into(),
2536 body: None,
2537 stream: false,
2538 body_bytes: Some(bytes::Bytes::from(vec![b'x'; MAX_PROXY_REQUEST_BYTES + 1])),
2539 };
2540 match req.validate() {
2541 Err(BrokerError::Denied(msg)) => assert!(msg.contains("byte"), "got: {msg}"),
2542 other => panic!("oversize body_bytes must be denied, got {other:?}"),
2543 }
2544 }
2545
2546 #[test]
2550 fn post_json_exact_sets_coherent_body_and_bytes() {
2551 let value = serde_json::json!({"model": "m", "stream": true});
2552 let (req, digest_bytes) =
2553 ProxyRequest::post_json_exact("groq", "/chat/completions", value.clone(), true)
2554 .expect("serializable body must construct");
2555 assert_eq!(req.method, ProxyMethod::Post);
2556 assert!(req.stream);
2557 assert_eq!(req.body.as_ref(), Some(&value));
2558 let stored = req.body_bytes.as_ref().expect("bytes must be set");
2559 assert_eq!(stored, &digest_bytes, "caller digest bytes == wire bytes");
2560 assert_eq!(
2561 serde_json::from_slice::<serde_json::Value>(stored).unwrap(),
2562 value
2563 );
2564 req.validate().expect("coherent request must validate");
2565 }
2566
2567 #[test]
2571 fn mismatched_body_bytes_rejected_by_debug_validation() {
2572 let req = ProxyRequest {
2573 provider: "groq".into(),
2574 method: ProxyMethod::Post,
2575 path: "/chat/completions".into(),
2576 body: Some(serde_json::json!({"model": "claimed"})),
2577 stream: true,
2578 body_bytes: Some(bytes::Bytes::from_static(b"{\"model\":\"actually-sent\"}")),
2579 };
2580 match req.validate() {
2581 Err(BrokerError::Denied(msg)) => {
2582 assert!(msg.contains("body_bytes"), "got: {msg}");
2583 }
2584 other => panic!("incoherent handoff must be denied, got {other:?}"),
2585 }
2586 }
2587
2588 #[test]
2589 fn proxy_request_validation_fails_closed() {
2590 let ok = ProxyRequest {
2591 provider: "groq".into(),
2592 method: ProxyMethod::Post,
2593 path: "/chat/completions".into(),
2594 body: None,
2595 stream: false,
2596 body_bytes: None,
2597 };
2598 assert!(ok.validate().is_ok());
2599
2600 let unknown = ProxyRequest {
2601 provider: "evil".into(),
2602 ..ok.clone()
2603 };
2604 assert!(matches!(
2605 unknown.validate(),
2606 Err(BrokerError::UnknownProvider(_))
2607 ));
2608
2609 let absolute = ProxyRequest {
2610 path: "https://evil.example/x".into(),
2611 ..ok.clone()
2612 };
2613 assert!(matches!(absolute.validate(), Err(BrokerError::Denied(_))));
2614
2615 let traversal = ProxyRequest {
2616 path: "/../secrets".into(),
2617 ..ok.clone()
2618 };
2619 assert!(matches!(traversal.validate(), Err(BrokerError::Denied(_))));
2620
2621 let relative = ProxyRequest {
2622 path: "chat/completions".into(),
2623 ..ok
2624 };
2625 assert!(matches!(relative.validate(), Err(BrokerError::Denied(_))));
2626 }
2627
2628 #[test]
2631 fn proxy_allows_only_pinned_kimi_code_paths() {
2632 let chat = ProxyRequest {
2633 provider: "kimi-code".into(),
2634 method: ProxyMethod::Post,
2635 path: "/chat/completions".into(),
2636 body: None,
2637 stream: true,
2638 body_bytes: None,
2639 };
2640 assert!(chat.validate().is_ok(), "kimi-code chat must be allowed");
2641
2642 for path in ["/models", "/me", "/usages"] {
2643 let get = ProxyRequest {
2644 provider: "kimi-code".into(),
2645 method: ProxyMethod::Get,
2646 path: path.into(),
2647 body: None,
2648 stream: false,
2649 body_bytes: None,
2650 };
2651 assert!(get.validate().is_ok(), "kimi-code GET {path} must pass");
2652 let post = ProxyRequest {
2654 method: ProxyMethod::Post,
2655 ..get
2656 };
2657 assert!(
2658 matches!(post.validate(), Err(BrokerError::Denied(_))),
2659 "kimi-code POST {path} must be denied"
2660 );
2661 }
2662
2663 for path in ["/feedback", "/search", "/fetch", "/completions"] {
2665 let denied = ProxyRequest {
2666 provider: "kimi-code".into(),
2667 method: ProxyMethod::Post,
2668 path: path.into(),
2669 body: None,
2670 stream: false,
2671 body_bytes: None,
2672 };
2673 assert!(
2674 matches!(denied.validate(), Err(BrokerError::Denied(_))),
2675 "kimi-code {path} must be denied"
2676 );
2677 }
2678
2679 let static_kimi = ProxyRequest {
2682 provider: "kimi".into(),
2683 method: ProxyMethod::Post,
2684 path: "/chat/completions".into(),
2685 body: None,
2686 stream: false,
2687 body_bytes: None,
2688 };
2689 assert!(static_kimi.validate().is_ok());
2690 }
2691
2692 #[test]
2696 fn proxy_rejects_oauth_providers_except_codex_catalog() {
2697 let claude = ProxyRequest {
2699 provider: "claude".into(),
2700 method: ProxyMethod::Get,
2701 path: "/models".into(),
2702 body: None,
2703 stream: false,
2704 body_bytes: None,
2705 };
2706 assert!(
2707 matches!(claude.validate(), Err(BrokerError::UnknownProvider(_))),
2708 "claude alias must not be proxyable"
2709 );
2710 let anth = ProxyRequest {
2712 provider: "anthropic".into(),
2713 method: ProxyMethod::Get,
2714 path: "/models".into(),
2715 body: None,
2716 stream: false,
2717 body_bytes: None,
2718 };
2719 assert!(
2720 matches!(anth.validate(), Err(BrokerError::Denied(_))),
2721 "anthropic /models must be denied"
2722 );
2723 let bare = ProxyRequest {
2725 provider: "openai-codex".into(),
2726 method: ProxyMethod::Get,
2727 path: "/models".into(),
2728 body: None,
2729 stream: false,
2730 body_bytes: None,
2731 };
2732 assert!(
2733 matches!(bare.validate(), Err(BrokerError::Denied(_))),
2734 "openai-codex /models without client_version must be denied"
2735 );
2736 }
2737
2738 #[test]
2741 fn proxy_allows_only_pinned_openai_codex_models_path() {
2742 let allowed = ProxyRequest {
2743 provider: "openai-codex".into(),
2744 method: ProxyMethod::Get,
2745 path: "/models?client_version=0.6.0".into(),
2746 body: None,
2747 stream: false,
2748 body_bytes: None,
2749 };
2750 assert!(
2751 allowed.validate().is_ok(),
2752 "exact codex models path must be allowed"
2753 );
2754
2755 for path in [
2756 "/models",
2757 "/models?client_version=",
2758 "/models?foo=1",
2759 "/models?client_version=0.6.0&extra=1",
2760 "/codex/responses",
2761 "/backend-api/models?client_version=0.6.0",
2762 "/v1/models?client_version=0.6.0",
2763 "/chat/completions",
2764 ] {
2765 let req = ProxyRequest {
2766 provider: "openai-codex".into(),
2767 method: ProxyMethod::Get,
2768 path: path.into(),
2769 body: None,
2770 stream: false,
2771 body_bytes: None,
2772 };
2773 assert!(
2774 matches!(req.validate(), Err(BrokerError::Denied(_))),
2775 "openai-codex path {path} must be denied"
2776 );
2777 }
2778
2779 let post = ProxyRequest {
2781 provider: "openai-codex".into(),
2782 method: ProxyMethod::Post,
2783 path: "/models?client_version=0.6.0".into(),
2784 body: None,
2785 stream: false,
2786 body_bytes: None,
2787 };
2788 assert!(
2789 matches!(post.validate(), Err(BrokerError::Denied(_))),
2790 "POST openai-codex models must be denied"
2791 );
2792 }
2793
2794 #[test]
2796 fn proxy_allows_only_pinned_github_copilot_paths() {
2797 let models = ProxyRequest {
2798 provider: "github-copilot".into(),
2799 method: ProxyMethod::Get,
2800 path: "/models".into(),
2801 body: None,
2802 stream: false,
2803 body_bytes: None,
2804 };
2805 assert!(models.validate().is_ok());
2806
2807 let chat = ProxyRequest {
2808 provider: "github-copilot".into(),
2809 method: ProxyMethod::Post,
2810 path: "/chat/completions".into(),
2811 body: None,
2812 stream: false,
2813 body_bytes: None,
2814 };
2815 assert!(chat.validate().is_ok());
2816 let responses = ProxyRequest {
2817 provider: "github-copilot".into(),
2818 method: ProxyMethod::Post,
2819 path: "/responses".into(),
2820 body: None,
2821 stream: false,
2822 body_bytes: None,
2823 };
2824 assert!(responses.validate().is_ok());
2825 for path in ["/v1/messages", "/models?x=1", "/embeddings"] {
2826 let request = ProxyRequest {
2827 provider: "github-copilot".into(),
2828 method: ProxyMethod::Post,
2829 path: path.into(),
2830 body: None,
2831 stream: false,
2832 body_bytes: None,
2833 };
2834 assert!(matches!(request.validate(), Err(BrokerError::Denied(_))));
2835 }
2836 }
2837
2838 #[test]
2840 fn proxy_allows_only_pinned_anthropic_models_path() {
2841 for path in [
2842 "/v1/models",
2843 "/v1/models?limit=100",
2844 "/v1/models?limit=100&after_id=claude-opus-4-7",
2845 ] {
2846 let req = ProxyRequest {
2847 provider: "anthropic".into(),
2848 method: ProxyMethod::Get,
2849 path: path.into(),
2850 body: None,
2851 stream: false,
2852 body_bytes: None,
2853 };
2854 assert!(req.validate().is_ok(), "{path} must be allowed");
2855 }
2856 for path in [
2857 "/v1/messages",
2858 "/models",
2859 "/v1/models?foo=1",
2860 "/v1/models?limit=",
2861 "/api/oauth/usage",
2862 "/v1/models?limit=100&evil=1",
2863 "/v1/models?limit=100&limit=50",
2864 "/v1/models?after_id=a&after_id=b",
2865 "/v1/models?limit=100&after_id=a&limit=50",
2866 ] {
2867 let req = ProxyRequest {
2868 provider: "anthropic".into(),
2869 method: ProxyMethod::Get,
2870 path: path.into(),
2871 body: None,
2872 stream: false,
2873 body_bytes: None,
2874 };
2875 assert!(
2876 matches!(req.validate(), Err(BrokerError::Denied(_))),
2877 "anthropic path {path} must be denied"
2878 );
2879 }
2880 }
2881
2882 #[test]
2883 fn anthropic_oauth_catalog_headers_are_bearer_version_beta_not_x_api_key() {
2884 let headers = anthropic_oauth_catalog_request_headers();
2885 let names: Vec<_> = headers.iter().map(|(n, _)| *n).collect();
2886 assert!(names.contains(&"anthropic-version"));
2887 assert!(names.contains(&"anthropic-beta"));
2888 assert!(names.contains(&"accept"));
2889 assert!(!names.iter().any(|n| n.eq_ignore_ascii_case("x-api-key")));
2890 assert!(!names
2891 .iter()
2892 .any(|n| n.eq_ignore_ascii_case("authorization")));
2893 let beta = headers
2894 .iter()
2895 .find(|(n, _)| *n == "anthropic-beta")
2896 .map(|(_, v)| *v)
2897 .expect("beta");
2898 assert_eq!(beta, "oauth-2025-04-20");
2899 }
2900
2901 #[test]
2902 fn anthropic_models_path_round_trip_and_rejects_duplicates() {
2903 assert!(is_allowed_anthropic_path("/v1/models"));
2904 assert!(is_allowed_anthropic_path("/v1/models?limit=100"));
2905 assert!(is_allowed_anthropic_path(
2906 "/v1/models?limit=100&after_id=claude-opus-4-7"
2907 ));
2908 assert!(!is_allowed_anthropic_path("/v1/models?limit=100&limit=50"));
2909 assert!(!is_allowed_anthropic_path(
2910 "/v1/models?after_id=a&after_id=b"
2911 ));
2912 let page0 = "/v1/models?limit=100";
2914 let page1 = "/v1/models?limit=100&after_id=model-1";
2915 assert!(is_allowed_anthropic_path(page0));
2916 assert!(is_allowed_anthropic_path(page1));
2917 }
2918
2919 #[test]
2921 fn proxy_allows_only_pinned_google_gemini_paths() {
2922 for path in [
2923 "/v1internal:loadCodeAssist",
2924 "/v1internal:onboardUser",
2925 "/v1internal:streamGenerateContent",
2926 "/v1internal:countTokens",
2927 "/v1internal/operations/op-12345",
2928 ] {
2929 let req = ProxyRequest {
2930 provider: "google-gemini".into(),
2931 method: ProxyMethod::Post,
2932 path: path.into(),
2933 body: None,
2934 stream: false,
2935 body_bytes: None,
2936 };
2937 assert!(req.validate().is_ok(), "{path} must be allowed");
2938 }
2939 for path in [
2940 "/v1internal:listExperiments",
2942 "/v1internal:fetchAdminControls",
2943 "/v1internal:setCodeAssistGlobalUserSetting",
2944 "/v1internal:generateContent",
2945 "/v2/models",
2947 "/v1beta/models",
2948 "/",
2950 "/v1internal:",
2951 ] {
2952 let req = ProxyRequest {
2953 provider: "google-gemini".into(),
2954 method: ProxyMethod::Post,
2955 path: path.into(),
2956 body: None,
2957 stream: false,
2958 body_bytes: None,
2959 };
2960 assert!(
2961 matches!(req.validate(), Err(BrokerError::Denied(_))),
2962 "{path} must be denied"
2963 );
2964 }
2965 }
2966
2967 #[test]
2968 fn access_token_type_has_no_refresh_field() {
2969 let t: AccessToken =
2972 serde_json::from_str(r#"{"token":"sk-x","expires":123,"refresh":"MUST-NOT-EXIST"}"#)
2973 .unwrap();
2974 let round = serde_json::to_value(&t).unwrap();
2975 assert_eq!(round.get("refresh"), None);
2976 assert_eq!(round["token"], "sk-x");
2977 }
2978
2979 #[test]
2980 fn broker_error_display_never_echoes_values() {
2981 let e = BrokerError::NotConfigured("groq".into());
2983 assert!(!format!("{e}").contains("sk-"));
2984 let e = BrokerError::Unauthorized;
2985 assert_eq!(format!("{e}"), "broker rejected machine auth");
2986 }
2987
2988 #[test]
2989 fn mask_key_is_short_and_lossy() {
2990 assert_eq!(mask_key("gsk-live-1234567890abcdef"), "gsk-…cdef");
2991 assert_eq!(mask_key("short"), "…");
2992 }
2993
2994 #[test]
2995 fn proxy_request_serde_roundtrip() {
2996 let req = ProxyRequest {
2997 provider: "openrouter".into(),
2998 method: ProxyMethod::Get,
2999 path: "/models".into(),
3000 body: Some(serde_json::json!({"a": 1})),
3001 stream: true,
3002 body_bytes: None,
3003 };
3004 let json = serde_json::to_string(&req).unwrap();
3005 let back: ProxyRequest = serde_json::from_str(&json).unwrap();
3006 assert_eq!(back.provider, "openrouter");
3007 assert!(back.stream);
3008 assert!(matches!(back.method, ProxyMethod::Get));
3009 }
3010
3011 async fn spawn_upstream(app: axum::Router) -> String {
3014 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3015 let addr = listener.local_addr().unwrap();
3016 tokio::spawn(async move {
3017 axum::serve(listener, app).await.unwrap();
3018 });
3019 format!("http://{addr}")
3020 }
3021
3022 #[tokio::test]
3026 async fn local_broker_streams_local_endpoint_and_applies_key() {
3027 use axum::routing::post;
3028 let seen_auth = Arc::new(std::sync::Mutex::new(String::new()));
3029 let seen = seen_auth.clone();
3030 let app = axum::Router::new().route(
3031 "/chat/completions",
3032 post(move |headers: axum::http::HeaderMap| {
3033 let seen = seen.clone();
3034 async move {
3035 *seen.lock().unwrap() = headers
3036 .get("authorization")
3037 .and_then(|v| v.to_str().ok())
3038 .unwrap_or("")
3039 .to_string();
3040 (
3041 [("content-type", "text/event-stream")],
3042 "data: {\"x\":1}\n\ndata: [DONE]\n\n",
3043 )
3044 }
3045 }),
3046 );
3047 let url = spawn_upstream(app).await;
3048 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3049 let mut stream = broker
3050 .proxy_stream(ProxyRequest {
3051 provider: LOCAL_PROVIDER_KEY.into(),
3052 method: ProxyMethod::Post,
3053 path: "/chat/completions".into(),
3054 body: Some(serde_json::json!({"model": "m"})),
3055 stream: true,
3056 body_bytes: None,
3057 })
3058 .await
3059 .expect("stream must open");
3060 use futures::StreamExt;
3061 let mut collected = Vec::new();
3062 while let Some(chunk) = stream.next().await {
3063 collected.extend_from_slice(&chunk.unwrap());
3064 }
3065 let text = String::from_utf8(collected).unwrap();
3066 assert!(text.contains("data: {\"x\":1}"));
3067 assert!(text.contains("[DONE]"));
3068 assert_eq!(&*seen_auth.lock().unwrap(), "Bearer local");
3070 }
3071
3072 #[tokio::test]
3075 async fn local_broker_stream_error_is_typed_and_keyless() {
3076 use axum::routing::post;
3077 let app = axum::Router::new().route(
3078 "/chat/completions",
3079 post(|| async {
3080 (
3081 axum::http::StatusCode::UNAUTHORIZED,
3082 "{\"error\":\"bad key\"}",
3083 )
3084 }),
3085 );
3086 let url = spawn_upstream(app).await;
3087 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3088 let err = match broker
3089 .proxy_stream(ProxyRequest {
3090 provider: LOCAL_PROVIDER_KEY.into(),
3091 method: ProxyMethod::Post,
3092 path: "/chat/completions".into(),
3093 body: None,
3094 stream: true,
3095 body_bytes: None,
3096 })
3097 .await
3098 {
3099 Err(e) => e,
3100 Ok(_) => panic!("non-2xx upstream must not yield a stream"),
3101 };
3102 let msg = format!("{err}");
3103 assert!(msg.contains("401"), "got: {msg}");
3104 assert!(
3105 !msg.to_lowercase().contains("bearer"),
3106 "no auth material in errors"
3107 );
3108 }
3109
3110 #[tokio::test]
3112 async fn local_broker_proxy_returns_status_and_body() {
3113 use axum::routing::get;
3114 let app =
3115 axum::Router::new().route("/models", get(|| async { "{\"data\":[{\"id\":\"m1\"}]}" }));
3116 let url = spawn_upstream(app).await;
3117 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3118 let resp = broker
3119 .proxy(ProxyRequest {
3120 provider: LOCAL_PROVIDER_KEY.into(),
3121 method: ProxyMethod::Get,
3122 path: "/models".into(),
3123 body: None,
3124 stream: false,
3125 body_bytes: None,
3126 })
3127 .await
3128 .unwrap();
3129 assert_eq!(resp.status, 200);
3130 assert!(resp.body.contains("m1"));
3131 }
3132
3133 #[tokio::test]
3136 async fn local_broker_proxy_rejects_oversize_response() {
3137 use axum::routing::get;
3138 let app = axum::Router::new().route("/models", get(|| async { "x".repeat(4096) }));
3139 let url = spawn_upstream(app).await;
3140 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
3141 .with_max_response_bytes(64);
3142 let err = broker
3143 .proxy(ProxyRequest {
3144 provider: LOCAL_PROVIDER_KEY.into(),
3145 method: ProxyMethod::Get,
3146 path: "/models".into(),
3147 body: None,
3148 stream: false,
3149 body_bytes: None,
3150 })
3151 .await
3152 .expect_err("oversize body must be rejected");
3153 let msg = format!("{err}");
3154 assert!(msg.contains("limit"), "got: {msg}");
3155 }
3156
3157 #[tokio::test]
3160 async fn local_broker_stream_error_is_status_only() {
3161 use axum::routing::post;
3162 let app = axum::Router::new().route(
3163 "/chat/completions",
3164 post(|| async {
3165 (
3166 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
3167 format!("\x1b[2J{}", "E".repeat(64 * 1024)),
3168 )
3169 }),
3170 );
3171 let url = spawn_upstream(app).await;
3172 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3173 let err = broker
3174 .proxy_stream(ProxyRequest {
3175 provider: LOCAL_PROVIDER_KEY.into(),
3176 method: ProxyMethod::Post,
3177 path: "/chat/completions".into(),
3178 body: None,
3179 stream: true,
3180 body_bytes: None,
3181 })
3182 .await
3183 .err()
3184 .expect("non-2xx upstream must not yield a stream");
3185 let msg = format!("{err}");
3186 assert!(
3187 msg.len() < 700,
3188 "error must be bounded, got {} bytes",
3189 msg.len()
3190 );
3191 assert!(!msg.contains('\x1b'), "control chars must be stripped");
3192 assert!(msg.contains("500"));
3193 }
3194
3195 #[tokio::test]
3198 async fn local_broker_buffered_request_times_out() {
3199 use axum::routing::get;
3200 let app = axum::Router::new().route(
3201 "/models",
3202 get(|| async {
3203 tokio::time::sleep(Duration::from_secs(30)).await;
3204 "too late"
3205 }),
3206 );
3207 let url = spawn_upstream(app).await;
3208 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url)
3209 .with_request_timeout(Duration::from_millis(200));
3210 let started = std::time::Instant::now();
3211 let err = broker
3212 .proxy(ProxyRequest {
3213 provider: LOCAL_PROVIDER_KEY.into(),
3214 method: ProxyMethod::Get,
3215 path: "/models".into(),
3216 body: None,
3217 stream: false,
3218 body_bytes: None,
3219 })
3220 .await
3221 .expect_err("hung upstream must time out");
3222 assert!(matches!(err, BrokerError::Transport(_)));
3223 assert!(
3224 started.elapsed() < Duration::from_secs(5),
3225 "timeout must be enforced promptly"
3226 );
3227 }
3228
3229 #[tokio::test]
3232 async fn global_broker_defaults_to_in_process_local() {
3233 let broker = global_broker();
3234 let caps = broker
3235 .capabilities()
3236 .await
3237 .expect("in-process broker needs no daemon");
3238 assert!(caps.iter().any(|c| c.key == LOCAL_PROVIDER_KEY));
3239 let again = global_broker();
3241 assert!(again.capabilities().await.is_ok());
3242 }
3243
3244 const HOSTILE_SENTINEL: &str = "ZX9-HOSTILE-SENTINEL-7Q";
3254
3255 fn hostile_body() -> String {
3258 format!(
3259 "{{\"error\":{{\"message\":\"ECHOED {HOSTILE_SENTINEL} \
3260 {{\\\"messages\\\":[{{\\\"content\\\":\\\"my secret prompt\\\"}}],\
3261 \\\"tools\\\":[{{\\\"input_schema\\\":{{}}}}]}}\"}}}}"
3262 )
3263 }
3264
3265 fn hostile_broker_json() -> String {
3268 format!(
3269 "{{\"error\":\"ECHOED {HOSTILE_SENTINEL} input_schema \
3270 my secret prompt\"}}"
3271 )
3272 }
3273
3274 fn assert_no_hostile_bytes(msg: &str) {
3275 assert!(
3276 !msg.contains(HOSTILE_SENTINEL),
3277 "sentinel leaked into error: {msg}"
3278 );
3279 assert!(!msg.contains("ECHOED"), "echoed body leaked: {msg}");
3280 assert!(
3281 !msg.contains("input_schema"),
3282 "request-shaped data leaked: {msg}"
3283 );
3284 assert!(
3285 !msg.contains("secret prompt"),
3286 "prompt content leaked: {msg}"
3287 );
3288 }
3289
3290 #[tokio::test]
3295 async fn local_broker_stream_error_omits_upstream_body() {
3296 use axum::routing::post;
3297 let app = axum::Router::new().route(
3298 "/chat/completions",
3299 post(|| async {
3300 (
3301 axum::http::StatusCode::TOO_MANY_REQUESTS,
3302 format!("\x1b[2J{}", hostile_body()),
3303 )
3304 }),
3305 );
3306 let url = spawn_upstream(app).await;
3307 let broker = LocalBroker::with_local_base_url(reqwest::Client::new(), url);
3308 let err = broker
3309 .proxy_stream(ProxyRequest {
3310 provider: LOCAL_PROVIDER_KEY.into(),
3311 method: ProxyMethod::Post,
3312 path: "/chat/completions".into(),
3313 body: None,
3314 stream: true,
3315 body_bytes: None,
3316 })
3317 .await
3318 .err()
3319 .expect("non-2xx upstream must not yield a stream");
3320 assert!(matches!(err, BrokerError::Transport(_)));
3321 let msg = format!("{err}");
3322 assert!(
3323 msg.contains("provider request failed: 429 Too Many Requests"),
3324 "status prefix must survive for engine classification: {msg}"
3325 );
3326 assert!(
3327 !msg.contains('\x1b'),
3328 "escapes must not reach errors: {msg}"
3329 );
3330 assert_no_hostile_bytes(&msg);
3331 }
3332
3333 #[tokio::test]
3337 #[serial_test::serial]
3338 async fn local_broker_usage_error_omits_upstream_body() {
3339 use axum::routing::get;
3340 let app = axum::Router::new().route(
3341 "/usage",
3342 get(|| async {
3343 (
3344 axum::http::StatusCode::INTERNAL_SERVER_ERROR,
3345 hostile_body(),
3346 )
3347 }),
3348 );
3349 let url = spawn_upstream(app).await;
3350
3351 let dir = tempfile::tempdir().unwrap();
3357 let base = dir.path().join(".synaps-cli");
3358 std::fs::create_dir_all(&base).unwrap();
3359 let creds = crate::auth::OAuthCredentials {
3360 auth_type: "oauth".into(),
3361 refresh: String::new(),
3362 access: "test-access-token".into(),
3363 expires: crate::epoch_millis() + 3_600_000,
3364 account_id: None,
3365 };
3366 super::super::storage::save_provider_auth_at_test_hook(
3367 &base.join("auth.json"),
3368 "anthropic",
3369 &creds,
3370 )
3371 .unwrap();
3372 let old_base = std::env::var("SYNAPS_BASE_DIR").ok();
3373 let old_home = std::env::var("HOME").ok();
3374 std::env::remove_var("SYNAPS_BASE_DIR");
3375 std::env::set_var("HOME", dir.path());
3376
3377 let broker = LocalBroker::new(reqwest::Client::new())
3378 .with_anthropic_usage_url(format!("{url}/usage"));
3379 let result = broker.anthropic_usage().await;
3380
3381 match old_home {
3382 Some(v) => std::env::set_var("HOME", v),
3383 None => std::env::remove_var("HOME"),
3384 }
3385 if let Some(v) = old_base {
3386 std::env::set_var("SYNAPS_BASE_DIR", v);
3387 }
3388
3389 let err = result.expect_err("non-2xx usage response must fail");
3390 assert!(matches!(err, BrokerError::Transport(_)));
3391 let msg = format!("{err}");
3392 assert!(
3393 msg.contains("usage request failed") && msg.contains("500"),
3394 "operation + status must survive: {msg}"
3395 );
3396 assert_no_hostile_bytes(&msg);
3397 }
3398
3399 async fn spawn_remote_broker_error(status: u16) -> RemoteBroker {
3400 use axum::routing::{get, post};
3401 let code = axum::http::StatusCode::from_u16(status).unwrap();
3402 let app = axum::Router::new()
3403 .route(
3404 "/proxy",
3405 post(move || async move { (code, hostile_broker_json()) }),
3406 )
3407 .route(
3408 "/usage",
3409 get(move || async move { (code, hostile_broker_json()) }),
3410 );
3411 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3412 let addr = listener.local_addr().unwrap();
3413 tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
3414 RemoteBroker::new(
3415 format!("http://{addr}"),
3416 "opaque-machine-token",
3417 reqwest::Client::new(),
3418 super::super::TokenCache::new(),
3419 )
3420 }
3421
3422 fn groq_request(stream: bool) -> ProxyRequest {
3423 ProxyRequest {
3424 provider: "groq".into(),
3425 method: ProxyMethod::Post,
3426 path: "/chat/completions".into(),
3427 body: None,
3428 stream,
3429 body_bytes: None,
3430 }
3431 }
3432
3433 #[tokio::test]
3436 async fn remote_broker_denied_omits_broker_error_body() {
3437 for status in [400u16, 403] {
3438 let broker = spawn_remote_broker_error(status).await;
3439 let err = broker
3440 .proxy(groq_request(false))
3441 .await
3442 .expect_err("4xx broker response must fail");
3443 assert!(
3444 matches!(err, BrokerError::Denied(_)),
3445 "typed category must survive, got {err:?}"
3446 );
3447 let msg = format!("{err}");
3448 assert!(msg.contains(&status.to_string()), "status lost: {msg}");
3449 assert_no_hostile_bytes(&msg);
3450 }
3451 }
3452
3453 #[tokio::test]
3456 async fn remote_broker_stream_error_omits_broker_body() {
3457 let broker = spawn_remote_broker_error(500).await;
3458 let err = broker
3459 .proxy_stream(groq_request(true))
3460 .await
3461 .err()
3462 .expect("non-2xx broker response must not yield a stream");
3463 assert!(matches!(err, BrokerError::Transport(_)));
3464 let msg = format!("{err}");
3465 assert!(
3466 msg.contains("broker proxy stream failed") && msg.contains("500"),
3467 "operation + status must survive: {msg}"
3468 );
3469 assert_no_hostile_bytes(&msg);
3470 }
3471
3472 #[tokio::test]
3475 async fn remote_broker_usage_error_omits_broker_body() {
3476 let broker = spawn_remote_broker_error(502).await;
3477 let err = broker
3478 .anthropic_usage()
3479 .await
3480 .expect_err("non-2xx broker response must fail");
3481 assert!(matches!(err, BrokerError::Transport(_)));
3482 let msg = format!("{err}");
3483 assert!(
3484 msg.contains("broker usage returned HTTP 502"),
3485 "operation + status must survive: {msg}"
3486 );
3487 assert_no_hostile_bytes(&msg);
3488 }
3489
3490 #[tokio::test]
3493 async fn capabilities_expose_no_secret_fields() {
3494 let broker = LocalBroker::new(reqwest::Client::new());
3495 let caps = broker.capabilities().await.unwrap();
3496 assert!(caps
3497 .iter()
3498 .any(|c| c.key == "anthropic" && c.kind == CredentialKind::OAuth));
3499 assert!(caps
3500 .iter()
3501 .any(|c| c.key == "groq" && c.kind == CredentialKind::StaticKey));
3502 let json = serde_json::to_string(&caps).unwrap();
3503 for field in ["key\":", "name\":", "kind\":", "configured\":"] {
3504 assert!(json.contains(field));
3505 }
3506 assert!(!json.contains("refresh"));
3507 assert!(!json.contains("access"));
3508 }
3509}