feat(jimeng): add request signing, body validation and polling state machine

src/services/jimeng/mod.rs

@@ -1,9 +1,387 @@
 //! Jimeng (Volcengine Visual / CV) async generation provider.
 //!
-//! Both image and video capabilities share one submit→poll→store state
+//! The provider uses [`volc_visual_sdk::VisualClient`] purely as a request
-//! machine and differ only by `req_key` and a few poll parameters (see
+//! *signer* (`build_signed_request`) — the actual HTTP is performed over the
-//! [`types::JimengCapability`]). This module hosts the capability mapping,
+//! shared async `reqwest::Client` in [`AppState`] so we never block a tokio
-//! request/response types and business error-code classification; the submit
+//! worker. Both image and video capabilities share one submit→poll→store
-//! and polling state machine are layered on top in later changes.
+//! state machine and differ only by `req_key` and a few poll parameters
+//! (see [`JimengCapability`]).
 
+pub mod poller;
 pub mod types;
+
+use crate::config::JimengConfig;
+use serde_json::{json, Value};
+use types::{
+    classify, CodeClass, GetResultResponse, ImageRequest, JimengCapability, SubmitResponse,
+    ACTION_GET_RESULT, ACTION_SUBMIT, ALLOWED_ASPECT_RATIOS, ALLOWED_FRAMES, API_VERSION,
+};
+use volc_visual_sdk::VisualClient;
+
+use self::types::VideoRequest;
+
+/// Errors produced by the Jimeng provider.
+#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
+pub enum JimengError {
+    /// Provider is disabled in config.
+    #[error("jimeng provider is not enabled")]
+    Disabled,
+    /// Request validation failed (caller error).
+    #[error("invalid request: {0}")]
+    Invalid(String),
+    /// Signing the request failed (bad credentials / serialization).
+    #[error("signing error: {0}")]
+    Sign(String),
+    /// Network / transport failure talking to Volcengine.
+    #[error("http error: {0}")]
+    Http(String),
+    /// Response body could not be decoded.
+    #[error("decode error: {0}")]
+    Decode(String),
+    /// Volcengine returned a non-success business `code`.
+    #[error("api error (code {code}): {message}")]
+    Api {
+        /// The business error code.
+        code: i64,
+        /// The accompanying message.
+        message: String,
+        /// Whether the code is classified as retryable.
+        retryable: bool,
+    },
+    /// The task id was not found (or aged past 12h retention).
+    #[error("task not found")]
+    NotFound,
+    /// The task expired before completion.
+    #[error("task expired")]
+    Expired,
+    /// A finished task produced no usable output.
+    #[error("task completed without output")]
+    EmptyOutput,
+    /// Polling exceeded the configured timeout.
+    #[error("timed out after {0}s")]
+    Timeout(u64),
+    /// The local task was cancelled while polling.
+    #[error("cancelled")]
+    Cancelled,
+}
+
+/// A signer-backed Jimeng provider. Cheap to construct; holds only the signing
+/// client (no HTTP client of its own — that lives in [`AppState`]).
+#[derive(Clone)]
+pub struct JimengProvider {
+    client: VisualClient,
+}
+
+impl JimengProvider {
+    /// Build a provider from config. Returns [`JimengError::Disabled`] when the
+    /// provider is turned off so callers can short-circuit uniformly.
+    pub fn from_config(cfg: &JimengConfig) -> Result<Self, JimengError> {
+        if !cfg.enabled {
+            return Err(JimengError::Disabled);
+        }
+        let mut client = VisualClient::new(cfg.access_key.clone(), cfg.secret_key.clone())
+            .with_region(cfg.region.clone())
+            .with_host(cfg.host.clone());
+        if let Some(token) = cfg.session_token.as_ref().filter(|t| !t.is_empty()) {
+            client = client.with_security_token(token.clone());
+        }
+        Ok(Self { client })
+    }
+
+    /// Submit a request body for `capability` and return the Volcengine task id.
+    pub async fn submit(
+        &self,
+        http: &reqwest::Client,
+        capability: JimengCapability,
+        body: &Value,
+    ) -> Result<String, JimengError> {
+        let text = self.post_signed(http, ACTION_SUBMIT, body).await?;
+        let parsed: SubmitResponse = serde_json::from_str(&text)
+            .map_err(|e| JimengError::Decode(format!("{e}; raw: {text}")))?;
+        match classify(parsed.code, capability) {
+            CodeClass::Success => parsed
+                .data
+                .map(|d| d.task_id)
+                .ok_or(JimengError::EmptyOutput),
+            class => Err(JimengError::Api {
+                code: parsed.code,
+                message: parsed.message.unwrap_or_else(|| "unknown error".into()),
+                retryable: class == CodeClass::Retryable,
+            }),
+        }
+    }
+
+    /// Poll a task's result for `capability`. Returns the parsed envelope; the
+    /// caller (poller) interprets `code`/`status`.
+    pub async fn get_result(
+        &self,
+        http: &reqwest::Client,
+        capability: JimengCapability,
+        volc_task_id: &str,
+    ) -> Result<GetResultResponse, JimengError> {
+        let body = build_get_result_body(capability, volc_task_id);
+        let text = self.post_signed(http, ACTION_GET_RESULT, &body).await?;
+        serde_json::from_str(&text).map_err(|e| JimengError::Decode(format!("{e}; raw: {text}")))
+    }
+
+    /// Sign `body` for `action` and POST it over the shared async client,
+    /// returning the raw response text. The exact bytes that were signed are
+    /// the bytes that are sent (`prepared.body`).
+    async fn post_signed(
+        &self,
+        http: &reqwest::Client,
+        action: &str,
+        body: &Value,
+    ) -> Result<String, JimengError> {
+        let prepared = self
+            .client
+            .build_signed_request(action, API_VERSION, body)
+            .map_err(|e| JimengError::Sign(e.to_string()))?;
+
+        let mut req = http
+            .post(&prepared.url)
+            .header("Content-Type", prepared.signed.content_type.as_str())
+            .header("Host", prepared.signed.host.as_str())
+            .header("X-Date", prepared.signed.x_date.as_str())
+            .header("X-Content-Sha256", prepared.signed.x_content_sha256.as_str())
+            .header("Authorization", prepared.signed.authorization.as_str());
+        if let Some(token) = &prepared.signed.x_security_token {
+            req = req.header("X-Security-Token", token.as_str());
+        }
+
+        let resp = req
+            .body(prepared.body)
+            .send()
+            .await
+            .map_err(|e| JimengError::Http(e.to_string()))?;
+        resp.text()
+            .await
+            .map_err(|e| JimengError::Http(e.to_string()))
+    }
+}
+
+/// Build the GetResult request body for a capability. Images request URLs back
+/// (`return_url:true`); video returns its single url directly.
+pub fn build_get_result_body(capability: JimengCapability, volc_task_id: &str) -> Value {
+    match capability {
+        JimengCapability::Image => json!({
+            "req_key": capability.req_key(),
+            "task_id": volc_task_id,
+            "req_json": "{\"return_url\":true}",
+        }),
+        JimengCapability::Video => json!({
+            "req_key": capability.req_key(),
+            "task_id": volc_task_id,
+        }),
+    }
+}
+
+/// Validate and build the submit body for an image generation.
+pub fn build_image_body(req: &ImageRequest) -> Result<Value, JimengError> {
+    if req.prompt.trim().is_empty() {
+        return Err(JimengError::Invalid("prompt is required".into()));
+    }
+    let mut body = json!({
+        "req_key": JimengCapability::Image.req_key(),
+        "prompt": req.prompt,
+    });
+    let obj = body.as_object_mut().expect("object");
+    if let Some(width) = req.width {
+        obj.insert("width".into(), json!(width));
+    }
+    if let Some(height) = req.height {
+        obj.insert("height".into(), json!(height));
+    }
+    if let Some(scale) = req.scale {
+        obj.insert("scale".into(), json!(scale));
+    }
+    if let Some(force_single) = req.force_single {
+        obj.insert("force_single".into(), json!(force_single));
+    }
+    Ok(body)
+}
+
+/// Validate and build the submit body for a video generation.
+///
+/// Requires either a non-empty `prompt` or at least one first-frame
+/// `image_urls` entry; `frames` must be one of [`ALLOWED_FRAMES`] and
+/// `aspect_ratio`, when present, must be in [`ALLOWED_ASPECT_RATIOS`].
+pub fn build_video_body(req: &VideoRequest) -> Result<Value, JimengError> {
+    let has_prompt = req.prompt.as_deref().map(str::trim).is_some_and(|p| !p.is_empty());
+    let has_image = !req.image_urls.is_empty();
+    if !has_prompt && !has_image {
+        return Err(JimengError::Invalid(
+            "either prompt or image_urls is required".into(),
+        ));
+    }
+    if !ALLOWED_FRAMES.contains(&req.frames) {
+        return Err(JimengError::Invalid(format!(
+            "frames must be one of {ALLOWED_FRAMES:?} (got {})",
+            req.frames
+        )));
+    }
+    if let Some(ratio) = req.aspect_ratio.as_deref() {
+        if !ALLOWED_ASPECT_RATIOS.contains(&ratio) {
+            return Err(JimengError::Invalid(format!(
+                "aspect_ratio must be one of {ALLOWED_ASPECT_RATIOS:?} (got {ratio})"
+            )));
+        }
+    }
+
+    let mut body = json!({
+        "req_key": JimengCapability::Video.req_key(),
+        "frames": req.frames,
+    });
+    let obj = body.as_object_mut().expect("object");
+    if let Some(prompt) = req.prompt.as_ref().filter(|p| !p.trim().is_empty()) {
+        obj.insert("prompt".into(), json!(prompt));
+    }
+    if has_image {
+        obj.insert("image_urls".into(), json!(req.image_urls));
+    }
+    if let Some(ratio) = &req.aspect_ratio {
+        obj.insert("aspect_ratio".into(), json!(ratio));
+    }
+    if let Some(seed) = req.seed {
+        obj.insert("seed".into(), json!(seed));
+    }
+    Ok(body)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn image_body_requires_prompt() {
+        let req = ImageRequest {
+            prompt: "  ".into(),
+            width: None,
+            height: None,
+            scale: None,
+            force_single: None,
+        };
+        assert!(matches!(build_image_body(&req), Err(JimengError::Invalid(_))));
+    }
+
+    #[test]
+    fn image_body_includes_optional_fields() {
+        let req = ImageRequest {
+            prompt: "a cat".into(),
+            width: Some(2048),
+            height: Some(2048),
+            scale: Some(0.5),
+            force_single: Some(true),
+        };
+        let body = build_image_body(&req).unwrap();
+        assert_eq!(body["req_key"], "jimeng_t2i_v40");
+        assert_eq!(body["prompt"], "a cat");
+        assert_eq!(body["width"], 2048);
+        assert_eq!(body["height"], 2048);
+        assert_eq!(body["scale"], 0.5);
+        assert_eq!(body["force_single"], true);
+    }
+
+    #[test]
+    fn image_body_omits_absent_optional_fields() {
+        let req = ImageRequest {
+            prompt: "a cat".into(),
+            width: None,
+            height: None,
+            scale: None,
+            force_single: None,
+        };
+        let body = build_image_body(&req).unwrap();
+        let obj = body.as_object().unwrap();
+        assert!(!obj.contains_key("width"));
+        assert!(!obj.contains_key("scale"));
+    }
+
+    #[test]
+    fn video_body_requires_prompt_or_image() {
+        let req = VideoRequest {
+            prompt: None,
+            image_urls: vec![],
+            frames: 121,
+            aspect_ratio: None,
+            seed: None,
+        };
+        assert!(matches!(build_video_body(&req), Err(JimengError::Invalid(_))));
+    }
+
+    #[test]
+    fn video_body_accepts_prompt_only() {
+        let req = VideoRequest {
+            prompt: Some("horses".into()),
+            image_urls: vec![],
+            frames: 121,
+            aspect_ratio: Some("16:9".into()),
+            seed: Some(-1),
+        };
+        let body = build_video_body(&req).unwrap();
+        assert_eq!(body["req_key"], "jimeng_ti2v_v30_pro");
+        assert_eq!(body["prompt"], "horses");
+        assert_eq!(body["frames"], 121);
+        assert_eq!(body["aspect_ratio"], "16:9");
+        assert_eq!(body["seed"], -1);
+        assert!(!body.as_object().unwrap().contains_key("image_urls"));
+    }
+
+    #[test]
+    fn video_body_accepts_image_only() {
+        let req = VideoRequest {
+            prompt: None,
+            image_urls: vec!["https://x/a.jpg".into()],
+            frames: 241,
+            aspect_ratio: None,
+            seed: None,
+        };
+        let body = build_video_body(&req).unwrap();
+        assert_eq!(body["frames"], 241);
+        assert_eq!(body["image_urls"][0], "https://x/a.jpg");
+        assert!(!body.as_object().unwrap().contains_key("prompt"));
+    }
+
+    #[test]
+    fn video_body_rejects_bad_frames() {
+        let req = VideoRequest {
+            prompt: Some("x".into()),
+            image_urls: vec![],
+            frames: 100,
+            aspect_ratio: None,
+            seed: None,
+        };
+        assert!(matches!(build_video_body(&req), Err(JimengError::Invalid(_))));
+    }
+
+    #[test]
+    fn video_body_rejects_bad_aspect_ratio() {
+        let req = VideoRequest {
+            prompt: Some("x".into()),
+            image_urls: vec![],
+            frames: 121,
+            aspect_ratio: Some("21:9".into()),
+            seed: None,
+        };
+        assert!(matches!(build_video_body(&req), Err(JimengError::Invalid(_))));
+    }
+
+    #[test]
+    fn get_result_body_shapes_differ_by_capability() {
+        let img = build_get_result_body(JimengCapability::Image, "tid");
+        assert_eq!(img["task_id"], "tid");
+        assert_eq!(img["req_json"], "{\"return_url\":true}");
+        let vid = build_get_result_body(JimengCapability::Video, "tid");
+        assert_eq!(vid["task_id"], "tid");
+        assert!(!vid.as_object().unwrap().contains_key("req_json"));
+    }
+
+    #[test]
+    fn disabled_config_yields_disabled_error() {
+        let cfg = JimengConfig::default(); // enabled = false
+        assert!(matches!(
+            JimengProvider::from_config(&cfg),
+            Err(JimengError::Disabled)
+        ));
+    }
+}

src/services/jimeng/poller.rs

@@ -0,0 +1,352 @@
+//! Generic polling state machine for Jimeng async tasks.
+//!
+//! The decision logic ([`decide`]) is a pure function over a parsed
+//! [`GetResultResponse`] so it can be unit tested with mock sequences. The
+//! async driver ([`poll_until_done`]) wraps it with sleep/backoff, a timeout
+//! deadline and a cancellation check, and is generic over the fetch and
+//! cancel closures so tests can drive it without any network or database.
+
+use super::types::{
+    classify, CodeClass, GetResultData, GetResultResponse, JimengCapability, JimengOutput,
+    TaskStatus,
+};
+use super::JimengError;
+use std::future::Future;
+use std::time::{Duration, Instant};
+
+/// The next action the poll loop should take after inspecting one response.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum PollDecision {
+    /// Task still running — wait `interval` then poll again.
+    Wait,
+    /// Transient failure — back off then retry (carries the offending code).
+    Backoff(i64),
+    /// Task finished successfully with the given output.
+    Done(JimengOutput),
+    /// Permanent failure — stop and report.
+    Fail(JimengError),
+}
+
+/// Decide the next action from a single GetResult response.
+///
+/// Enforces the cardinal rule from the integration notes: inspect the business
+/// `code` *before* `data.status` (on error `data` is null). Retryable codes
+/// back off; fatal codes fail immediately; on `code == 10000` the `status`
+/// drives the outcome.
+pub fn decide(capability: JimengCapability, resp: &GetResultResponse) -> PollDecision {
+    match classify(resp.code, capability) {
+        CodeClass::Retryable => return PollDecision::Backoff(resp.code),
+        CodeClass::Fatal => {
+            return PollDecision::Fail(JimengError::Api {
+                code: resp.code,
+                message: resp
+                    .message
+                    .clone()
+                    .unwrap_or_else(|| "unknown error".into()),
+                retryable: false,
+            });
+        }
+        CodeClass::Success => {}
+    }
+
+    match resp.task_status() {
+        TaskStatus::InQueue | TaskStatus::Generating | TaskStatus::Unknown => PollDecision::Wait,
+        TaskStatus::Done => match extract_output(capability, resp.data.as_ref()) {
+            Some(output) => PollDecision::Done(output),
+            None => PollDecision::Fail(JimengError::EmptyOutput),
+        },
+        TaskStatus::NotFound => PollDecision::Fail(JimengError::NotFound),
+        TaskStatus::Expired => PollDecision::Fail(JimengError::Expired),
+    }
+}
+
+/// Extract the finished output from a `done` response, if present.
+fn extract_output(capability: JimengCapability, data: Option<&GetResultData>) -> Option<JimengOutput> {
+    let data = data?;
+    match capability {
+        JimengCapability::Image => {
+            let urls = data.image_urls.clone().unwrap_or_default();
+            let base64 = data.binary_data_base64.clone().unwrap_or_default();
+            if urls.is_empty() && base64.is_empty() {
+                None
+            } else {
+                Some(JimengOutput::Image { urls, base64 })
+            }
+        }
+        JimengCapability::Video => data
+            .video_url
+            .as_ref()
+            .filter(|u| !u.is_empty())
+            .map(|u| JimengOutput::Video { url: u.clone() }),
+    }
+}
+
+/// Maximum backoff multiplier applied to the base interval on repeated
+/// retryable codes (exponential, capped).
+const MAX_BACKOFF_MULTIPLIER: u32 = 8;
+
+/// Drive the poll loop until the task reaches a terminal state, the timeout
+/// elapses, or the task is cancelled.
+///
+/// `fetch` is called to obtain each [`GetResultResponse`]; `is_cancelled` is
+/// checked before every poll. Both are async closures so the route supplies
+/// real provider/database calls while tests supply mocks. `sleep` performs the
+/// inter-poll delay (tests can pass a no-op).
+pub async fn poll_until_done<F, Fut, C, CFut, S, SFut>(
+    capability: JimengCapability,
+    interval: Duration,
+    timeout: Duration,
+    mut fetch: F,
+    mut is_cancelled: C,
+    mut sleep: S,
+) -> Result<JimengOutput, JimengError>
+where
+    F: FnMut() -> Fut,
+    Fut: Future<Output = Result<GetResultResponse, JimengError>>,
+    C: FnMut() -> CFut,
+    CFut: Future<Output = bool>,
+    S: FnMut(Duration) -> SFut,
+    SFut: Future<Output = ()>,
+{
+    let deadline = Instant::now() + timeout;
+    let mut backoff: u32 = 1;
+    loop {
+        if is_cancelled().await {
+            return Err(JimengError::Cancelled);
+        }
+        if Instant::now() >= deadline {
+            return Err(JimengError::Timeout(timeout.as_secs()));
+        }
+
+        let resp = fetch().await?;
+        match decide(capability, &resp) {
+            PollDecision::Done(output) => return Ok(output),
+            PollDecision::Fail(err) => return Err(err),
+            PollDecision::Wait => {
+                backoff = 1;
+                sleep(interval).await;
+            }
+            PollDecision::Backoff(_) => {
+                let delay = interval.saturating_mul(backoff);
+                backoff = (backoff * 2).min(MAX_BACKOFF_MULTIPLIER);
+                sleep(delay).await;
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::cell::RefCell;
+
+    fn resp(code: i64, status: Option<&str>, video_url: Option<&str>) -> GetResultResponse {
+        let data = if code == 10000 {
+            Some(GetResultData {
+                status: status.map(str::to_string),
+                image_urls: None,
+                binary_data_base64: None,
+                video_url: video_url.map(str::to_string),
+            })
+        } else {
+            None
+        };
+        GetResultResponse {
+            code,
+            message: Some("m".into()),
+            data,
+        }
+    }
+
+    fn image_done() -> GetResultResponse {
+        GetResultResponse {
+            code: 10000,
+            message: None,
+            data: Some(GetResultData {
+                status: Some("done".into()),
+                image_urls: Some(vec!["https://x/a.png".into()]),
+                binary_data_base64: None,
+                video_url: None,
+            }),
+        }
+    }
+
+    #[test]
+    fn decide_waits_while_generating() {
+        let r = resp(10000, Some("generating"), None);
+        assert_eq!(decide(JimengCapability::Image, &r), PollDecision::Wait);
+        let r = resp(10000, Some("in_queue"), None);
+        assert_eq!(decide(JimengCapability::Image, &r), PollDecision::Wait);
+    }
+
+    #[test]
+    fn decide_done_image_returns_output() {
+        let r = image_done();
+        match decide(JimengCapability::Image, &r) {
+            PollDecision::Done(JimengOutput::Image { urls, .. }) => {
+                assert_eq!(urls, vec!["https://x/a.png".to_string()]);
+            }
+            other => panic!("expected image output, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn decide_done_video_returns_url() {
+        let r = resp(10000, Some("done"), Some("https://x/v.mp4"));
+        assert_eq!(
+            decide(JimengCapability::Video, &r),
+            PollDecision::Done(JimengOutput::Video {
+                url: "https://x/v.mp4".into()
+            })
+        );
+    }
+
+    #[test]
+    fn decide_done_without_output_fails_empty() {
+        let r = resp(10000, Some("done"), None);
+        assert!(matches!(
+            decide(JimengCapability::Video, &r),
+            PollDecision::Fail(JimengError::EmptyOutput)
+        ));
+    }
+
+    #[test]
+    fn decide_expired_and_not_found_fail() {
+        assert!(matches!(
+            decide(JimengCapability::Image, &resp(10000, Some("expired"), None)),
+            PollDecision::Fail(JimengError::Expired)
+        ));
+        assert!(matches!(
+            decide(JimengCapability::Image, &resp(10000, Some("not_found"), None)),
+            PollDecision::Fail(JimengError::NotFound)
+        ));
+    }
+
+    #[test]
+    fn decide_retryable_code_backs_off() {
+        // 50429 = QPS limit, retryable.
+        assert_eq!(
+            decide(JimengCapability::Image, &resp(50429, None, None)),
+            PollDecision::Backoff(50429)
+        );
+    }
+
+    #[test]
+    fn decide_fatal_code_fails() {
+        // 50411 = input moderation, fatal.
+        assert!(matches!(
+            decide(JimengCapability::Image, &resp(50411, None, None)),
+            PollDecision::Fail(JimengError::Api { code: 50411, .. })
+        ));
+    }
+
+    #[test]
+    fn decide_unknown_status_waits() {
+        // code ok but unrecognized status: keep waiting (timeout bounds it).
+        let r = resp(10000, Some("weird"), None);
+        assert_eq!(decide(JimengCapability::Image, &r), PollDecision::Wait);
+    }
+
+    // --- async driver tests (single-threaded, no real sleeping) ---
+
+    async fn run_sequence(
+        capability: JimengCapability,
+        seq: Vec<GetResultResponse>,
+        timeout: Duration,
+    ) -> Result<JimengOutput, JimengError> {
+        let idx = RefCell::new(0usize);
+        let seq = RefCell::new(seq);
+        poll_until_done(
+            capability,
+            Duration::from_millis(0),
+            timeout,
+            || async {
+                let mut i = idx.borrow_mut();
+                let s = seq.borrow();
+                let item = s.get(*i).cloned().unwrap_or_else(|| {
+                    // Exhausted: keep returning the last to avoid panics.
+                    s.last().cloned().unwrap()
+                });
+                *i += 1;
+                Ok(item)
+            },
+            || async { false },
+            |_d| async {},
+        )
+        .await
+    }
+
+    #[tokio::test]
+    async fn driver_polls_until_done() {
+        let seq = vec![
+            resp(10000, Some("in_queue"), None),
+            resp(10000, Some("generating"), None),
+            resp(10000, Some("generating"), None),
+            image_done(),
+        ];
+        let out = run_sequence(JimengCapability::Image, seq, Duration::from_secs(30))
+            .await
+            .unwrap();
+        assert!(matches!(out, JimengOutput::Image { .. }));
+    }
+
+    #[tokio::test]
+    async fn driver_backoff_then_done() {
+        let seq = vec![
+            resp(50429, None, None), // retryable
+            resp(10000, Some("generating"), None),
+            resp(10000, Some("done"), Some("https://x/v.mp4")),
+        ];
+        let out = run_sequence(JimengCapability::Video, seq, Duration::from_secs(30))
+            .await
+            .unwrap();
+        assert_eq!(
+            out,
+            JimengOutput::Video {
+                url: "https://x/v.mp4".into()
+            }
+        );
+    }
+
+    #[tokio::test]
+    async fn driver_expired_fails() {
+        let seq = vec![
+            resp(10000, Some("generating"), None),
+            resp(10000, Some("expired"), None),
+        ];
+        let err = run_sequence(JimengCapability::Image, seq, Duration::from_secs(30))
+            .await
+            .unwrap_err();
+        assert!(matches!(err, JimengError::Expired));
+    }
+
+    #[tokio::test]
+    async fn driver_times_out() {
+        // Zero timeout: deadline is now, loop returns Timeout on first check.
+        let seq = vec![resp(10000, Some("generating"), None)];
+        let err = run_sequence(JimengCapability::Image, seq, Duration::from_secs(0))
+            .await
+            .unwrap_err();
+        assert!(matches!(err, JimengError::Timeout(0)));
+    }
+
+    #[tokio::test]
+    async fn driver_honors_cancellation() {
+        let idx = RefCell::new(0usize);
+        let err = poll_until_done(
+            JimengCapability::Image,
+            Duration::from_millis(0),
+            Duration::from_secs(30),
+            || async { Ok(resp(10000, Some("generating"), None)) },
+            || async {
+                let mut i = idx.borrow_mut();
+                *i += 1;
+                *i > 2 // cancel on the third check
+            },
+            |_d| async {},
+        )
+        .await
+        .unwrap_err();
+        assert!(matches!(err, JimengError::Cancelled));
+    }
+}