1use std::sync::Arc;
14use std::time::Duration;
15
16use crate::ZaiResult;
17use crate::client::endpoint::EndpointConfig;
18use crate::client::secret::ApiSecret;
19
20#[derive(Clone)]
26pub struct ZaiClient {
27 inner: Arc<ClientInner>,
28}
29
30pub(crate) struct ClientInner {
32 pub(crate) endpoints: EndpointConfig,
34 pub(crate) transport: HttpTransportConfig,
36 pub(crate) sender: crate::client::transport::Transport,
38}
39
40impl ZaiClient {
41 pub fn builder(api_key: impl Into<String>) -> ZaiClientBuilder {
43 ZaiClientBuilder {
44 api_key: api_key.into(),
45 endpoints: EndpointConfig::builder(),
46 transport: HttpTransportConfig::default(),
47 allow_insecure: false,
48 }
49 }
50
51 pub fn from_env() -> ZaiResult<Self> {
53 let key = std::env::var("ZHIPU_API_KEY").map_err(|_| crate::ZaiError::ApiError {
54 code: crate::client::error::codes::SDK_CONFIG,
55 message: "ZHIPU_API_KEY environment variable not set".to_string(),
56 })?;
57 if key.trim().is_empty() {
58 return Err(crate::ZaiError::ApiError {
59 code: crate::client::error::codes::SDK_CONFIG,
60 message: "ZHIPU_API_KEY environment variable is empty".to_string(),
61 });
62 }
63 Self::builder(key).build()
64 }
65
66 pub fn endpoints(&self) -> &EndpointConfig {
68 &self.inner.endpoints
69 }
70
71 pub fn transport(&self) -> &HttpTransportConfig {
73 &self.inner.transport
74 }
75
76 pub(crate) async fn send_json<B, R>(
78 &self,
79 method: &'static str,
80 url: String,
81 body: &B,
82 ) -> ZaiResult<R>
83 where
84 B: serde::Serialize + ?Sized,
85 R: serde::de::DeserializeOwned,
86 {
87 let bytes = bytes::Bytes::from(serde_json::to_vec(body).map_err(crate::ZaiError::from)?);
88 let request = crate::client::transport::request::PreparedRequest {
89 method,
90 url,
91 body: crate::client::transport::request::BodyKind::Bytes(&bytes),
92 retry_safety: crate::client::transport::retry::RetrySafety::for_method(method),
93 retry_override: None,
94 response_mode: crate::client::transport::request::ResponseMode::Json,
95 route_template: "typed-api-request",
96 };
97 self.inner.sender.send(&request).await?.json()
98 }
99
100 pub(crate) async fn send_empty<R>(&self, method: &'static str, url: String) -> ZaiResult<R>
102 where
103 R: serde::de::DeserializeOwned,
104 {
105 let request = crate::client::transport::request::PreparedRequest {
106 method,
107 url,
108 body: crate::client::transport::request::BodyKind::None,
109 retry_safety: crate::client::transport::retry::RetrySafety::for_method(method),
110 retry_override: None,
111 response_mode: crate::client::transport::request::ResponseMode::Json,
112 route_template: "typed-api-request",
113 };
114 self.inner.sender.send(&request).await?.json()
115 }
116
117 pub(crate) async fn send_json_bytes<B: serde::Serialize + ?Sized>(
119 &self,
120 method: &'static str,
121 url: String,
122 body: &B,
123 ) -> ZaiResult<bytes::Bytes> {
124 let bytes = bytes::Bytes::from(serde_json::to_vec(body).map_err(crate::ZaiError::from)?);
125 let request = crate::client::transport::request::PreparedRequest {
126 method,
127 url,
128 body: crate::client::transport::request::BodyKind::Bytes(&bytes),
129 retry_safety: crate::client::transport::retry::RetrySafety::for_method(method),
130 retry_override: None,
131 response_mode: crate::client::transport::request::ResponseMode::Audio,
132 route_template: "binary-api-request",
133 };
134 self.inner.sender.send(&request).await?.bytes()
135 }
136
137 pub(crate) async fn send_empty_bytes(
139 &self,
140 method: &'static str,
141 url: String,
142 ) -> ZaiResult<bytes::Bytes> {
143 let request = crate::client::transport::request::PreparedRequest {
144 method,
145 url,
146 body: crate::client::transport::request::BodyKind::None,
147 retry_safety: crate::client::transport::retry::RetrySafety::for_method(method),
148 retry_override: None,
149 response_mode: crate::client::transport::request::ResponseMode::File,
150 route_template: "binary-api-request",
151 };
152 self.inner.sender.send(&request).await?.bytes()
153 }
154
155 pub(crate) async fn send_multipart<R: serde::de::DeserializeOwned>(
157 &self,
158 method: &'static str,
159 url: String,
160 factory: &crate::client::transport::multipart::MultipartBodyFactory,
161 ) -> ZaiResult<R> {
162 let request = crate::client::transport::request::PreparedRequest {
163 method,
164 url,
165 body: crate::client::transport::request::BodyKind::Multipart(factory),
166 retry_safety: crate::client::transport::retry::RetrySafety::for_method(method),
167 retry_override: None,
168 response_mode: crate::client::transport::request::ResponseMode::Json,
169 route_template: "multipart-api-request",
170 };
171 self.inner.sender.send(&request).await?.json()
172 }
173
174 pub(crate) async fn send_sse_json<B: serde::Serialize + ?Sized>(
178 &self,
179 method: &'static str,
180 url: String,
181 body: &B,
182 ) -> ZaiResult<crate::client::transport::SseByteStream> {
183 let bytes = bytes::Bytes::from(serde_json::to_vec(body).map_err(crate::ZaiError::from)?);
184 let request = crate::client::transport::request::PreparedRequest {
185 method,
186 url,
187 body: crate::client::transport::request::BodyKind::Bytes(&bytes),
188 retry_safety: crate::client::transport::retry::RetrySafety::NonIdempotent,
189 retry_override: None,
190 response_mode: crate::client::transport::request::ResponseMode::Json,
191 route_template: "sse-api-request",
192 };
193 self.inner.sender.send_sse(&request).await
194 }
195
196 pub(crate) async fn send_sse_multipart(
198 &self,
199 method: &'static str,
200 url: String,
201 factory: &crate::client::transport::multipart::MultipartBodyFactory,
202 ) -> ZaiResult<crate::client::transport::SseByteStream> {
203 let request = crate::client::transport::request::PreparedRequest {
204 method,
205 url,
206 body: crate::client::transport::request::BodyKind::Multipart(factory),
207 retry_safety: crate::client::transport::retry::RetrySafety::NonIdempotent,
208 retry_override: None,
209 response_mode: crate::client::transport::request::ResponseMode::Json,
210 route_template: "multipart-sse-api-request",
211 };
212 self.inner.sender.send_sse(&request).await
213 }
214}
215
216impl std::fmt::Debug for ZaiClient {
217 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
218 f.debug_struct("ZaiClient")
220 .field("credentials", &"[REDACTED]")
221 .field("endpoints", &self.inner.endpoints)
222 .field("transport", &self.inner.transport)
223 .finish_non_exhaustive()
224 }
225}
226
227pub struct ZaiClientBuilder {
229 api_key: String,
230 endpoints: crate::client::endpoint::EndpointConfigBuilder,
231 transport: HttpTransportConfig,
232 allow_insecure: bool,
233}
234
235impl ZaiClientBuilder {
236 pub fn endpoint(
241 mut self,
242 family: crate::client::endpoint::ApiFamily,
243 base: impl Into<String>,
244 ) -> Self {
245 use crate::client::endpoint::ApiFamily::*;
246 match family {
247 PaasV4 => self.endpoints = self.endpoints.paas_v4(base),
248 CodingPaasV4 => self.endpoints = self.endpoints.coding_paas_v4(base),
249 AgentV1 => self.endpoints = self.endpoints.agent_v1(base),
250 LlmApplication | ApplicationV2 | ApplicationV3 => {
251 self.endpoints = self.endpoints.llm_application(base)
252 },
253 Zrag => self.endpoints = self.endpoints.zrag(base),
254 Monitor => self.endpoints = self.endpoints.monitor(base),
255 Realtime => self.endpoints = self.endpoints.realtime(base),
256 }
257 self
258 }
259
260 pub fn transport(mut self, transport: HttpTransportConfig) -> Self {
262 self.transport = transport;
263 self
264 }
265
266 pub fn allow_insecure_transport(mut self, allow: bool) -> Self {
271 self.allow_insecure = allow;
272 self
273 }
274
275 pub fn build(self) -> ZaiResult<ZaiClient> {
277 if self.api_key.trim().is_empty() {
278 return Err(crate::ZaiError::ApiError {
279 code: crate::client::error::codes::SDK_CONFIG,
280 message: "ZaiClient requires a non-empty api_key".to_string(),
281 });
282 }
283 if self.api_key.trim() != self.api_key {
284 return Err(crate::ZaiError::ApiError {
285 code: crate::client::error::codes::SDK_CONFIG,
286 message: "ZaiClient api_key must not contain surrounding whitespace".to_string(),
287 });
288 }
289 if !self.api_key.bytes().all(|byte| byte.is_ascii_graphic()) {
290 return Err(crate::ZaiError::ApiError {
291 code: crate::client::error::codes::SDK_CONFIG,
292 message: "ZaiClient api_key must contain printable ASCII without whitespace"
293 .to_string(),
294 });
295 }
296 self.transport.validate()?;
297 let endpoints = self.endpoints.build(self.allow_insecure)?;
298 let reqwest = build_reqwest_client(&self.transport)?;
299 let sender = crate::client::transport::Transport::new(
300 reqwest,
301 ApiSecret::new(self.api_key),
302 &self.transport,
303 );
304 let inner = Arc::new(ClientInner {
305 endpoints,
306 transport: self.transport,
307 sender,
308 });
309 Ok(ZaiClient { inner })
310 }
311}
312
313fn build_reqwest_client(transport: &HttpTransportConfig) -> ZaiResult<reqwest::Client> {
318 let mut builder = reqwest::Client::builder()
319 .redirect(reqwest::redirect::Policy::none())
320 .pool_max_idle_per_host(8)
321 .pool_idle_timeout(Some(Duration::from_secs(90)))
322 .tcp_keepalive(Some(Duration::from_secs(60)))
323 .connect_timeout(transport.connect_timeout);
324 builder = builder.gzip(transport.enable_compression);
328 builder.build().map_err(crate::ZaiError::from)
329}
330
331const ALLOWED_HEADER_NAMES: &[&str] = &["Accept-Language", "X-Correlation-ID", "X-Test-Client"];
335
336#[derive(Clone)]
338pub struct AdditionalHeader {
339 name: &'static str,
340 value: String,
341}
342
343impl std::fmt::Debug for AdditionalHeader {
344 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345 formatter
346 .debug_struct("AdditionalHeader")
347 .field("name", &self.name)
348 .field("value", &"[REDACTED]")
349 .finish()
350 }
351}
352
353impl AdditionalHeader {
354 pub fn new(name: &str, value: &str) -> ZaiResult<Self> {
357 let Some(static_name) = ALLOWED_HEADER_NAMES
358 .iter()
359 .copied()
360 .find(|candidate| candidate.eq_ignore_ascii_case(name))
361 else {
362 return Err(crate::ZaiError::ApiError {
363 code: crate::client::error::codes::SDK_CONFIG,
364 message: format!("header name {name:?} is not allow-listed"),
365 });
366 };
367 if value.len() > 1024 {
368 return Err(crate::ZaiError::ApiError {
369 code: crate::client::error::codes::SDK_CONFIG,
370 message: "additional header value exceeds 1024 bytes".to_string(),
371 });
372 }
373 if !value.bytes().all(|byte| (0x20..=0x7e).contains(&byte)) {
374 return Err(crate::ZaiError::ApiError {
375 code: crate::client::error::codes::SDK_CONFIG,
376 message: "additional header value must contain printable ASCII only".to_string(),
377 });
378 }
379 Ok(Self {
380 name: static_name,
381 value: value.to_string(),
382 })
383 }
384
385 pub fn name(&self) -> &'static str {
387 self.name
388 }
389 pub fn value(&self) -> &str {
391 &self.value
392 }
393}
394
395#[derive(Debug, Clone, Copy, PartialEq, Eq)]
399pub enum RetryOverride {
400 AssumeIdempotent,
403}
404
405#[derive(Debug, Clone)]
412pub struct HttpTransportConfig {
413 pub connect_timeout: Duration,
415 pub request_timeout: Duration,
417 pub enable_compression: bool,
419 pub max_attempts: u8,
421 pub additional_headers: Vec<AdditionalHeader>,
423}
424
425impl Default for HttpTransportConfig {
426 fn default() -> Self {
427 Self {
428 connect_timeout: Duration::from_secs(10),
429 request_timeout: Duration::from_secs(60),
430 enable_compression: true,
431 max_attempts: 3,
432 additional_headers: Vec::new(),
433 }
434 }
435}
436
437impl HttpTransportConfig {
438 pub fn builder() -> HttpTransportConfigBuilder {
440 HttpTransportConfigBuilder {
441 config: Self::default(),
442 }
443 }
444
445 pub fn validate(&self) -> ZaiResult<()> {
448 if self.connect_timeout.is_zero() || self.connect_timeout > Duration::from_secs(10) {
449 return Err(crate::ZaiError::ApiError {
450 code: crate::client::error::codes::SDK_CONFIG,
451 message: "connect_timeout must be in 1ns..=10s".to_string(),
452 });
453 }
454 if self.request_timeout.is_zero() || self.request_timeout > Duration::from_secs(60) {
455 return Err(crate::ZaiError::ApiError {
456 code: crate::client::error::codes::SDK_CONFIG,
457 message: "request_timeout must be in 1ns..=60s".to_string(),
458 });
459 }
460 if !(1..=3).contains(&self.max_attempts) {
461 return Err(crate::ZaiError::ApiError {
462 code: crate::client::error::codes::SDK_CONFIG,
463 message: "max_attempts must be 1, 2 or 3".to_string(),
464 });
465 }
466 let mut names = std::collections::HashSet::with_capacity(self.additional_headers.len());
467 for header in &self.additional_headers {
468 if !names.insert(header.name()) {
469 return Err(crate::ZaiError::ApiError {
470 code: crate::client::error::codes::SDK_CONFIG,
471 message: format!(
472 "additional header {:?} must not be configured more than once",
473 header.name()
474 ),
475 });
476 }
477 }
478 Ok(())
479 }
480
481 pub fn with_request_timeout(mut self, d: Duration) -> ZaiResult<Self> {
484 if d.is_zero() || d > Duration::from_secs(60) {
485 return Err(crate::ZaiError::ApiError {
486 code: crate::client::error::codes::SDK_CONFIG,
487 message: "request_timeout must be in 1ns..=60s".to_string(),
488 });
489 }
490 self.request_timeout = d;
491 Ok(self)
492 }
493
494 pub fn with_connect_timeout(mut self, d: Duration) -> ZaiResult<Self> {
496 if d.is_zero() || d > Duration::from_secs(10) {
497 return Err(crate::ZaiError::ApiError {
498 code: crate::client::error::codes::SDK_CONFIG,
499 message: "connect_timeout must be in 1ns..=10s".to_string(),
500 });
501 }
502 self.connect_timeout = d;
503 Ok(self)
504 }
505
506 pub fn with_max_attempts(mut self, n: u8) -> ZaiResult<Self> {
508 if n == 0 || n > 3 {
509 return Err(crate::ZaiError::ApiError {
510 code: crate::client::error::codes::SDK_CONFIG,
511 message: "max_attempts must be 1, 2 or 3".to_string(),
512 });
513 }
514 self.max_attempts = n;
515 Ok(self)
516 }
517
518 pub fn with_additional_header(mut self, header: AdditionalHeader) -> Self {
520 self.additional_headers.push(header);
521 self
522 }
523}
524
525#[derive(Debug, Clone)]
527pub struct HttpTransportConfigBuilder {
528 config: HttpTransportConfig,
529}
530
531impl HttpTransportConfigBuilder {
532 pub fn request_timeout(mut self, d: Duration) -> ZaiResult<Self> {
534 self.config = self.config.with_request_timeout(d)?;
535 Ok(self)
536 }
537 pub fn connect_timeout(mut self, d: Duration) -> ZaiResult<Self> {
539 self.config = self.config.with_connect_timeout(d)?;
540 Ok(self)
541 }
542 pub fn max_attempts(mut self, n: u8) -> ZaiResult<Self> {
544 self.config = self.config.with_max_attempts(n)?;
545 Ok(self)
546 }
547 pub fn additional_header(mut self, header: AdditionalHeader) -> Self {
549 self.config.additional_headers.push(header);
550 self
551 }
552 pub fn build(self) -> HttpTransportConfig {
554 self.config
555 }
556}
557
558#[cfg(test)]
559mod tests {
560 use super::*;
561
562 #[test]
563 fn builder_rejects_blank_key() {
564 assert!(ZaiClient::builder(" ").build().is_err());
565 assert!(ZaiClient::builder("").build().is_err());
566 }
567
568 #[test]
569 fn builder_rejects_keys_that_cannot_be_safe_header_credentials() {
570 assert!(ZaiClient::builder("abc.def\nghi").build().is_err());
571 assert!(ZaiClient::builder("abc.def ghi").build().is_err());
572 assert!(ZaiClient::builder("密钥.abcdefghij").build().is_err());
573 }
574
575 #[test]
576 fn clone_shares_inner_no_secret_leak() {
577 let c = ZaiClient::builder("abcdefghij.0123456789abcdef")
578 .build()
579 .unwrap();
580 let c2 = c.clone();
581 let dbg = format!("{c2:?}");
583 assert!(dbg.contains("[REDACTED]"));
584 assert!(!dbg.contains("abcdefghij"));
585 }
586
587 #[test]
588 fn additional_header_allow_list() {
589 assert!(AdditionalHeader::new("X-Test-Client", "preserved").is_ok());
590 assert_eq!(
591 AdditionalHeader::new("x-test-client", "preserved")
592 .unwrap()
593 .name(),
594 "X-Test-Client"
595 );
596 assert!(AdditionalHeader::new("Authorization", "nope").is_err());
597 assert!(AdditionalHeader::new("Cookie", "nope").is_err());
598 assert!(AdditionalHeader::new("Proxy-Authorization", "nope").is_err());
599 }
600
601 #[test]
602 fn additional_header_value_limits() {
603 let long = "x".repeat(1025);
604 assert!(AdditionalHeader::new("X-Test-Client", &long).is_err());
605 assert!(AdditionalHeader::new("X-Test-Client", "ok\0bad").is_err());
606 assert!(AdditionalHeader::new("X-Test-Client", "非 ASCII").is_err());
607 }
608
609 #[test]
610 fn transport_only_tightens() {
611 assert!(
613 HttpTransportConfig::default()
614 .with_request_timeout(Duration::from_secs(120))
615 .is_err()
616 );
617 assert!(
619 HttpTransportConfig::default()
620 .with_request_timeout(Duration::from_secs(5))
621 .is_ok()
622 );
623 assert!(HttpTransportConfig::default().with_max_attempts(0).is_err());
625 assert!(HttpTransportConfig::default().with_max_attempts(4).is_err());
626 assert!(HttpTransportConfig::default().with_max_attempts(2).is_ok());
627 assert!(
628 HttpTransportConfig::default()
629 .with_request_timeout(Duration::ZERO)
630 .is_err()
631 );
632 }
633
634 #[test]
635 fn client_build_validates_direct_transport_fields() {
636 let invalid = HttpTransportConfig {
637 max_attempts: 0,
638 ..HttpTransportConfig::default()
639 };
640 assert!(
641 ZaiClient::builder("abcdefghij.0123456789abcdef")
642 .transport(invalid)
643 .build()
644 .is_err()
645 );
646 assert!(
647 ZaiClient::builder(" abcdefghij.0123456789abcdef ")
648 .build()
649 .is_err()
650 );
651
652 let duplicate_headers = HttpTransportConfig::default()
653 .with_additional_header(AdditionalHeader::new("X-Test-Client", "a").unwrap())
654 .with_additional_header(AdditionalHeader::new("x-test-client", "b").unwrap());
655 assert!(duplicate_headers.validate().is_err());
656 }
657}