From d2d56d722095ef7bfa950e852c9f7b7af0d7b676 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 11 May 2026 17:09:30 -0400 Subject: [PATCH] polling --- .../fal_client/src/error/fal_error_plus.rs | 3 + crates/api_clients/fal_client/src/lib.rs | 2 + .../api_clients/fal_client/src/polling/mod.rs | 1 + .../fal_client/src/polling/poll_job_status.rs | 87 +++++++++++++++++++ .../src/requests/core_api/fal_request.rs | 6 +- .../api_clients/fal_client/src/types/mod.rs | 0 6 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 crates/api_clients/fal_client/src/polling/mod.rs create mode 100644 crates/api_clients/fal_client/src/polling/poll_job_status.rs create mode 100644 crates/api_clients/fal_client/src/types/mod.rs diff --git a/crates/api_clients/fal_client/src/error/fal_error_plus.rs b/crates/api_clients/fal_client/src/error/fal_error_plus.rs index eaecf25970..01bac0afe9 100644 --- a/crates/api_clients/fal_client/src/error/fal_error_plus.rs +++ b/crates/api_clients/fal_client/src/error/fal_error_plus.rs @@ -12,6 +12,8 @@ pub enum FalErrorPlus { FalApiKeyError(String), /// The fal account has a billing issue FalBillingError(String), + /// Error with an invalid polling URL. + InvalidPollingUrl(String), /// Another error we didn't handle. AnyhowError(anyhow::Error), /// URL parse errors. @@ -28,6 +30,7 @@ impl Display for FalErrorPlus { Self::FalError(err) => write!(f, "FalErrorPlus::FalError: {:?}", err), Self::FalApiKeyError(reason) => write!(f, "FalErrorPlus::FalApiKeyError: {}", reason), Self::FalBillingError(reason) => write!(f, "FalErrorPlus::FalBillingError: {}", reason), + Self::InvalidPollingUrl(url) => write!(f, "FalErrorPlus::InvalidPollingUrl: {}", url), Self::AnyhowError(err) => write!(f, "FalErrorPlus::AnyhowError: {:?}", err), Self::UrlParseError(err) => write!(f, "FalErrorPlus::UrlParseError: {:?}", err), Self::UnhandledEndpoint(endpoint) => write!(f, "FalErrorPlus::UnhandledEndpoint: {:?}", endpoint), diff --git a/crates/api_clients/fal_client/src/lib.rs b/crates/api_clients/fal_client/src/lib.rs index b79e7ee89b..3b5f695d94 100644 --- a/crates/api_clients/fal_client/src/lib.rs +++ b/crates/api_clients/fal_client/src/lib.rs @@ -2,6 +2,8 @@ pub mod creds; pub mod error; +pub mod polling; pub mod requests; +pub mod types; pub mod utils; pub mod webhook_payload; diff --git a/crates/api_clients/fal_client/src/polling/mod.rs b/crates/api_clients/fal_client/src/polling/mod.rs new file mode 100644 index 0000000000..22bc850576 --- /dev/null +++ b/crates/api_clients/fal_client/src/polling/mod.rs @@ -0,0 +1 @@ +pub mod poll_job_status; diff --git a/crates/api_clients/fal_client/src/polling/poll_job_status.rs b/crates/api_clients/fal_client/src/polling/poll_job_status.rs new file mode 100644 index 0000000000..3bf4285093 --- /dev/null +++ b/crates/api_clients/fal_client/src/polling/poll_job_status.rs @@ -0,0 +1,87 @@ +use crate::creds::fal_api_key::FalApiKey; +use crate::error::fal_error_plus::FalErrorPlus; +use log::info; +use url::Url; + +const EXPECTED_HOST: &str = "queue.fal.run"; + +pub struct PollJobStatusArgs<'a> { + pub status_url: &'a str, + pub api_key: &'a FalApiKey, +} + +/// Poll the status of a queued FAL job. +/// +/// The `status_url` must point to `queue.fal.run`. Returns the raw response +/// body as a string so the caller can deserialize it as needed. +pub async fn poll_job_status(args: PollJobStatusArgs<'_>) -> Result { + let parsed = Url::parse(args.status_url)?; + + let host = parsed.host_str().unwrap_or(""); + if host != EXPECTED_HOST { + return Err(FalErrorPlus::InvalidPollingUrl(format!( + "Expected host '{}' but got '{}' in status URL: {}", + EXPECTED_HOST, + host, + args.status_url, + ))); + } + + info!("Polling FAL job status: {}", args.status_url); + + let response = reqwest::Client::new() + .get(args.status_url) + .header("Authorization", format!("Key {}", args.api_key.0)) + .send() + .await?; + + let status = response.status(); + let body = response.text().await?; + + if !status.is_success() { + return Err(FalErrorPlus::AnyhowError(anyhow::anyhow!( + "FAL poll returned HTTP {}: {}", + status, + body, + ))); + } + + Ok(body) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn rejects_wrong_host() { + let api_key = FalApiKey::from_str("test-key"); + let args = PollJobStatusArgs { + status_url: "https://evil.example.com/fal-ai/flux/requests/abc123", + api_key: &api_key, + }; + let result = poll_job_status(args).await; + assert!(result.is_err()); + let err = format!("{}", result.unwrap_err()); + assert!(err.contains("evil.example.com"), "error should mention the bad host: {}", err); + } + + #[tokio::test] + async fn rejects_invalid_url() { + let api_key = FalApiKey::from_str("test-key"); + let args = PollJobStatusArgs { + status_url: "not a url at all", + api_key: &api_key, + }; + let result = poll_job_status(args).await; + assert!(result.is_err()); + } + + #[test] + fn accepts_valid_host() { + // Just verify URL parsing succeeds — the actual HTTP call will fail without a real key, + // but we can check that the host validation passes. + let parsed = Url::parse("https://queue.fal.run/fal-ai/flux/requests/019e18d8-8c36-7bc1-aa77-2bc2f70268c6").unwrap(); + assert_eq!(parsed.host_str(), Some(EXPECTED_HOST)); + } +} diff --git a/crates/api_clients/fal_client/src/requests/core_api/fal_request.rs b/crates/api_clients/fal_client/src/requests/core_api/fal_request.rs index 8f4c087133..99da819e6d 100644 --- a/crates/api_clients/fal_client/src/requests/core_api/fal_request.rs +++ b/crates/api_clients/fal_client/src/requests/core_api/fal_request.rs @@ -153,7 +153,11 @@ impl FalRequest return Err(error.into()); } - let payload: QueueResponse = response.error_for_status()?.json().await?; + let body = response.text().await?; + + println!("Queue response: {}", body); + + let payload: QueueResponse = serde_json::from_str(&body)?; Ok(payload) } diff --git a/crates/api_clients/fal_client/src/types/mod.rs b/crates/api_clients/fal_client/src/types/mod.rs new file mode 100644 index 0000000000..e69de29bb2