1use std::time::Duration;
4use std::time::Instant;
5
6use bytes::Bytes;
7use bytes::BytesMut;
8use ferrin_provider_util::http::HttpRequest;
9use ferrin_provider_util::http::RequestBody;
10use ferrin_provider_util::http::ResponseHandlers;
11use ferrin_provider_util::http::delete;
12use ferrin_provider_util::http::get;
13use ferrin_provider_util::http::json_response_handler;
14use ferrin_provider_util::http::post_json;
15use ferrin_provider_util::http::send;
16use ferrin_provider_util::http::text_response_handler;
17use ferrin_provider_util::secure_url::validate_url;
18use ferrin_spec::Headers;
19use ferrin_spec::JsonObject;
20use ferrin_spec::JsonValue;
21use ferrin_spec::MediaType;
22use ferrin_spec::ProviderId;
23use ferrin_spec::ProviderReference;
24use ferrin_spec::error::ApiCallError;
25use ferrin_spec::error::InvalidArgumentError;
26use ferrin_spec::error::InvalidResponseDataError;
27use ferrin_spec::error::ProviderError;
28use ferrin_spec::files::DeleteFileResult;
29use ferrin_spec::files::FileMetadataResult;
30use ferrin_spec::files::FileReferenceOptions;
31use ferrin_spec::files::Files;
32use ferrin_spec::files::UploadData;
33use ferrin_spec::files::UploadFileOptions;
34use ferrin_spec::files::UploadFileResult;
35use futures_util::StreamExt;
36use futures_util::future::Either;
37use futures_util::future::select;
38use serde::Deserialize;
39use serde_json::json;
40use tokio_util::sync::CancellationToken;
41use url::Url;
42
43use crate::api_types::deserialize_count;
44use crate::config::CANONICAL_OPTIONS_KEY;
45use crate::config::SharedConfig;
46use crate::config::UPLOAD_PATH;
47use crate::convert_prompt::resolve_reference;
48use crate::error::failed_response_handler;
49use crate::options::parse_merged;
50use crate::output::OutputMapper;
51
52pub const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
54
55pub const DEFAULT_POLL_TIMEOUT_MS: u64 = 300_000;
57
58#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
60#[serde(rename_all = "camelCase")]
61pub struct GoogleFilesOptions {
62 #[serde(default)]
64 pub display_name: Option<String>,
65 #[serde(default)]
67 pub poll_interval_ms: Option<u64>,
68 #[serde(default)]
70 pub poll_timeout_ms: Option<u64>,
71}
72
73#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
75#[serde(rename_all = "camelCase")]
76pub struct GoogleFile {
77 pub name: String,
79 #[serde(default)]
81 pub display_name: Option<String>,
82 #[serde(default)]
84 pub mime_type: Option<String>,
85 #[serde(default, deserialize_with = "deserialize_count")]
87 pub size_bytes: Option<u64>,
88 #[serde(default)]
90 pub create_time: Option<String>,
91 #[serde(default)]
93 pub update_time: Option<String>,
94 #[serde(default)]
96 pub expiration_time: Option<String>,
97 #[serde(default)]
99 pub sha256_hash: Option<String>,
100 #[serde(default)]
102 pub uri: Option<String>,
103 #[serde(default)]
105 pub state: Option<String>,
106}
107
108#[derive(Debug, Deserialize)]
110#[serde(untagged)]
111enum FileResponse {
112 Envelope { file: GoogleFile },
113 Bare(GoogleFile),
114}
115
116impl FileResponse {
117 fn into_file(self) -> GoogleFile {
118 match self {
119 Self::Envelope { file } | Self::Bare(file) => file,
120 }
121 }
122}
123
124#[derive(Debug, Clone)]
126pub struct UploadRequest {
127 pub data: Bytes,
129 pub media_type: String,
131 pub display_name: Option<String>,
133 pub poll_interval: Duration,
135 pub poll_timeout: Duration,
137 pub headers: Headers,
139 pub cancellation: CancellationToken,
141}
142
143impl UploadRequest {
144 #[must_use]
146 pub fn new(data: Bytes, media_type: impl Into<String>) -> Self {
147 Self {
148 data,
149 media_type: media_type.into(),
150 display_name: None,
151 poll_interval: Duration::from_millis(DEFAULT_POLL_INTERVAL_MS),
152 poll_timeout: Duration::from_millis(DEFAULT_POLL_TIMEOUT_MS),
153 headers: Headers::new(),
154 cancellation: CancellationToken::new(),
155 }
156 }
157}
158
159async fn collect(data: UploadData) -> Result<Bytes, ProviderError> {
160 match data {
161 UploadData::Bytes(bytes) => Ok(bytes),
162 UploadData::Text(text) => Ok(Bytes::from(text)),
163 UploadData::Stream(stream) => {
164 let mut buffer = BytesMut::new();
165 let mut stream = stream;
166 while let Some(chunk) = stream.next().await {
167 buffer.extend_from_slice(&chunk?);
168 }
169 Ok(buffer.freeze())
170 }
171 #[allow(unreachable_patterns, reason = "UploadData is non-exhaustive")]
172 _ => Err(ProviderError::unsupported("upload data type")),
173 }
174}
175
176fn parse_time(value: Option<&str>) -> Option<chrono::DateTime<chrono::Utc>> {
177 value
178 .and_then(|time| chrono::DateTime::parse_from_rfc3339(time).ok())
179 .map(|time| time.with_timezone(&chrono::Utc))
180}
181
182fn file_name(reference: &str) -> String {
185 if let Some((_, rest)) = reference.rsplit_once("/files/") {
186 return format!("files/{rest}");
187 }
188 if reference.starts_with("files/") {
189 return reference.to_owned();
190 }
191 format!("files/{reference}")
192}
193
194#[derive(Debug, Clone)]
196pub struct GoogleFiles {
197 config: SharedConfig,
198 provider: ProviderId,
199}
200
201impl GoogleFiles {
202 #[must_use]
204 pub fn new(config: SharedConfig) -> Self {
205 Self {
206 provider: ProviderId::new(config.name.clone()),
207 config,
208 }
209 }
210
211 fn reference(&self, file: &GoogleFile) -> ProviderReference {
212 let value = file.uri.clone().unwrap_or_else(|| file.name.clone());
213 let mut reference = ProviderReference::new();
214 reference.insert(CANONICAL_OPTIONS_KEY.to_owned(), value.clone());
215 reference.insert(self.config.name.clone(), value);
216 reference
217 }
218
219 #[must_use]
221 pub fn to_result(&self, file: &GoogleFile) -> UploadFileResult {
222 let mut meta = JsonObject::new();
223 let string =
224 |value: &Option<String>| value.clone().map_or(JsonValue::Null, JsonValue::from);
225 meta.insert("name".to_owned(), JsonValue::from(file.name.clone()));
226 meta.insert("displayName".to_owned(), string(&file.display_name));
227 meta.insert("mimeType".to_owned(), string(&file.mime_type));
228 meta.insert(
229 "sizeBytes".to_owned(),
230 file.size_bytes
231 .map_or(JsonValue::Null, |size| JsonValue::from(size.to_string())),
232 );
233 meta.insert("state".to_owned(), string(&file.state));
234 meta.insert("uri".to_owned(), string(&file.uri));
235 for (key, value) in [
236 ("createTime", &file.create_time),
237 ("updateTime", &file.update_time),
238 ("expirationTime", &file.expiration_time),
239 ("sha256Hash", &file.sha256_hash),
240 ] {
241 if let Some(value) = value {
242 meta.insert(key.to_owned(), json!(value));
243 }
244 }
245 let mapper = OutputMapper::new(self.config.clone(), Default::default());
246 UploadFileResult {
247 provider_reference: self.reference(file),
248 media_type: file.mime_type.as_deref().map(MediaType::new),
249 filename: file.display_name.clone(),
250 byte_size: file.size_bytes,
251 created_at: parse_time(file.create_time.as_deref()),
252 expires_at: parse_time(file.expiration_time.as_deref()),
253 provider_metadata: Some(mapper.metadata(meta)),
254 warnings: Vec::new(),
255 }
256 }
257
258 fn resource_url(&self, name: &str) -> Result<Url, ProviderError> {
259 let mut url = self.config.base_url.clone();
260 let mut segments = url.path_segments_mut().map_err(|()| {
261 InvalidArgumentError::new("base_url", "file base URL cannot contain path segments")
262 })?;
263 segments.pop_if_empty();
264 let id = if let Some(id) = name.strip_prefix("files/").filter(|id| !id.contains('/')) {
265 segments.push("files");
266 id
267 } else {
268 name
269 };
270 segments.push(match id {
272 "." => "%2E",
273 ".." => "%2E%2E",
274 id => id,
275 });
276 drop(segments);
277 Ok(url)
278 }
279
280 pub async fn fetch_file(
286 &self,
287 name: &str,
288 headers: &Headers,
289 cancellation: CancellationToken,
290 ) -> Result<GoogleFile, ProviderError> {
291 let handlers = ResponseHandlers::new(
292 json_response_handler::<FileResponse>(),
293 failed_response_handler(),
294 );
295 let response = get(
296 self.config.transport.as_ref(),
297 self.resource_url(name)?,
298 self.config.headers(headers)?,
299 &handlers,
300 cancellation,
301 )
302 .await?;
303 Ok(response.value.into_file())
304 }
305
306 #[tracing::instrument(skip_all, fields(media_type = %request.media_type, bytes = request.data.len()))]
316 pub async fn upload_bytes(&self, request: UploadRequest) -> Result<GoogleFile, ProviderError> {
317 let start_url = self.config.origin_url(UPLOAD_PATH);
318 let start_headers = self
319 .config
320 .headers(&request.headers)?
321 .with("x-goog-upload-protocol", "resumable")
322 .with("x-goog-upload-command", "start")
323 .with(
324 "x-goog-upload-header-content-length",
325 &request.data.len().to_string(),
326 )
327 .with("x-goog-upload-header-content-type", &request.media_type);
328 let mut file = JsonObject::new();
329 if let Some(name) = &request.display_name {
330 file.insert("display_name".to_owned(), JsonValue::from(name.as_str()));
331 }
332 let start_handlers =
333 ResponseHandlers::new(text_response_handler(), failed_response_handler());
334 let started = post_json(
335 self.config.transport.as_ref(),
336 start_url,
337 start_headers,
338 &json!({"file": file}),
339 &start_handlers,
340 request.cancellation.clone(),
341 )
342 .await?;
343 let upload_url = started
344 .response_headers
345 .get_str("x-goog-upload-url")
346 .and_then(|value| Url::parse(value).ok())
347 .ok_or_else(|| {
348 ProviderError::InvalidResponseData(Box::new(InvalidResponseDataError::new(
349 "google did not return a resumable upload URL",
350 JsonValue::Null,
351 )))
352 })?;
353 let validated = validate_url(&upload_url, &self.config.url_policy)
354 .await
355 .map_err(|error| {
356 InvalidResponseDataError::new(
357 format!("google returned an unsafe upload URL: {error}"),
358 JsonValue::Null,
359 )
360 })?;
361 let mut finalize_headers = if upload_url.origin() == self.config.base_url.origin()
362 || self.config.url_policy.is_credentialed(&upload_url)
363 {
364 self.config.unauthenticated_headers(&request.headers)
365 } else {
366 Headers::new().with_user_agent_suffix([crate::config::USER_AGENT])
367 };
368 finalize_headers.remove(crate::config::API_KEY_HEADER);
369 let finalize_headers = finalize_headers
370 .with("x-goog-upload-offset", "0")
371 .with("x-goog-upload-command", "upload, finalize");
372 let handlers = ResponseHandlers::new(
373 json_response_handler::<FileResponse>()
374 .with_max_bytes(self.config.url_policy.max_body_bytes),
375 failed_response_handler().with_max_bytes(self.config.url_policy.max_body_bytes),
376 );
377 let finalize = HttpRequest::post(upload_url.clone())
378 .with_headers(finalize_headers)
379 .with_body(RequestBody::Bytes {
380 content_type: request.media_type.clone(),
381 data: request.data,
382 })
383 .with_cancellation(request.cancellation.clone())
384 .with_pinned_addresses(validated.addresses);
385 let uploaded = send(self.config.transport.as_ref(), finalize, None, &handlers).await?;
386 let mut file = uploaded.value.into_file();
387 let started_at = Instant::now();
388 while file.state.as_deref() == Some("PROCESSING") {
389 if started_at.elapsed() > request.poll_timeout {
390 return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
391 format!(
392 "file processing timed out after {}ms",
393 request.poll_timeout.as_millis()
394 ),
395 upload_url,
396 ))));
397 }
398 let sleep = Box::pin(tokio::time::sleep(request.poll_interval));
399 let cancelled = Box::pin(request.cancellation.cancelled());
400 if let Either::Right(_) = select(sleep, cancelled).await {
401 return Err(ProviderError::Cancelled);
402 }
403 file = self
404 .fetch_file(&file.name, &request.headers, request.cancellation.clone())
405 .await?;
406 }
407 if file.state.as_deref() == Some("FAILED") {
408 return Err(ProviderError::ApiCall(Box::new(ApiCallError::new(
409 format!("file processing failed for {}", file.name),
410 upload_url,
411 ))));
412 }
413 Ok(file)
414 }
415}
416
417impl Files for GoogleFiles {
418 fn provider(&self) -> &ProviderId {
419 &self.provider
420 }
421
422 async fn upload_file(
423 &self,
424 options: UploadFileOptions,
425 ) -> Result<UploadFileResult, ProviderError> {
426 let google = parse_merged::<GoogleFilesOptions>(
427 &self.config,
428 &options.provider_options,
429 |mut canonical, custom| {
430 if custom.display_name.is_some() {
431 canonical.display_name = custom.display_name;
432 }
433 if custom.poll_interval_ms.is_some() {
434 canonical.poll_interval_ms = custom.poll_interval_ms;
435 }
436 if custom.poll_timeout_ms.is_some() {
437 canonical.poll_timeout_ms = custom.poll_timeout_ms;
438 }
439 canonical
440 },
441 )?;
442 if google.poll_interval_ms == Some(0) || google.poll_timeout_ms == Some(0) {
443 return Err(InvalidArgumentError::new(
444 "provider_options",
445 "file polling intervals and timeouts must be positive",
446 )
447 .into());
448 }
449 let ignored_filename = options.filename.is_some();
450 let data = match select(
451 Box::pin(collect(options.data)),
452 Box::pin(options.cancellation.cancelled()),
453 )
454 .await
455 {
456 Either::Left((data, _)) => data?,
457 Either::Right(_) => return Err(ProviderError::Cancelled),
458 };
459 let mut request = UploadRequest::new(data, options.media_type.as_str());
460 request.display_name = google.display_name;
461 if let Some(interval) = google.poll_interval_ms {
462 request.poll_interval = Duration::from_millis(interval);
463 }
464 if let Some(timeout) = google.poll_timeout_ms {
465 request.poll_timeout = Duration::from_millis(timeout);
466 }
467 request.headers = options.headers;
468 request.cancellation = options.cancellation;
469 let file = self.upload_bytes(request).await?;
470 let mut result = self.to_result(&file);
471 if ignored_filename {
472 result
473 .warnings
474 .push(ferrin_spec::Warning::unsupported("filename"));
475 }
476 if result.media_type.is_none() {
477 result.media_type = Some(options.media_type);
478 }
479 Ok(result)
480 }
481
482 fn supports_get_file_metadata(&self) -> bool {
483 true
484 }
485
486 async fn get_file_metadata(
487 &self,
488 options: FileReferenceOptions,
489 ) -> Result<FileMetadataResult, ProviderError> {
490 let name = file_name(resolve_reference(&self.config, &options.file)?);
491 let file = self
492 .fetch_file(&name, &options.headers, options.cancellation)
493 .await?;
494 Ok(self.to_result(&file))
495 }
496
497 fn supports_delete_file(&self) -> bool {
498 true
499 }
500
501 async fn delete_file(
502 &self,
503 options: FileReferenceOptions,
504 ) -> Result<DeleteFileResult, ProviderError> {
505 let name = file_name(resolve_reference(&self.config, &options.file)?);
506 let handlers = ResponseHandlers::new(text_response_handler(), failed_response_handler());
507 delete(
508 self.config.transport.as_ref(),
509 self.resource_url(&name)?,
510 self.config.headers(&options.headers)?,
511 &handlers,
512 options.cancellation,
513 )
514 .await?;
515 Ok(DeleteFileResult {
516 provider_reference: options.file,
517 deleted: true,
518 provider_metadata: None,
519 warnings: Vec::new(),
520 })
521 }
522}