a3s 0.7.4

//! The `runtime` tool — A3S Runtime offload for the a3s-code TUI.
//!
//! Registered into the session **only after the user logs in to OS (OS)**
//! (see `Tui::sync_runtime_tool`), so the model never sees it while signed out.
//!
//! When the model calls it, it fans the given subtasks out to OS's
//! Function-as-a-Service batch API (`POST /api/v1/functions/<worker>/batch`) for
//! **parallel** execution on the OS pod substrate, polls the batch to completion,
//! collects each invocation's output, and returns the aggregated results — so a
//! decomposed job (e.g. deep-research sub-questions) runs concurrently on remote
//! compute instead of serially in-process. See `docs/function-compute.md` in the
//! OS repo for the wire contract.
//!
//! Mechanics (each live-validated against the deployed OS):
//! - `worker` accepts a tool-kind agent asset **UUID**, or an asset **name**
//!   which is resolved via the assets API (the OS itself only takes UUIDs —
//!   a raw name fails server-side with `RUNNER_ERROR: invalid input syntax
//!   for type uuid`). Unknown names fail fast listing the available workers.
//! - Polling streams live progress into the TUI (`ToolStreamEvent`), backs off
//!   exponentially (1.5s → 6s cap), and tolerates transient network errors.
//! - On completion, results are collected **concurrently**; on budget expiry
//!   the finished subset is still returned (`partial: true` + `batchId`).

use a3s_code_core::tools::{Tool, ToolContext, ToolOutput, ToolStreamEvent};
use anyhow::Result;
use async_trait::async_trait;
use serde_json::{json, Value};
use std::time::Duration;

use crate::a3s_os::{os_origin, StoredOsSession};

const HTTP_TIMEOUT: Duration = Duration::from_secs(30);
/// Overall wait cap for a batch to finish (ceiling; callers can lower via `timeout_ms`).
const DEFAULT_BATCH_TIMEOUT_MS: u64 = 600_000;
/// Poll backoff: start fast (short batches return quickly), then ease off so a
/// 10-minute batch costs ~110 polls instead of 400.
const POLL_START: Duration = Duration::from_millis(1500);
const POLL_CAP: Duration = Duration::from_millis(6000);
/// Consecutive poll failures tolerated before giving up — one flaky HTTP tick
/// must not abandon an entire running batch.
const MAX_POLL_FAILURES: u32 = 3;

pub(crate) struct RuntimeTool {
    /// OS origin (`scheme://host[:port]`), derived from the login session address.
    origin: String,
    /// OS bearer token (the OAuth access token captured at login/refresh).
    token: String,
    /// Poll pacing — fields (not consts) so tests can run with tiny intervals.
    poll_start: Duration,
    poll_cap: Duration,
}

impl RuntimeTool {
    /// Build from the active OS session. Re-created on every login/refresh so the
    /// captured token stays current (the TUI rebuilds the session on auth change).
    pub(crate) fn new(session: &StoredOsSession) -> Self {
        Self {
            origin: os_origin(&session.address),
            token: session.access_token.clone(),
            poll_start: POLL_START,
            poll_cap: POLL_CAP,
        }
    }

    fn client(&self) -> Result<reqwest::Client> {
        let mut builder = reqwest::Client::builder().timeout(HTTP_TIMEOUT);
        if is_loopback_origin(&self.origin) {
            builder = builder.no_proxy();
        }
        Ok(builder.build()?)
    }

    /// Unwrap the shared OS response envelope `{code,status,message,data,...}` and
    /// return `data` (the real payload). Errors carry the wire `message` so the
    /// model sees why a call failed.
    fn unwrap_envelope(body: &str, status: u16) -> Result<Value> {
        let v: Value = serde_json::from_str(body).map_err(|e| {
            anyhow::anyhow!(
                "Non-JSON response (HTTP {status}): {e}: {}",
                truncate(body, 200)
            )
        })?;
        let code = v
            .get("code")
            .and_then(Value::as_u64)
            .unwrap_or(status as u64);
        if code >= 400 || status >= 400 {
            let msg = v
                .get("message")
                .and_then(Value::as_str)
                .unwrap_or("request failed");
            anyhow::bail!("A3S Runtime returned {code}: {msg}");
        }
        Ok(v.get("data").cloned().unwrap_or(v))
    }
}

fn is_loopback_origin(origin: &str) -> bool {
    origin.starts_with("http://127.")
        || origin.starts_with("https://127.")
        || origin.starts_with("http://localhost")
        || origin.starts_with("https://localhost")
        || origin.starts_with("http://[::1]")
        || origin.starts_with("https://[::1]")
}

#[async_trait]
impl Tool for RuntimeTool {
    fn name(&self) -> &str {
        "runtime"
    }

    fn description(&self) -> &str {
        "Offload independent subtasks to OS A3S Runtime for parallel remote \
         execution, stream progress while they run, then return a combined result. \
         Use it for decomposable work such as multiple deep-research subquestions. \
         `worker` is a tool-kind agent asset UUID or name. Names auto-resolve; \
         invalid names list the available workers."
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "tasks": {
                    "type": "array",
                    "description": "Independent subtasks to run in parallel. Each item is passed to the worker as input, usually a subquestion string or an object matching the worker inputSchema.",
                    "items": { "type": ["string", "object"] },
                    "minItems": 1
                },
                "worker": {
                    "type": "string",
                    "description": "Worker that runs each subtask: a tool-kind agent asset UUID or name. Required."
                },
                "timeout_ms": {
                    "type": "integer",
                    "description": "Maximum wait for the full batch in milliseconds. Defaults to 10 minutes; on timeout, returns completed results."
                }
            },
            "required": ["tasks", "worker"]
        })
    }

    async fn execute(&self, args: &Value, ctx: &ToolContext) -> Result<ToolOutput> {
        let tasks: Vec<Value> = match args.get("tasks").and_then(Value::as_array) {
            Some(a) if !a.is_empty() => a.clone(),
            _ => return Ok(ToolOutput::error("`tasks` must be a non-empty array")),
        };
        let worker = match args.get("worker").and_then(Value::as_str) {
            Some(w) if !w.trim().is_empty() => w.trim().to_string(),
            _ => {
                return Ok(ToolOutput::error(
                    "`worker` is required: use a tool-kind agent asset UUID or name",
                ))
            }
        };
        let budget_ms = args
            .get("timeout_ms")
            .and_then(Value::as_u64)
            .unwrap_or(DEFAULT_BATCH_TIMEOUT_MS);

        match self.run_batch(&worker, tasks, budget_ms, ctx).await {
            Ok(out) => Ok(ToolOutput::success(out)),
            // A3S Runtime / network failure is a tool failure, not a crash — surface it.
            Err(e) => Ok(ToolOutput::error(format!(
                "A3S Runtime offload failed: {e}"
            ))),
        }
    }
}

impl RuntimeTool {
    async fn run_batch(
        &self,
        worker: &str,
        tasks: Vec<Value>,
        budget_ms: u64,
        ctx: &ToolContext,
    ) -> Result<String> {
        let client = self.client()?;
        let n = tasks.len();
        let progress = |msg: String| {
            if let Some(tx) = &ctx.event_tx {
                // Best-effort: progress must never block or fail the batch.
                let _ = tx.try_send(ToolStreamEvent::OutputDelta(msg));
            }
        };

        // 0. Resolve a worker NAME to its asset UUID (the OS API only accepts
        //    UUIDs). UUIDs pass straight through.
        let worker_id = if looks_like_uuid(worker) {
            worker.to_string()
        } else {
            let id = self.resolve_worker_name(&client, worker).await?;
            progress(format!("worker {worker} -> {id}\n"));
            id
        };

        // 1. Fan out. idempotencyKey (hash of worker + task set) makes a retry
        //    re-attach to the same batch instead of double-spending.
        let idem = idempotency_key(&worker_id, &tasks);
        let submit_url = format!("{}/api/v1/functions/{}/batch", self.origin, worker_id);
        let resp = client
            .post(&submit_url)
            .bearer_auth(&self.token)
            .json(&json!({ "inputs": tasks, "agentKind": "tool", "idempotencyKey": idem }))
            .send()
            .await?;
        let status = resp.status().as_u16();
        let data = Self::unwrap_envelope(&resp.text().await?, status)?;
        let batch_id = data
            .get("batchId")
            .and_then(Value::as_str)
            .ok_or_else(|| anyhow::anyhow!("batch response did not include batchId"))?
            .to_string();
        let invocation_ids: Vec<String> = data
            .get("invocationIds")
            .and_then(Value::as_array)
            .map(|a| {
                a.iter()
                    .filter_map(|v| v.as_str().map(String::from))
                    .collect()
            })
            .unwrap_or_default();
        progress(format!(
            "{n} parallel subtasks submitted (batch {batch_id})\n"
        ));

        // 2. Poll until every member is terminal or the budget expires — with
        //    exponential backoff, live progress, and transient-failure tolerance.
        let poll_url = format!("{}/api/v1/functions/batches/{}", self.origin, batch_id);
        let mut waited = 0u64;
        let mut interval = self.poll_start;
        let mut consecutive_failures = 0u32;
        let mut last_report = String::new();
        let mut timed_out_pending = 0u64;
        loop {
            let poll = async {
                let resp = client
                    .get(&poll_url)
                    .bearer_auth(&self.token)
                    .send()
                    .await?;
                let status = resp.status().as_u16();
                Self::unwrap_envelope(&resp.text().await?, status)
            }
            .await;
            match poll {
                Ok(bd) => {
                    consecutive_failures = 0;
                    let counts = bd.get("counts").cloned().unwrap_or_else(|| json!({}));
                    let c = |k: &str| counts.get(k).and_then(Value::as_u64).unwrap_or(0);
                    let (queued, running) = (c("queued"), c("running"));
                    let done = c("succeeded") + c("failed") + c("canceled") + c("unknown");
                    let pending = queued + running;
                    // Emit progress only when the picture changes (no spam).
                    let report =
                        format!("⏳ {done}/{n} done · {running} running · {queued} queued\n");
                    if report != last_report {
                        progress(report.clone());
                        last_report = report;
                    }
                    if pending == 0 {
                        break;
                    }
                    if waited >= budget_ms {
                        timed_out_pending = pending;
                        break;
                    }
                }
                Err(e) => {
                    // Tolerate flaky ticks: one failed GET must not abandon a
                    // whole running batch.
                    consecutive_failures += 1;
                    if consecutive_failures > MAX_POLL_FAILURES {
                        return Err(e.context(format!(
                            "polling batch {batch_id} failed {consecutive_failures} consecutive times"
                        )));
                    }
                    progress(format!(
                        "poll failed (attempt {consecutive_failures}; retrying)\n"
                    ));
                }
            }
            tokio::time::sleep(interval).await;
            waited += interval.as_millis() as u64;
            interval = (interval * 3 / 2).min(self.poll_cap);
        }

        // 3. Collect every member's result CONCURRENTLY (one RTT, not N).
        //    On timeout this still runs: finished members' outputs are returned
        //    (partial) instead of being thrown away.
        let fetches = invocation_ids.iter().map(|id| {
            let url = format!("{}/api/v1/functions/invocations/{}", self.origin, id);
            let client = client.clone();
            let token = self.token.clone();
            async move {
                let inv = async {
                    let resp = client.get(&url).bearer_auth(&token).send().await?;
                    let status = resp.status().as_u16();
                    Self::unwrap_envelope(&resp.text().await?, status)
                }
                .await
                .unwrap_or_else(|e| json!({ "error": e.to_string() }));
                let result = inv.get("result").cloned().unwrap_or(inv);
                json!({
                    "invocationId": id,
                    "state": result.get("status").cloned().unwrap_or_else(|| json!("unknown")),
                    "output": result.get("output").cloned().unwrap_or(Value::Null),
                    "error": result.get("error").cloned().unwrap_or(Value::Null),
                })
            }
        });
        let results: Vec<Value> = futures::future::join_all(fetches).await;

        let mut summary = json!({
            "batchId": batch_id,
            "worker": worker_id,
            "count": n,
            "results": results,
        });
        if timed_out_pending > 0 {
            summary["partial"] = json!(true);
            summary["note"] = json!(format!(
                "Timed out after {waited}ms with {timed_out_pending} subtasks still pending; \
                 completed results were returned. Query batchId={batch_id} later for unfinished items."
            ));
        }
        Ok(serde_json::to_string_pretty(&summary)?)
    }

    /// Resolve a worker asset NAME to its UUID via the assets API. Fails with
    /// the list of available tool-kind workers so the model can self-correct.
    async fn resolve_worker_name(&self, client: &reqwest::Client, name: &str) -> Result<String> {
        let url = format!("{}/api/v1/assets?category=agent&limit=100", self.origin);
        let resp = client.get(&url).bearer_auth(&self.token).send().await?;
        let status = resp.status().as_u16();
        let data = Self::unwrap_envelope(&resp.text().await?, status)?;
        let items = data
            .get("items")
            .or_else(|| data.get("list"))
            .and_then(Value::as_array)
            .cloned()
            .unwrap_or_default();
        let tools: Vec<(&str, &str)> = items
            .iter()
            .filter(|a| a.get("agentKind").and_then(Value::as_str) == Some("tool"))
            .filter_map(|a| {
                Some((
                    a.get("id").and_then(Value::as_str)?,
                    a.get("name").and_then(Value::as_str)?,
                ))
            })
            .collect();
        if let Some((id, _)) = tools
            .iter()
            .find(|(_, n)| n.eq_ignore_ascii_case(name.trim()))
        {
            return Ok(id.to_string());
        }
        let available: Vec<String> = tools
            .iter()
            .take(10)
            .map(|(id, n)| format!("{n} ({id})"))
            .collect();
        anyhow::bail!(
            "No tool-kind worker named \"{name}\". Available workers: {}",
            if available.is_empty() {
                "none; create a tool-kind agent asset in the OS Asset Center first".to_string()
            } else {
                available.join(", ")
            }
        )
    }
}

/// A canonical hyphenated UUID (the only `ref` form the OS batch API accepts).
fn looks_like_uuid(s: &str) -> bool {
    let b = s.as_bytes();
    b.len() == 36
        && b.iter().enumerate().all(|(i, c)| match i {
            8 | 13 | 18 | 23 => *c == b'-',
            _ => c.is_ascii_hexdigit(),
        })
}

/// Deterministic idempotency key from the worker + task set (sha256, truncated) —
/// a retry of the same fan-out re-attaches to the existing batch instead of
/// double-spending, while the same tasks on a DIFFERENT worker stay distinct.
fn idempotency_key(worker: &str, tasks: &[Value]) -> String {
    use sha2::{Digest, Sha256};
    let mut h = Sha256::new();
    h.update(worker.as_bytes());
    h.update([0u8]);
    h.update(serde_json::to_vec(tasks).unwrap_or_default());
    let hex: String = h.finalize().iter().map(|b| format!("{:02x}", b)).collect();
    format!("a3s-code-runtime-{}", &hex[..24])
}

fn truncate(s: &str, max: usize) -> String {
    if s.chars().count() <= max {
        s.to_string()
    } else {
        // Char-based to avoid panicking on a multibyte boundary.
        format!("{}…", s.chars().take(max).collect::<String>())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Mutex};
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpListener;

    // ── pure logic ──────────────────────────────────────────────────────────

    #[test]
    fn unwrap_envelope_returns_data_and_surfaces_errors() {
        let ok = RuntimeTool::unwrap_envelope(
            r#"{"code":200,"status":"OK","data":{"batchId":"b1"}}"#,
            200,
        )
        .unwrap();
        assert_eq!(ok.get("batchId").unwrap(), "b1");
        let err =
            RuntimeTool::unwrap_envelope(r#"{"code":403,"message":"Forbidden"}"#, 200).unwrap_err();
        assert!(err.to_string().contains("403") && err.to_string().contains("Forbidden"));
        assert!(RuntimeTool::unwrap_envelope(r#"{"message":"x"}"#, 500).is_err());
        assert!(RuntimeTool::unwrap_envelope("<html>502</html>", 502).is_err());
        let bare = RuntimeTool::unwrap_envelope(r#"{"batchId":"b2"}"#, 200).unwrap();
        assert_eq!(bare.get("batchId").unwrap(), "b2");
    }

    #[test]
    fn idempotency_key_covers_worker_and_task_set() {
        let a = idempotency_key("w1", &[json!("q1"), json!("q2")]);
        assert_eq!(a, idempotency_key("w1", &[json!("q1"), json!("q2")]));
        // Same tasks on a DIFFERENT worker must be a different batch.
        assert_ne!(a, idempotency_key("w2", &[json!("q1"), json!("q2")]));
        assert_ne!(a, idempotency_key("w1", &[json!("q2"), json!("q1")]));
        assert!(a.starts_with("a3s-code-runtime-") && a.len() == "a3s-code-runtime-".len() + 24);
    }

    #[test]
    fn uuid_detection_is_strict() {
        assert!(looks_like_uuid("57989959-0b1d-41da-974c-31ad8101df37"));
        assert!(!looks_like_uuid("risk-reporter"));
        assert!(!looks_like_uuid("57989959-0b1d-41da-974c-31ad8101df3")); // 35 chars
        assert!(!looks_like_uuid("g7989959-0b1d-41da-974c-31ad8101df37")); // non-hex
    }

    // ── mock A3S Runtime speaking the exact OS contract ─────────────────────

    /// Scripted mock state: each poll consumes the next `poll_plan` item and
    /// repeats the final item after the plan is exhausted.
    struct MockState {
        submit_path: Option<String>,
        submit_body: Option<String>,
        /// Each entry: Some(counts json) → 200 with those counts; None → HTTP 500.
        poll_plan: Vec<Option<String>>,
        poll_idx: usize,
    }

    async fn spawn_mock(state: Arc<Mutex<MockState>>) -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let origin = format!("http://{}", listener.local_addr().unwrap());
        tokio::spawn(async move {
            loop {
                let Ok((mut sock, _)) = listener.accept().await else {
                    return;
                };
                let st = state.clone();
                tokio::spawn(async move {
                    let mut buf = vec![0u8; 16384];
                    let n = sock.read(&mut buf).await.unwrap_or(0);
                    let req = String::from_utf8_lossy(&buf[..n]).into_owned();
                    let line = req.lines().next().unwrap_or("").to_string();
                    let body = req.split("\r\n\r\n").nth(1).unwrap_or("").to_string();
                    let (status, payload) = route(&st, &line, &body);
                    let resp = format!(
                        "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{payload}",
                        payload.len()
                    );
                    let _ = sock.write_all(resp.as_bytes()).await;
                    let _ = sock.flush().await;
                });
            }
        });
        origin
    }

    fn route(st: &Arc<Mutex<MockState>>, line: &str, body: &str) -> (&'static str, String) {
        let env = |data: &str| format!(r#"{{"code":200,"status":"OK","data":{data}}}"#);
        // Specific paths BEFORE the generic "/batch" (substring overlap).
        if line.contains("/api/v1/assets") {
            return (
                "200 OK",
                env(r#"{"items":[
                    {"id":"57989959-0b1d-41da-974c-31ad8101df37","name":"risk-reporter","agentKind":"tool"},
                    {"id":"74af5078-7b53-4857-bf69-fc59c9fdce06","name":"shangfei-poc3","agentKind":"tool"},
                    {"id":"02f3b08e-9358-43aa-97c3-05981b57a1a2","name":"some-app","agentKind":"application"}
                ]}"#),
            );
        }
        if line.contains("/batches/") {
            let mut s = st.lock().unwrap();
            let i = s.poll_idx.min(s.poll_plan.len().saturating_sub(1));
            s.poll_idx += 1;
            return match s.poll_plan.get(i).cloned().flatten() {
                Some(counts) => (
                    "200 OK",
                    env(&format!(r#"{{"batchId":"batch-1","counts":{counts}}}"#)),
                ),
                None => ("500 Internal Server Error", "boom".to_string()),
            };
        }
        if line.contains("/invocations/inv-1") {
            return (
                "200 OK",
                env(
                    r#"{"status":"succeeded","result":{"status":"succeeded","output":{"answer":"alpha"},"error":null}}"#,
                ),
            );
        }
        if line.contains("/invocations/") {
            return ("200 OK", env(r#"{"status":"running","result":null}"#));
        }
        if line.contains("/batch") {
            let mut s = st.lock().unwrap();
            s.submit_path = Some(line.split_whitespace().nth(1).unwrap_or("").to_string());
            s.submit_body = Some(body.to_string());
            return (
                "200 OK",
                env(r#"{"batchId":"batch-1","invocationIds":["inv-1","inv-2"]}"#),
            );
        }
        ("404 Not Found", "{}".to_string())
    }

    fn fast_tool(origin: String) -> RuntimeTool {
        RuntimeTool {
            origin,
            token: "test-token".into(),
            poll_start: Duration::from_millis(10),
            poll_cap: Duration::from_millis(20),
        }
    }

    fn state(poll_plan: Vec<Option<&str>>) -> Arc<Mutex<MockState>> {
        Arc::new(Mutex::new(MockState {
            submit_path: None,
            submit_body: None,
            poll_plan: poll_plan.into_iter().map(|p| p.map(String::from)).collect(),
            poll_idx: 0,
        }))
    }

    /// A ToolContext with a progress channel; returns (ctx, drained-events fn).
    fn ctx_with_progress() -> (ToolContext, tokio::sync::mpsc::Receiver<ToolStreamEvent>) {
        let (tx, rx) = tokio::sync::mpsc::channel(64);
        let ctx = ToolContext::new(std::env::temp_dir()).with_event_tx(tx);
        (ctx, rx)
    }

    #[tokio::test]
    async fn full_flow_streams_progress_and_aggregates() {
        // Two live ticks then terminal — exercises backoff + change-only progress.
        let st = state(vec![
            Some(r#"{"queued":1,"running":1,"succeeded":0,"failed":0}"#),
            Some(r#"{"queued":0,"running":1,"succeeded":1,"failed":0}"#),
            Some(r#"{"queued":0,"running":0,"succeeded":2,"failed":0}"#),
        ]);
        let origin = spawn_mock(st.clone()).await;
        let tool = fast_tool(origin);
        let (ctx, mut rx) = ctx_with_progress();
        let out = tool
            .execute(
                &json!({ "tasks": ["a", "b"], "worker": "57989959-0b1d-41da-974c-31ad8101df37" }),
                &ctx,
            )
            .await
            .unwrap();
        assert!(out.success, "{}", out.content);

        // Request contract: inputs + agentKind + worker-scoped idempotency key.
        let sent: Value =
            serde_json::from_str(st.lock().unwrap().submit_body.as_ref().unwrap()).unwrap();
        assert_eq!(sent["inputs"], json!(["a", "b"]));
        assert_eq!(sent["agentKind"], "tool");
        assert_eq!(
            sent["idempotencyKey"].as_str().unwrap(),
            idempotency_key(
                "57989959-0b1d-41da-974c-31ad8101df37",
                &[json!("a"), json!("b")]
            )
        );

        // Progress streamed: submit line + one line per distinct counts picture.
        let mut deltas = Vec::new();
        while let Ok(ev) = rx.try_recv() {
            let ToolStreamEvent::OutputDelta(s) = ev;
            deltas.push(s);
        }
        let all = deltas.join("");
        assert!(all.contains("2 parallel subtasks submitted"), "{all}");
        assert!(
            all.contains("0/2 done") && all.contains("2/2 done"),
            "{all}"
        );

        // Aggregation: both results, in invocation order.
        let agg: Value = serde_json::from_str(&out.content).unwrap();
        assert_eq!(agg["count"], 2);
        assert_eq!(agg["results"][0]["output"]["answer"], "alpha");
        assert!(
            agg.get("partial").is_none(),
            "terminal batch is not partial"
        );
    }

    #[tokio::test]
    async fn transient_poll_failure_is_tolerated_but_persistent_is_not() {
        // One 500 tick between two good ones → still succeeds.
        let st = state(vec![
            Some(r#"{"queued":0,"running":1,"succeeded":1,"failed":0}"#),
            None, // 500
            Some(r#"{"queued":0,"running":0,"succeeded":2,"failed":0}"#),
        ]);
        let origin = spawn_mock(st).await;
        let tool = fast_tool(origin);
        let (ctx, _rx) = ctx_with_progress();
        let out = tool
            .execute(
                &json!({ "tasks": ["a", "b"], "worker": "57989959-0b1d-41da-974c-31ad8101df37" }),
                &ctx,
            )
            .await
            .unwrap();
        assert!(out.success, "one flaky tick must not abandon the batch");

        // 4+ consecutive failures → gives up with the poll error surfaced.
        let st2 = state(vec![
            Some(r#"{"queued":0,"running":2,"succeeded":0,"failed":0}"#),
            None,
            None,
            None,
            None,
        ]);
        let origin2 = spawn_mock(st2).await;
        let tool2 = fast_tool(origin2);
        let (ctx2, _rx2) = ctx_with_progress();
        let out2 = tool2
            .execute(
                &json!({ "tasks": ["a", "b"], "worker": "57989959-0b1d-41da-974c-31ad8101df37" }),
                &ctx2,
            )
            .await
            .unwrap();
        assert!(!out2.success);
        assert!(
            out2.content.contains("consecutive times"),
            "{}",
            out2.content
        );
    }

    #[tokio::test]
    async fn timeout_returns_partial_results_not_nothing() {
        // Batch never finishes; budget expires after the first tick. The
        // finished member (inv-1) must still come back, flagged partial.
        let st = state(vec![Some(
            r#"{"queued":0,"running":1,"succeeded":1,"failed":0}"#,
        )]);
        let origin = spawn_mock(st).await;
        let tool = fast_tool(origin);
        let (ctx, _rx) = ctx_with_progress();
        let out = tool
            .execute(
                &json!({ "tasks": ["a", "b"], "worker": "57989959-0b1d-41da-974c-31ad8101df37", "timeout_ms": 1 }),
                &ctx,
            )
            .await
            .unwrap();
        assert!(out.success, "{}", out.content);
        let agg: Value = serde_json::from_str(&out.content).unwrap();
        assert_eq!(agg["partial"], true);
        assert!(agg["note"].as_str().unwrap().contains("batch-1"));
        assert_eq!(agg["results"][0]["output"]["answer"], "alpha"); // finished one kept
        assert_eq!(agg["results"][1]["state"], "unknown"); // unfinished: no result yet
    }

    #[tokio::test]
    async fn worker_name_resolves_to_uuid_and_unknown_names_list_options() {
        let st = state(vec![Some(
            r#"{"queued":0,"running":0,"succeeded":2,"failed":0}"#,
        )]);
        let origin = spawn_mock(st.clone()).await;
        let tool = fast_tool(origin.clone());
        let (ctx, _rx) = ctx_with_progress();
        // Name → UUID: the submit URL must target the resolved asset id.
        let out = tool
            .execute(
                &json!({ "tasks": ["a", "b"], "worker": "risk-reporter" }),
                &ctx,
            )
            .await
            .unwrap();
        assert!(out.success, "{}", out.content);
        let path = st.lock().unwrap().submit_path.clone().unwrap();
        assert!(
            path.contains("/functions/57989959-0b1d-41da-974c-31ad8101df37/batch"),
            "{path}"
        );

        // Unknown name → error listing available tool workers (not applications).
        let (ctx2, _rx2) = ctx_with_progress();
        let out2 = tool
            .execute(
                &json!({ "tasks": ["a"], "worker": "no-such-worker" }),
                &ctx2,
            )
            .await
            .unwrap();
        assert!(!out2.success);
        assert!(out2.content.contains("risk-reporter"), "{}", out2.content);
        assert!(
            !out2.content.contains("some-app"),
            "application-kind assets are not workers"
        );
    }

    #[tokio::test]
    async fn bad_args_are_tool_errors_not_requests() {
        let tool = fast_tool("http://127.0.0.1:1".into()); // never contacted
        let ctx = ToolContext::new(std::env::temp_dir());
        let e1 = tool
            .execute(&json!({ "tasks": [], "worker": "w" }), &ctx)
            .await
            .unwrap();
        assert!(!e1.success);
        let e2 = tool
            .execute(&json!({ "tasks": ["x"] }), &ctx)
            .await
            .unwrap();
        assert!(!e2.success && e2.content.contains("worker"));
    }
}