adk_gcp/
client.rs

1//! Authenticated, bounded HTTP transport for Google Cloud REST APIs.
2
3use crate::error::{GcpErrorContext, truncate_for_error};
4use adk_core::Result;
5use google_cloud_auth::credentials::{self, CacheableResource, Credentials};
6use reqwest::header::HeaderMap;
7use reqwest::{Client, Method, RequestBuilder, StatusCode, Url};
8use serde_json::{Map, Value};
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::sync::RwLock;
12
13const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
14const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
15const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
16const DEFAULT_AUTH_TIMEOUT: Duration = Duration::from_secs(30);
17const DEFAULT_MAX_RESPONSE_BYTES: usize = 64 * 1024 * 1024;
18const DEFAULT_API_VERSION: &str = "v1beta1";
19
20/// Builder for [`GcpHttpClient`].
21///
22/// Defaults mirror the workspace's Vertex backends: 10 s connect timeout,
23/// 120 s request timeout, 30 s credential-acquisition timeout, 64 MiB
24/// response cap, API version `v1beta1`, and the `cloud-platform` OAuth scope.
25#[derive(Debug)]
26pub struct GcpHttpClientBuilder {
27    errors: GcpErrorContext,
28    endpoint: String,
29    api_version: String,
30    credentials: Option<Credentials>,
31    scopes: Vec<String>,
32    connect_timeout: Duration,
33    request_timeout: Duration,
34    auth_timeout: Duration,
35    max_response_bytes: usize,
36}
37
38impl GcpHttpClientBuilder {
39    /// Sets the API version segment prefixed to every request path.
40    #[must_use]
41    pub fn api_version(mut self, api_version: impl Into<String>) -> Self {
42        self.api_version = api_version.into();
43        self
44    }
45
46    /// Uses explicit credentials instead of Application Default Credentials.
47    #[must_use]
48    pub fn credentials(mut self, credentials: Credentials) -> Self {
49        self.credentials = Some(credentials);
50        self
51    }
52
53    /// Overrides the OAuth scopes requested when building ADC.
54    ///
55    /// Ignored when explicit [`credentials`](Self::credentials) are set.
56    #[must_use]
57    pub fn scopes<I, S>(mut self, scopes: I) -> Self
58    where
59        I: IntoIterator<Item = S>,
60        S: Into<String>,
61    {
62        self.scopes = scopes.into_iter().map(Into::into).collect();
63        self
64    }
65
66    /// Sets the TCP connect timeout.
67    #[must_use]
68    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
69        self.connect_timeout = timeout;
70        self
71    }
72
73    /// Sets the whole-request timeout, including body streaming.
74    #[must_use]
75    pub fn request_timeout(mut self, timeout: Duration) -> Self {
76        self.request_timeout = timeout;
77        self
78    }
79
80    /// Sets the credential-header acquisition timeout.
81    #[must_use]
82    pub fn auth_timeout(mut self, timeout: Duration) -> Self {
83        self.auth_timeout = timeout;
84        self
85    }
86
87    /// Sets the maximum accepted response size in bytes.
88    ///
89    /// Responses declaring or streaming more than this are rejected. The
90    /// caller is responsible for selecting a bound appropriate for the
91    /// deployment.
92    #[must_use]
93    pub fn max_response_bytes(mut self, max_response_bytes: usize) -> Self {
94        self.max_response_bytes = max_response_bytes;
95        self
96    }
97
98    /// Validates the endpoint, resolves credentials, and builds the client.
99    ///
100    /// # Errors
101    ///
102    /// Returns an error when the endpoint is not a valid secure origin
103    /// (HTTPS, or HTTP to a loopback host, with no userinfo, path, query,
104    /// or fragment), ADC cannot be constructed, or the redirect-disabled
105    /// HTTP client cannot be built.
106    pub fn build(self) -> Result<GcpHttpClient> {
107        let errors = self.errors;
108        let base_url = validate_endpoint(&self.endpoint, &errors)?;
109
110        let credentials = match self.credentials {
111            Some(credentials) => credentials,
112            None => credentials::Builder::default().with_scopes(&self.scopes).build().map_err(
113                |error| {
114                    let error = truncate_for_error(&error.to_string());
115                    errors.unauthorized(format!(
116                        "failed to build {} ADC credentials: {error}",
117                        errors.subject(),
118                    ))
119                },
120            )?,
121        };
122
123        let http_client = Client::builder()
124            .connect_timeout(self.connect_timeout)
125            .timeout(self.request_timeout)
126            .redirect(reqwest::redirect::Policy::none())
127            .build()
128            .map_err(|error| {
129                let error = truncate_for_error(&error.without_url().to_string());
130                errors.invalid_response(format!(
131                    "failed to build bounded {} HTTP client: {error}",
132                    errors.subject(),
133                ))
134            })?;
135
136        Ok(GcpHttpClient {
137            http_client,
138            base_url,
139            api_version: self.api_version,
140            credentials,
141            auth_headers: Arc::new(RwLock::new(None)),
142            request_timeout: self.request_timeout,
143            auth_timeout: self.auth_timeout,
144            max_response_bytes: self.max_response_bytes,
145            errors,
146        })
147    }
148}
149
150/// Authenticated, bounded HTTP client for Google Cloud REST APIs.
151///
152/// Consolidates the pattern shared by every Vertex backend in the
153/// workspace: cached ADC auth headers (honoring `CacheableResource`
154/// semantics), a redirect-disabled `reqwest` client with connect and
155/// request timeouts, HTTPS-or-loopback endpoint validation, and bounded
156/// JSON response reads with typed, consumer-branded errors.
157///
158/// # Example
159///
160/// ```rust,no_run
161/// use adk_core::ErrorComponent;
162/// use adk_gcp::{GcpErrorCodes, GcpErrorContext, GcpHttpClient};
163/// use serde_json::json;
164///
165/// const CODES: GcpErrorCodes = GcpErrorCodes {
166///     invalid_input: "memory.vertex.invalid_input",
167///     unauthorized: "memory.vertex.unauthorized",
168///     forbidden: "memory.vertex.forbidden",
169///     not_found: "memory.vertex.not_found",
170///     rate_limited: "memory.vertex.rate_limited",
171///     timeout: "memory.vertex.timeout",
172///     unavailable: "memory.vertex.unavailable",
173///     credentials_unavailable: "memory.vertex.credentials_unavailable",
174///     invalid_response: "memory.vertex.invalid_response",
175///     invalid_request: "memory.vertex.invalid_request",
176///     upstream_error: "memory.vertex.upstream_error",
177///     operation_failed: "memory.vertex.operation_failed",
178/// };
179///
180/// # async fn call() -> adk_core::Result<()> {
181/// let client = GcpHttpClient::builder(
182///     GcpErrorContext::new(ErrorComponent::Memory, CODES, "vertex memory"),
183///     "https://us-central1-aiplatform.googleapis.com",
184/// )
185/// .build()?;
186///
187/// let request = client
188///     .request(reqwest::Method::POST, "projects/p/locations/l/reasoningEngines/1/memories:generate")
189///     .await?
190///     .json(&json!({ "scope": { "app_name": "app" } }));
191/// let operation = client.send_value(request).await?;
192/// # let _ = operation;
193/// # Ok(())
194/// # }
195/// ```
196#[derive(Debug)]
197pub struct GcpHttpClient {
198    http_client: Client,
199    base_url: Url,
200    api_version: String,
201    credentials: Credentials,
202    auth_headers: Arc<RwLock<Option<HeaderMap>>>,
203    request_timeout: Duration,
204    auth_timeout: Duration,
205    max_response_bytes: usize,
206    errors: GcpErrorContext,
207}
208
209impl GcpHttpClient {
210    /// Starts building a client for the given error identity and endpoint.
211    pub fn builder(errors: GcpErrorContext, endpoint: impl Into<String>) -> GcpHttpClientBuilder {
212        GcpHttpClientBuilder {
213            errors,
214            endpoint: endpoint.into(),
215            api_version: DEFAULT_API_VERSION.to_string(),
216            credentials: None,
217            scopes: vec![CLOUD_PLATFORM_SCOPE.to_string()],
218            connect_timeout: DEFAULT_CONNECT_TIMEOUT,
219            request_timeout: DEFAULT_REQUEST_TIMEOUT,
220            auth_timeout: DEFAULT_AUTH_TIMEOUT,
221            max_response_bytes: DEFAULT_MAX_RESPONSE_BYTES,
222        }
223    }
224
225    /// The error identity this client stamps on failures.
226    pub fn errors(&self) -> &GcpErrorContext {
227        &self.errors
228    }
229
230    /// The configured maximum response size in bytes.
231    pub fn max_response_bytes(&self) -> usize {
232        self.max_response_bytes
233    }
234
235    /// Overrides the maximum accepted response size after construction.
236    ///
237    /// Lets consumers expose their own post-construction bound override
238    /// without rebuilding credentials. The caller is responsible for
239    /// selecting a bound appropriate for the deployment.
240    #[must_use]
241    pub fn with_max_response_bytes(mut self, max_response_bytes: usize) -> Self {
242        self.max_response_bytes = max_response_bytes;
243        self
244    }
245
246    /// Builds the absolute URL for an API path (version prefix applied).
247    ///
248    /// # Errors
249    ///
250    /// Returns an error when the resulting URL cannot be constructed.
251    pub fn url(&self, path: &str) -> Result<Url> {
252        let mut url = self.base_url.clone();
253        url.set_path(&format!("/{}/{}", self.api_version, path.trim_start_matches('/')));
254        Ok(url)
255    }
256
257    /// Builds an authorized request for an API path.
258    ///
259    /// # Errors
260    ///
261    /// Returns an error when credential headers cannot be acquired within
262    /// the configured auth timeout.
263    pub async fn request(&self, method: Method, path: &str) -> Result<RequestBuilder> {
264        let url = self.url(path)?;
265        let request = self.http_client.request(method, url);
266        self.apply_auth(request).await
267    }
268
269    /// Applies cached credential headers to a request.
270    ///
271    /// # Errors
272    ///
273    /// Returns an error when credential headers cannot be acquired within
274    /// the configured auth timeout.
275    pub async fn apply_auth(&self, request: RequestBuilder) -> Result<RequestBuilder> {
276        let headers = self.auth_headers().await?;
277        Ok(request.headers(headers))
278    }
279
280    /// Acquires credential headers, serving cached headers on `NotModified`.
281    ///
282    /// # Errors
283    ///
284    /// Returns an error when acquisition times out, the credential source
285    /// fails, or the source reports `NotModified` before any headers were
286    /// cached.
287    pub async fn auth_headers(&self) -> Result<HeaderMap> {
288        let cacheable_headers =
289            tokio::time::timeout(self.auth_timeout, self.credentials.headers(Default::default()))
290                .await
291                .map_err(|_| {
292                    self.errors.timeout(format!(
293                        "{} credential header acquisition timed out after {} seconds",
294                        self.errors.subject(),
295                        self.auth_timeout.as_secs_f64(),
296                    ))
297                })?
298                .map_err(|error| self.errors.credentials_error(&error))?;
299
300        match cacheable_headers {
301            CacheableResource::New { data, .. } => {
302                *self.auth_headers.write().await = Some(data.clone());
303                Ok(data)
304            }
305            CacheableResource::NotModified => {
306                self.auth_headers.read().await.clone().ok_or_else(|| {
307                    self.errors.unauthorized(
308                        "google cloud credentials returned NotModified before any cached auth headers were available",
309                    )
310                })
311            }
312        }
313    }
314
315    /// Sends a request and parses the response as JSON.
316    ///
317    /// Empty or whitespace-only success bodies parse as an empty object.
318    ///
319    /// # Errors
320    ///
321    /// Returns an error when the request times out, transport fails, the
322    /// response exceeds the configured size bound, the status is not a
323    /// success, or the body is not valid JSON.
324    pub async fn send_value(&self, request: RequestBuilder) -> Result<Value> {
325        match self.send_value_internal(request, false).await? {
326            Some((value, _)) => Ok(value),
327            None => Ok(Value::Object(Map::new())),
328        }
329    }
330
331    /// Sends a request, returning the parsed JSON and the decoded body size.
332    ///
333    /// The byte count lets paginated callers enforce an aggregate response
334    /// bound across pages.
335    ///
336    /// # Errors
337    ///
338    /// Returns the same errors as [`send_value`](Self::send_value).
339    pub async fn send_value_counted(&self, request: RequestBuilder) -> Result<(Value, usize)> {
340        match self.send_value_internal(request, false).await? {
341            Some(value) => Ok(value),
342            None => Ok((Value::Object(Map::new()), 0)),
343        }
344    }
345
346    /// Sends a request, mapping `404 Not Found` to `Ok(None)`.
347    ///
348    /// # Errors
349    ///
350    /// Returns the same errors as [`send_value`](Self::send_value), except
351    /// a 404 status which becomes `Ok(None)`.
352    pub async fn send_value_allow_not_found(
353        &self,
354        request: RequestBuilder,
355    ) -> Result<Option<Value>> {
356        self.send_value_internal(request, true).await.map(|option| option.map(|(value, _)| value))
357    }
358
359    async fn send_value_internal(
360        &self,
361        request: RequestBuilder,
362        allow_not_found: bool,
363    ) -> Result<Option<(Value, usize)>> {
364        let (status, body) = tokio::time::timeout(self.request_timeout, async {
365            let mut response =
366                request.send().await.map_err(|error| self.errors.transport_error(error))?;
367            let status = response.status();
368            if let Some(declared) = response.content_length()
369                && declared > self.max_response_bytes as u64
370            {
371                return Err(self
372                    .errors
373                    .response_too_large(
374                        "response Content-Length",
375                        self.max_response_bytes,
376                        declared,
377                    )
378                    .with_upstream_status(status.as_u16()));
379            }
380            let capacity = response
381                .content_length()
382                .and_then(|length| usize::try_from(length).ok())
383                .unwrap_or_default();
384            let mut body = Vec::with_capacity(capacity);
385            while let Some(chunk) = response.chunk().await.map_err(|error| {
386                let timeout = error.is_timeout();
387                let detail = truncate_for_error(&error.without_url().to_string());
388                let error = if timeout {
389                    self.errors.timeout(format!(
390                        "{} response body timed out: {detail}",
391                        self.errors.subject(),
392                    ))
393                } else {
394                    self.errors.unavailable(format!(
395                        "failed to read {} response body: {detail}",
396                        self.errors.subject(),
397                    ))
398                };
399                error.with_upstream_status(status.as_u16())
400            })? {
401                let observed = body.len().checked_add(chunk.len()).ok_or_else(|| {
402                    self.errors
403                        .response_too_large("response body", self.max_response_bytes, u64::MAX)
404                        .with_upstream_status(status.as_u16())
405                })?;
406                if observed > self.max_response_bytes {
407                    return Err(self
408                        .errors
409                        .response_too_large(
410                            "response body",
411                            self.max_response_bytes,
412                            observed as u64,
413                        )
414                        .with_upstream_status(status.as_u16()));
415                }
416                body.extend_from_slice(&chunk);
417            }
418            Ok::<_, adk_core::AdkError>((status, body))
419        })
420        .await
421        .map_err(|_| {
422            self.errors.timeout(format!(
423                "{} request timed out after {} seconds",
424                self.errors.subject(),
425                self.request_timeout.as_secs(),
426            ))
427        })??;
428
429        if allow_not_found && status == StatusCode::NOT_FOUND {
430            return Ok(None);
431        }
432
433        if !status.is_success() {
434            let body = String::from_utf8_lossy(&body);
435            let body = if body.trim().is_empty() { "<empty body>" } else { body.as_ref() };
436            return Err(self.errors.status_error(status, body));
437        }
438
439        let body_len = body.len();
440        if body.iter().all(u8::is_ascii_whitespace) {
441            return Ok(Some((Value::Object(Map::new()), body_len)));
442        }
443
444        let value = serde_json::from_slice(&body).map_err(|error| {
445            let error = truncate_for_error(&error.to_string());
446            self.errors
447                .invalid_response(format!(
448                    "failed to parse {} response JSON: {error}",
449                    self.errors.subject(),
450                ))
451                .with_upstream_status(status.as_u16())
452        })?;
453        Ok(Some((value, body_len)))
454    }
455}
456
457/// Validates an endpoint as a bare secure origin.
458///
459/// Requires HTTPS (or HTTP to a loopback host, for tests) and rejects
460/// userinfo, path, query, and fragment components, so credentialed traffic
461/// can never be redirected or downgraded by configuration.
462fn validate_endpoint(endpoint: &str, errors: &GcpErrorContext) -> Result<Url> {
463    let url = Url::parse(endpoint).map_err(|error| {
464        let error = truncate_for_error(&error.to_string());
465        errors.invalid_input(format!("invalid {} endpoint URL: {error}", errors.subject()))
466    })?;
467    if url.scheme() != "https" && !(url.scheme() == "http" && is_loopback_host(&url)) {
468        return Err(errors.invalid_input(format!(
469            "{} endpoint must use HTTPS for secure transmission",
470            errors.subject(),
471        )));
472    }
473    if !url.username().is_empty()
474        || url.password().is_some()
475        || url.query().is_some()
476        || url.fragment().is_some()
477        || url.path() != "/"
478    {
479        return Err(errors.invalid_input(format!(
480            "{} endpoint must be an origin without userinfo, path, query, or fragment",
481            errors.subject(),
482        )));
483    }
484    Ok(url)
485}
486
487fn is_loopback_host(url: &Url) -> bool {
488    matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "::1"))
489}
490
491#[cfg(test)]
492mod tests {
493    use super::*;
494    use crate::error::GcpErrorCodes;
495    use adk_core::{ErrorCategory, ErrorComponent};
496
497    const CODES: GcpErrorCodes = GcpErrorCodes {
498        invalid_input: "session.vertex.invalid_input",
499        unauthorized: "session.vertex.unauthorized",
500        forbidden: "session.vertex.forbidden",
501        not_found: "session.vertex.not_found",
502        rate_limited: "session.vertex.rate_limited",
503        timeout: "session.vertex.timeout",
504        unavailable: "session.vertex.unavailable",
505        credentials_unavailable: "session.vertex.credentials_unavailable",
506        invalid_response: "session.vertex.invalid_response",
507        invalid_request: "session.vertex.invalid_request",
508        upstream_error: "session.vertex.upstream_error",
509        operation_failed: "session.vertex.operation_failed",
510    };
511
512    fn errors() -> GcpErrorContext {
513        GcpErrorContext::new(ErrorComponent::Session, CODES, "vertex session")
514    }
515
516    #[test]
517    fn endpoint_validation_requires_secure_origins() {
518        assert!(
519            validate_endpoint("https://us-central1-aiplatform.googleapis.com", &errors()).is_ok()
520        );
521        assert!(validate_endpoint("http://127.0.0.1:8080", &errors()).is_ok());
522        assert!(validate_endpoint("http://localhost:8080", &errors()).is_ok());
523
524        let rejected = [
525            "http://example.com",
526            "ftp://example.com",
527            "https://user:pass@example.com",
528            "https://example.com/path",
529            "https://example.com?query=1",
530            "https://example.com#fragment",
531            "not a url",
532        ];
533        for endpoint in rejected {
534            let error = validate_endpoint(endpoint, &errors()).unwrap_err();
535            assert_eq!(error.category, ErrorCategory::InvalidInput, "accepted {endpoint:?}");
536        }
537    }
538
539    #[tokio::test]
540    async fn url_prefixes_the_api_version() {
541        let client = GcpHttpClient::builder(errors(), "https://example.googleapis.com")
542            .api_version("v1")
543            .credentials(
544                google_cloud_auth::credentials::api_key_credentials::Builder::new("k").build(),
545            )
546            .build()
547            .unwrap();
548        let url = client.url("projects/p/locations/l/reasoningEngines/1").unwrap();
549        assert_eq!(
550            url.as_str(),
551            "https://example.googleapis.com/v1/projects/p/locations/l/reasoningEngines/1",
552        );
553    }
554}
adk_gcp/client.rs

adk_gcp/
client.rs
