adk_gcp/
lro.rs

1//! Long-running-operation polling with identity pinning.
2
3use crate::client::GcpHttpClient;
4use crate::error::truncate_for_error;
5use crate::resource::is_scoped_resource_name;
6use adk_core::Result;
7use reqwest::Method;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::time::Duration;
11use tokio::time::Instant;
12
13const DEFAULT_POLL_TIMEOUT: Duration = Duration::from_secs(120);
14const DEFAULT_INITIAL_DELAY: Duration = Duration::from_millis(100);
15const DEFAULT_MAX_DELAY: Duration = Duration::from_secs(2);
16
17/// A `google.longrunning.Operation` wire value.
18#[derive(Debug, Clone, Deserialize, Serialize)]
19pub struct Operation {
20    /// The server-assigned operation resource name.
21    pub name: String,
22    /// Whether the operation has reached a terminal state.
23    #[serde(default)]
24    pub done: bool,
25    /// The error result, set only on failed terminal operations.
26    #[serde(default)]
27    pub error: Option<OperationError>,
28    /// The success result, set only on completed terminal operations.
29    #[serde(default)]
30    pub response: Option<Value>,
31}
32
33/// The `google.rpc.Status` error carried by a failed operation.
34#[derive(Debug, Clone, Deserialize, Serialize)]
35pub struct OperationError {
36    /// The gRPC status code.
37    #[serde(default)]
38    pub code: i64,
39    /// The developer-facing error message.
40    #[serde(default)]
41    pub message: String,
42}
43
44/// Polls long-running operations to completion with capped backoff.
45///
46/// The poller pins the operation identity — a poll can never silently
47/// follow a different operation — and validates the operation name against
48/// the caller's project and location, so a compromised or buggy server
49/// cannot redirect polling elsewhere.
50///
51/// # Example
52///
53/// ```rust,no_run
54/// use adk_gcp::LroPoller;
55/// use std::time::Duration;
56///
57/// // Engine provisioning takes minutes; the deploy client widens the deadline.
58/// let poller = LroPoller::new()
59///     .with_poll_timeout(Duration::from_secs(900))
60///     .with_max_delay(Duration::from_secs(10));
61/// # let _ = poller;
62/// ```
63#[derive(Debug, Clone)]
64pub struct LroPoller {
65    poll_timeout: Duration,
66    initial_delay: Duration,
67    max_delay: Duration,
68}
69
70impl Default for LroPoller {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl LroPoller {
77    /// Creates a poller with the workspace defaults: 120 s deadline,
78    /// 100 ms initial delay, 2 s delay cap.
79    pub fn new() -> Self {
80        Self {
81            poll_timeout: DEFAULT_POLL_TIMEOUT,
82            initial_delay: DEFAULT_INITIAL_DELAY,
83            max_delay: DEFAULT_MAX_DELAY,
84        }
85    }
86
87    /// Sets the overall completion deadline.
88    #[must_use]
89    pub fn with_poll_timeout(mut self, timeout: Duration) -> Self {
90        self.poll_timeout = timeout;
91        self
92    }
93
94    /// Sets the delay before the first poll.
95    #[must_use]
96    pub fn with_initial_delay(mut self, delay: Duration) -> Self {
97        self.initial_delay = delay;
98        self
99    }
100
101    /// Sets the backoff delay cap.
102    #[must_use]
103    pub fn with_max_delay(mut self, delay: Duration) -> Self {
104        self.max_delay = delay;
105        self
106    }
107
108    /// Polls the operation in `initial` to completion.
109    ///
110    /// `operation_kind` names the operation in error messages (e.g.
111    /// `"memories generate"`). When `require_response` is set, a completed
112    /// operation without a `response` payload is an error.
113    ///
114    /// # Errors
115    ///
116    /// Returns an error when the initial value is not a valid operation,
117    /// the operation name is outside `project_id`/`location`, a poll
118    /// changes the operation identity, the operation completes with an
119    /// error result, or the deadline passes before completion.
120    pub async fn wait_for_operation(
121        &self,
122        client: &GcpHttpClient,
123        initial: Value,
124        operation_kind: &str,
125        require_response: bool,
126        project_id: &str,
127        location: &str,
128    ) -> Result<Option<Value>> {
129        let errors = client.errors();
130        let mut operation = parse_operation(client, initial, operation_kind)?;
131        if !is_scoped_resource_name(&operation.name, project_id, location) {
132            let name = truncate_for_error(&operation.name);
133            return Err(errors.invalid_response(format!(
134                "{} operation name '{name}' does not belong to projects/{project_id}/locations/{location}",
135                errors.subject(),
136            )));
137        }
138        let operation_name = operation.name.clone();
139        let deadline = Instant::now() + self.poll_timeout;
140        let mut delay = self.initial_delay;
141
142        loop {
143            if operation.done {
144                if let Some(error) = operation.error {
145                    return Err(errors.operation_error(
146                        operation_kind,
147                        &operation_name,
148                        error.code,
149                        &error.message,
150                    ));
151                }
152                if require_response && operation.response.is_none() {
153                    return Err(errors.invalid_response(format!(
154                        "{} {operation_kind} operation '{operation_name}' completed without a response; retry the request and inspect the operation in Google Cloud",
155                        errors.subject(),
156                    )));
157                }
158                return Ok(operation.response);
159            }
160            if Instant::now() >= deadline {
161                return Err(self.deadline_error(client, operation_kind, &operation_name));
162            }
163            let remaining = deadline.saturating_duration_since(Instant::now());
164            let poll = async {
165                let request = client.request(Method::GET, &operation_name).await?;
166                client.send_value(request).await
167            };
168            let value = tokio::time::timeout(remaining, poll)
169                .await
170                .map_err(|_| self.deadline_error(client, operation_kind, &operation_name))??;
171            let next = parse_operation(client, value, operation_kind)?;
172            if next.name != operation_name {
173                return Err(errors.invalid_response(format!(
174                    "{} {operation_kind} poll changed operation identity from '{operation_name}' to '{}'; refusing to follow a different operation",
175                    errors.subject(),
176                    truncate_for_error(&next.name),
177                )));
178            }
179            operation = next;
180            if !operation.done {
181                let remaining = deadline.saturating_duration_since(Instant::now());
182                if remaining.is_zero() {
183                    continue;
184                }
185                tokio::time::sleep(delay.min(remaining)).await;
186                delay = delay.saturating_mul(2).min(self.max_delay);
187            }
188        }
189    }
190
191    fn deadline_error(
192        &self,
193        client: &GcpHttpClient,
194        operation_kind: &str,
195        operation_name: &str,
196    ) -> adk_core::AdkError {
197        let errors = client.errors();
198        errors.timeout(format!(
199            "{} {operation_kind} operation '{operation_name}' did not complete within {} seconds; inspect the operation in Google Cloud before retrying",
200            errors.subject(),
201            self.poll_timeout.as_secs_f64(),
202        ))
203    }
204}
205
206/// Parses an operation wire value, rejecting inconsistent shapes.
207fn parse_operation(
208    client: &GcpHttpClient,
209    value: Value,
210    operation_kind: &str,
211) -> Result<Operation> {
212    let errors = client.errors();
213    let operation: Operation = serde_json::from_value(value).map_err(|error| {
214        let error = truncate_for_error(&error.to_string());
215        errors.invalid_response(format!(
216            "failed to parse {} {operation_kind} operation: {error}",
217            errors.subject(),
218        ))
219    })?;
220    if operation.name.trim().is_empty() {
221        return Err(errors.invalid_response(format!(
222            "{} {operation_kind} response did not contain an operation name",
223            errors.subject(),
224        )));
225    }
226    if operation.error.is_some() && operation.response.is_some() {
227        return Err(errors.invalid_response(format!(
228            "{} {operation_kind} operation '{}' contains both error and response results",
229            errors.subject(),
230            truncate_for_error(&operation.name),
231        )));
232    }
233    if !operation.done && (operation.error.is_some() || operation.response.is_some()) {
234        return Err(errors.invalid_response(format!(
235            "{} {operation_kind} operation '{}' contains a terminal result while done is false",
236            errors.subject(),
237            truncate_for_error(&operation.name),
238        )));
239    }
240    Ok(operation)
241}
adk_gcp/lro.rs

adk_gcp/
lro.rs
