gemini_rust/file_search/
operation_handle.rs

1use std::sync::Arc;
2use std::time::{Duration, Instant};
3use tracing::instrument;
4
5use crate::client::{Error, GeminiClient};
6use crate::file_search::model::{Operation, OperationResult};
7
8/// A handle for monitoring long-running file upload/import operations.
9///
10/// Operations represent asynchronous processing tasks like file chunking,
11/// embedding generation, and indexing. Use [`wait_until_done`](Self::wait_until_done)
12/// to poll until completion.
13#[derive(Debug, Clone)]
14pub struct OperationHandle {
15    client: Arc<GeminiClient>,
16    operation: Operation,
17}
18
19impl OperationHandle {
20    pub fn new(client: Arc<GeminiClient>, operation: Operation) -> Self {
21        Self { client, operation }
22    }
23
24    pub fn name(&self) -> &str {
25        &self.operation.name
26    }
27
28    pub fn is_done(&self) -> bool {
29        self.operation.done.unwrap_or(false)
30    }
31
32    pub fn result(&self) -> Option<&OperationResult> {
33        self.operation.result.as_ref()
34    }
35
36    #[instrument(skip_all, fields(operation.name = %self.operation.name))]
37    pub async fn refresh(&mut self) -> Result<(), Error> {
38        self.operation = self.client.get_operation(&self.operation.name).await?;
39        Ok(())
40    }
41
42    #[instrument(skip_all, fields(
43        operation.name = %self.operation.name,
44        poll.interval.secs = interval.as_secs(),
45        timeout.secs = timeout.as_ref().map(|d| d.as_secs()),
46    ))]
47    pub async fn wait_until_done(
48        &mut self,
49        interval: Duration,
50        timeout: Option<Duration>,
51    ) -> Result<(), Error> {
52        let start = Instant::now();
53
54        while !self.operation.done.unwrap_or(false) {
55            if let Some(timeout) = timeout {
56                if start.elapsed() >= timeout {
57                    return Err(Error::OperationTimeout {
58                        name: self.operation.name.clone(),
59                    });
60                }
61            }
62
63            tokio::time::sleep(interval).await;
64            self.refresh().await?;
65        }
66
67        if let Some(OperationResult::Error { error }) = &self.operation.result {
68            return Err(Error::OperationFailed {
69                name: self.operation.name.clone(),
70                code: error.code,
71                message: error.message.clone(),
72            });
73        }
74
75        Ok(())
76    }
77}
gemini_rust/file_search/operation_handle.rs

gemini_rust/file_search/
operation_handle.rs
