flow_execution/

ai_tools.rs

//! AI-node execution support: exposing bound adapters to a model as callable
//! tools, dispatching the model's tool calls back to those adapters, and
//! resolving image references for vision calls.
//!
//! This is the in-node tool-use path - distinct from whole-flow agentic
//! generation. The model emits tool calls, [`AdapterToolDispatcher`] runs each
//! through the *existing* adapter (so the fs/shell sandboxes apply unchanged),
//! and the result is fed back to the model.

use std::sync::Arc;

use chrono::Utc;
use flow_adapter_ai::{ToolDispatcher, ToolSpec};
use flow_domain::graph::{FlowNode, Position};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};

use crate::adapter::{AdapterRegistry, FieldType};
use crate::events::{EventSink, ExecutionEvent};

/// Separator between adapter name and action id in a tool name
/// (`fs__read-file`). `__` avoids clashing with the `-` inside action ids.
const SEP: &str = "__";

/// Default cap on the in-node tool-use loop (how many times the model may call
/// tools before the run gives up). Overridable per node via `maxToolIters`
/// ([`ai_max_tool_iters`]).
pub const DEFAULT_TOOL_ITERATIONS: usize = 8;

/// Observes each model tool call so the dispatcher can emit an
/// [`ExecutionEvent::ToolCall`] just before the tool runs - giving the UI live
/// per-tool-call granularity during an AI node's tool loop (the dispatcher
/// otherwise emits nothing per call). Cheap to clone (an `Arc` + two ids).
#[derive(Clone)]
pub struct ToolObserver {
    sink: Arc<dyn EventSink>,
    execution_id: String,
    node_id: String,
}

impl ToolObserver {
    pub fn new(sink: Arc<dyn EventSink>, execution_id: String, node_id: String) -> Self {
        Self {
            sink,
            execution_id,
            node_id,
        }
    }

    /// Emit a `ToolCall` event for the call about to run.
    fn observe(&self, name: &str, args: &Value) {
        self.sink.emit(ExecutionEvent::ToolCall {
            execution_id: self.execution_id.clone(),
            node_id: self.node_id.clone(),
            at: Utc::now(),
            name: name.to_string(),
            summary: summarize_tool_call(name, args),
        });
    }
}

/// One-line, human-readable summary of a model tool call for the `ToolCall`
/// event: the action id plus the most salient argument (path / command /
/// pattern / query / url), e.g. `edit-file src/foo.rs`. The salient value is
/// truncated so a giant `content`/`command` can't bloat a log line.
fn summarize_tool_call(name: &str, args: &Value) -> String {
    let action = name.split_once(SEP).map(|(_, a)| a).unwrap_or(name);
    let detail = ["path", "command", "args", "pattern", "query", "url"]
        .iter()
        .find_map(|k| args.get(*k).and_then(|v| v.as_str()))
        .map(str::trim)
        .filter(|s| !s.is_empty());
    match detail {
        Some(d) => {
            let mut short: String = d.chars().take(80).collect();
            if short.len() < d.len() {
                short.push('…');
            }
            format!("{action} {short}")
        }
        None => action.to_string(),
    }
}

/// Fields the dispatcher injects (routing + sandbox root), so they're omitted
/// from the model-facing tool schema.
const INJECTED: &[&str] = &["workspaceRoot", "cwd", "adapter", "actionId"];

/// Build model tool specs from the descriptors of the bound adapters. Each
/// adapter action becomes one tool `"{adapter}__{actionId}"`; its required +
/// optional fields become JSON-schema properties.
pub fn build_tool_specs(adapters: &AdapterRegistry, bound: &[String]) -> Vec<ToolSpec> {
    let mut specs = Vec::new();
    for adapter_name in bound {
        let Some(adapter) = adapters.get(adapter_name) else {
            continue;
        };
        let desc = adapter.descriptor();
        for action in &desc.actions {
            let mut props = Map::new();
            let mut required = Vec::new();
            for f in action
                .required_fields
                .iter()
                .chain(action.optional_fields.iter())
            {
                if INJECTED.contains(&f.name.as_str()) {
                    continue;
                }
                props.insert(
                    f.name.clone(),
                    json!({ "type": json_type(f.value_type), "description": f.description }),
                );
            }
            for f in &action.required_fields {
                if INJECTED.contains(&f.name.as_str()) {
                    continue;
                }
                required.push(Value::String(f.name.clone()));
            }
            specs.push(ToolSpec {
                name: format!("{}{SEP}{}", desc.name, action.id),
                description: action.summary.clone(),
                parameters: json!({
                    "type": "object",
                    "properties": props,
                    "required": required,
                }),
            });
        }
    }
    specs
}

fn json_type(t: FieldType) -> &'static str {
    match t {
        FieldType::Number => "number",
        FieldType::Bool => "boolean",
        FieldType::String | FieldType::Null => "string",
    }
}

/// Runs a model tool call by mapping `"{adapter}__{actionId}"` back to an
/// adapter action: it synthesizes a node from the model's args (plus the AI
/// node's workspace as `workspaceRoot`/`cwd`) and executes it through the bound
/// adapter, so all existing sandboxing/policy applies.
pub struct AdapterToolDispatcher {
    adapters: Arc<AdapterRegistry>,
    workspace: Option<String>,
    observer: Option<ToolObserver>,
}

impl AdapterToolDispatcher {
    pub fn new(adapters: Arc<AdapterRegistry>, workspace: Option<String>) -> Self {
        Self {
            adapters,
            workspace,
            observer: None,
        }
    }

    /// Attach a [`ToolObserver`] so each dispatched call emits a `ToolCall`
    /// event for live UI granularity.
    pub fn with_observer(mut self, observer: ToolObserver) -> Self {
        self.observer = Some(observer);
        self
    }
}

#[async_trait::async_trait]
impl ToolDispatcher for AdapterToolDispatcher {
    async fn call(&self, name: &str, args: &Value) -> Result<Value, String> {
        if let Some(o) = &self.observer {
            o.observe(name, args);
        }
        let (adapter_name, action_id) = name
            .split_once(SEP)
            .ok_or_else(|| format!("malformed tool name '{name}'"))?;
        let adapter = self
            .adapters
            .get(adapter_name)
            .ok_or_else(|| format!("adapter '{adapter_name}' not registered"))?;

        let mut data: Map<String, Value> = match args {
            Value::Object(m) => m.clone(),
            _ => Map::new(),
        };
        data.insert("adapter".into(), json!(adapter_name));
        data.insert("actionId".into(), json!(action_id));
        if let Some(ws) = &self.workspace {
            data.entry("workspaceRoot".to_string())
                .or_insert_with(|| json!(ws));
            data.entry("cwd".to_string()).or_insert_with(|| json!(ws));
        }

        let node = FlowNode {
            id: format!("tool-{adapter_name}-{action_id}"),
            node_type: "action".into(),
            position: Position { x: 0.0, y: 0.0 },
            data: Value::Object(data),
        };

        match adapter.execute(&node).await {
            Ok(out) => Ok(serde_json::to_value(out).unwrap_or(Value::Null)),
            Err(e) => Err(e.to_string()),
        }
    }
}

// ---------- staged (review-before-write) tool dispatch ----------

/// A single filesystem change the agent proposed during a turn, computed but
/// **not** applied. The IDE renders each as a diff (`before` → `after`) the
/// user accepts or rejects.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProposedEdit {
    /// Path relative to the workspace root.
    pub path: String,
    /// `"create"`, `"modify"`, or `"delete"`, derived from before/after.
    pub kind: String,
    /// File contents before the turn (`None` if the file did not exist).
    pub before: Option<String>,
    /// File contents after the proposed change (`None` if deleted).
    pub after: Option<String>,
}

/// Net staged edits for one agent turn: the per-path result plus first-touch
/// order, so the IDE shows one stable diff per file.
#[derive(Default)]
pub struct StagingState {
    order: Vec<String>,
    by_path: std::collections::HashMap<String, ProposedEdit>,
}

impl StagingState {
    /// The proposed edits in first-touch order, one per path.
    pub fn edits(&self) -> Vec<ProposedEdit> {
        self.order
            .iter()
            .filter_map(|p| self.by_path.get(p).cloned())
            .collect()
    }

    /// The staged result for `path` if the turn has touched it: `Some(Some(c))`
    /// = staged content, `Some(None)` = staged-deleted, `None` = not staged.
    fn staged(&self, path: &str) -> Option<Option<String>> {
        self.by_path.get(path).map(|e| e.after.clone())
    }

    fn record(&mut self, path: &str, before: Option<String>, after: Option<String>) {
        // Keep the *original* `before` across repeated edits to the same file.
        let before = match self.by_path.get(path) {
            Some(existing) => existing.before.clone(),
            None => {
                self.order.push(path.to_string());
                before
            }
        };
        let kind = if after.is_none() {
            "delete"
        } else if before.is_none() {
            "create"
        } else {
            "modify"
        };
        self.by_path.insert(
            path.to_string(),
            ProposedEdit {
                path: path.to_string(),
                kind: kind.to_string(),
                before,
                after,
            },
        );
    }
}

/// Shared handle to a turn's [`StagingState`]. `FlowApp::run_agent_turn`
/// creates one, the executor hands it to the dispatcher, and the caller reads
/// the proposed edits back out after the run.
pub type StagedEdits = Arc<std::sync::Mutex<StagingState>>;

/// Tool dispatcher for a coding-agent *turn*: filesystem **mutations**
/// (`fs` `write-file`/`edit-file`/`delete-file`) are staged into a buffer
/// instead of written to disk, so the turn can be reviewed before it lands.
/// Everything else - `fs` reads, `glob`/`grep`/`list-dir`, and every `shell`
/// tool - runs for real through [`AdapterToolDispatcher`]. A `read-file` of a
/// staged path is served from the overlay, so the model sees its own pending
/// edits within the same turn.
pub struct StagingToolDispatcher {
    inner: AdapterToolDispatcher,
    state: StagedEdits,
    observer: Option<ToolObserver>,
}

impl StagingToolDispatcher {
    pub fn new(
        adapters: Arc<AdapterRegistry>,
        workspace: Option<String>,
        state: StagedEdits,
    ) -> Self {
        // `inner` deliberately has no observer: this (outer) dispatcher emits
        // one `ToolCall` per model call at the top of `call`, so delegated and
        // internal-read passthroughs to `inner` must not emit again.
        Self {
            inner: AdapterToolDispatcher::new(adapters, workspace),
            state,
            observer: None,
        }
    }

    /// Attach a [`ToolObserver`] so each staged-or-delegated call emits a
    /// `ToolCall` event for live UI granularity.
    pub fn with_observer(mut self, observer: ToolObserver) -> Self {
        self.observer = Some(observer);
        self
    }

    /// Current effective content of `path`: the staged overlay if the turn has
    /// already touched it, otherwise what's on disk. `None` means absent
    /// everywhere (including staged-deleted).
    async fn effective_content(&self, path: &str) -> Option<String> {
        if let Ok(st) = self.state.lock() {
            if let Some(staged) = st.staged(path) {
                return staged;
            }
        }
        self.read_disk(path).await
    }

    /// Read `path` from disk through the real fs adapter (so the workspace jail
    /// applies). `None` on any error (missing / unreadable).
    async fn read_disk(&self, path: &str) -> Option<String> {
        match self.inner.call("fs__read-file", &json!({ "path": path })).await {
            Ok(v) => v
                .get("payload")
                .and_then(|p| p.get("content"))
                .and_then(|c| c.as_str())
                .map(str::to_string),
            Err(_) => None,
        }
    }

    fn record(&self, path: &str, before: Option<String>, after: Option<String>) {
        if let Ok(mut st) = self.state.lock() {
            st.record(path, before, after);
        }
    }

    async fn stage_write(&self, args: &Value) -> Result<Value, String> {
        let path = arg_str(args, "path")?;
        let content = arg_str(args, "content")?;
        let before = self.read_disk(&path).await;
        self.record(&path, before, Some(content.clone()));
        Ok(staged_ok(
            "write-file",
            &path,
            json!({ "bytesWritten": content.len() }),
        ))
    }

    async fn stage_edit(&self, args: &Value) -> Result<Value, String> {
        let path = arg_str(args, "path")?;
        let old = arg_str(args, "oldString")?;
        let new = arg_str(args, "newString")?;
        let replace_all = args
            .get("replaceAll")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        // Edit the current effective content (overlay first, then disk), so a
        // file created/edited earlier in the same turn edits correctly.
        let Some(base) = self.effective_content(&path).await else {
            return Err(format!("oldString not found in {path}"));
        };
        let count = base.matches(&old).count();
        if count == 0 {
            return Err(format!("oldString not found in {path}"));
        }
        if count > 1 && !replace_all {
            return Err(format!(
                "oldString occurs {count}× in {path}; pass replaceAll: true or add context to make it unique"
            ));
        }
        let updated = if replace_all {
            base.replace(&old, &new)
        } else {
            base.replacen(&old, &new, 1)
        };
        // `record` keeps the original `before` on repeat touches, so passing
        // the on-disk state here is correct for a first-touch edit and ignored
        // for a file already created/edited this turn.
        let before = self.read_disk(&path).await;
        self.record(&path, before, Some(updated));
        Ok(staged_ok(
            "edit-file",
            &path,
            json!({ "replacements": if replace_all { count } else { 1 } }),
        ))
    }

    async fn stage_delete(&self, args: &Value) -> Result<Value, String> {
        let path = arg_str(args, "path")?;
        if self.effective_content(&path).await.is_none() {
            return Err(format!("{path}: no such file"));
        }
        let before = self.read_disk(&path).await;
        self.record(&path, before, None);
        Ok(staged_ok("delete-file", &path, json!({})))
    }

    /// Serve a `read-file` from the overlay when the path is staged; `None`
    /// falls through to disk.
    fn overlay_read(&self, args: &Value) -> Option<Result<Value, String>> {
        let path = args.get("path").and_then(|v| v.as_str())?.to_string();
        let staged = self.state.lock().ok()?.staged(&path)?;
        Some(match staged {
            Some(content) => Ok(json!({
                "status": "succeeded",
                "payload": {
                    "action": "read-file",
                    "path": path,
                    "bytes": content.len(),
                    "content": content,
                    "staged": true,
                }
            })),
            None => Err(format!("read {path}: no such file (staged delete)")),
        })
    }
}

#[async_trait::async_trait]
impl ToolDispatcher for StagingToolDispatcher {
    async fn call(&self, name: &str, args: &Value) -> Result<Value, String> {
        if let Some(o) = &self.observer {
            o.observe(name, args);
        }
        if let Some((adapter, action)) = name.split_once(SEP) {
            if adapter == "fs" {
                match action {
                    "write-file" => return self.stage_write(args).await,
                    "edit-file" => return self.stage_edit(args).await,
                    "delete-file" => return self.stage_delete(args).await,
                    "read-file" => {
                        if let Some(res) = self.overlay_read(args) {
                            return res;
                        }
                    }
                    // glob / grep / list-dir read disk directly (MVP: the
                    // overlay only resolves read-file).
                    _ => {}
                }
            }
        }
        self.inner.call(name, args).await
    }
}

/// A success result shaped like the real fs adapter's, tagged `staged` so the
/// model (and logging) can tell the write was deferred.
fn staged_ok(action: &str, path: &str, extra: Value) -> Value {
    let mut payload = json!({ "action": action, "path": path, "staged": true });
    if let (Some(obj), Some(ex)) = (payload.as_object_mut(), extra.as_object()) {
        for (k, v) in ex {
            obj.insert(k.clone(), v.clone());
        }
    }
    json!({ "status": "succeeded", "payload": payload })
}

fn arg_str(args: &Value, key: &str) -> Result<String, String> {
    args.get(key)
        .and_then(|v| v.as_str())
        .filter(|s| !s.is_empty())
        .map(str::to_string)
        .ok_or_else(|| format!("missing required field '{key}'"))
}

/// Resolve an image reference into something a multimodal API accepts: an
/// `http(s)`/`data:` URL passes through unchanged; a local path is read and
/// encoded into a `data:<mime>;base64,…` URL (local inference servers can't
/// fetch arbitrary URLs). Returns an error string the caller surfaces.
pub fn resolve_image_ref(raw: &str) -> Result<String, String> {
    let r = raw.trim();
    if r.is_empty() {
        return Err("empty image reference".into());
    }
    if r.starts_with("http://") || r.starts_with("https://") || r.starts_with("data:") {
        return Ok(r.to_string());
    }
    let bytes = std::fs::read(r).map_err(|e| format!("read image {r}: {e}"))?;
    Ok(format!("data:{};base64,{}", mime_for(r), b64_encode(&bytes)))
}

fn mime_for(path: &str) -> &'static str {
    let ext = path.rsplit('.').next().unwrap_or("").to_ascii_lowercase();
    match ext.as_str() {
        "jpg" | "jpeg" => "image/jpeg",
        "gif" => "image/gif",
        "webp" => "image/webp",
        _ => "image/png",
    }
}

/// Minimal standard-alphabet base64 encoder (avoids pulling a crate for the
/// single use of building image `data:` URLs).
fn b64_encode(data: &[u8]) -> String {
    const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity(data.len().div_ceil(3) * 4);
    for chunk in data.chunks(3) {
        let b0 = chunk[0];
        let b1 = *chunk.get(1).unwrap_or(&0);
        let b2 = *chunk.get(2).unwrap_or(&0);
        let n = ((b0 as u32) << 16) | ((b1 as u32) << 8) | (b2 as u32);
        out.push(T[((n >> 18) & 63) as usize] as char);
        out.push(T[((n >> 12) & 63) as usize] as char);
        out.push(if chunk.len() > 1 {
            T[((n >> 6) & 63) as usize] as char
        } else {
            '='
        });
        out.push(if chunk.len() > 2 {
            T[(n & 63) as usize] as char
        } else {
            '='
        });
    }
    out
}

// ---------- node field readers + label matching ----------

/// The AI node's dispatch task: `"generate"` (default), `"embedding"`, or
/// `"classify"`. The inspector stamps it from the bound model's capability /
/// whether labels are set.
pub fn ai_task(node: &FlowNode) -> String {
    node.data
        .get("task")
        .and_then(|v| v.as_str())
        .filter(|s| !s.is_empty())
        .unwrap_or("generate")
        .to_string()
}

/// Per-request reasoning toggle from the node's `thinking` flag.
pub fn ai_reasoning(node: &FlowNode) -> Option<bool> {
    node.data.get("thinking").and_then(|v| v.as_bool())
}

/// Classification labels (trimmed, non-empty) from the node's `labels` array.
pub fn ai_labels(node: &FlowNode) -> Vec<String> {
    node.data
        .get("labels")
        .and_then(|v| v.as_array())
        .map(|a| {
            a.iter()
                .filter_map(|x| x.as_str().map(|s| s.trim().to_string()))
                .filter(|s| !s.is_empty())
                .collect()
        })
        .unwrap_or_default()
}

/// Adapter names the node exposes to the model as tools (its `tools` array).
pub fn ai_tool_adapters(node: &FlowNode) -> Vec<String> {
    node.data
        .get("tools")
        .and_then(|v| v.as_array())
        .map(|a| {
            a.iter()
                .filter_map(|x| x.as_str().map(str::to_string))
                .collect()
        })
        .unwrap_or_default()
}

/// Per-node cap on the in-node tool-use loop, read from the node's
/// `maxToolIters` field; falls back to [`DEFAULT_TOOL_ITERATIONS`]. Clamped to
/// at least 1 so a `0` can't silently disable tool use.
pub fn ai_max_tool_iters(node: &FlowNode) -> usize {
    node.data
        .get("maxToolIters")
        .and_then(|v| v.as_u64())
        .map(|v| (v as usize).max(1))
        .unwrap_or(DEFAULT_TOOL_ITERATIONS)
}

/// The structured-output JSON schema (raw text) from the node's `outputSchema`
/// field, used by the `"structured"` task. Stored as a string so it round-trips
/// through the DSL; the model is prompted to produce JSON matching it. `None`
/// when unset or blank.
pub fn ai_output_schema(node: &FlowNode) -> Option<String> {
    let s = match node.data.get("outputSchema")? {
        Value::String(s) => s.clone(),
        // Tolerate an object/array set programmatically by stringifying it.
        v @ (Value::Object(_) | Value::Array(_)) => v.to_string(),
        _ => return None,
    };
    let s = s.trim().to_string();
    if s.is_empty() {
        None
    } else {
        Some(s)
    }
}

/// The node's natural-language output expectation (`expect` field) for per-node
/// AI evaluation (roadmap E1). When set, the generate path self-evaluates the
/// output against it and fails the node - engaging the monitor/fix flow - if
/// unmet. `None` when unset/blank.
pub fn ai_expect(node: &FlowNode) -> Option<String> {
    node.data
        .get("expect")
        .and_then(|v| v.as_str())
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .map(str::to_string)
}

/// Best-effort parse of a model's text into a JSON value: a direct parse first,
/// else the substring between the first `{` and last `}` - tolerates markdown
/// fences or surrounding prose a model adds around structured output.
pub fn parse_json_lenient(text: &str) -> Option<Value> {
    if let Ok(v) = serde_json::from_str::<Value>(text.trim()) {
        return Some(v);
    }
    let start = text.find('{')?;
    let end = text.rfind('}')?;
    if end <= start {
        return None;
    }
    serde_json::from_str(&text[start..=end]).ok()
}

/// Resolve the node's `imageUrl` (plus an optional `images` array) into
/// ready-to-send references (URLs pass through; local paths become `data:`
/// URLs). Empty when the node carries no images.
pub fn resolve_node_images(node: &FlowNode) -> Result<Vec<String>, String> {
    let mut refs: Vec<String> = Vec::new();
    if let Some(s) = node
        .data
        .get("imageUrl")
        .and_then(|v| v.as_str())
        .filter(|s| !s.trim().is_empty())
    {
        refs.push(s.to_string());
    }
    if let Some(arr) = node.data.get("images").and_then(|v| v.as_array()) {
        for x in arr {
            if let Some(s) = x.as_str().filter(|s| !s.trim().is_empty()) {
                refs.push(s.to_string());
            }
        }
    }
    refs.iter().map(|r| resolve_image_ref(r)).collect()
}

/// Pick the label the model's output best corresponds to: an exact
/// (case-insensitive) match first, then a label contained in the output, else
/// the first label as a safe default.
pub fn match_label(output: &str, labels: &[String]) -> String {
    let o = output.trim().to_lowercase();
    if let Some(l) = labels.iter().find(|l| l.to_lowercase() == o) {
        return l.clone();
    }
    if let Some(l) = labels.iter().find(|l| o.contains(&l.to_lowercase())) {
        return l.clone();
    }
    labels.first().cloned().unwrap_or_default()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn matches_labels_by_exact_then_contains() {
        let labels = vec!["pass".to_string(), "fail".to_string()];
        assert_eq!(match_label("PASS", &labels), "pass");
        assert_eq!(match_label("the result is fail.", &labels), "fail");
        assert_eq!(match_label("???", &labels), "pass"); // first as default
    }

    #[test]
    fn max_tool_iters_defaults_and_overrides() {
        let mk = |data: Value| FlowNode {
            id: "n".into(),
            node_type: "ai".into(),
            position: Position { x: 0.0, y: 0.0 },
            data,
        };
        assert_eq!(ai_max_tool_iters(&mk(json!({}))), DEFAULT_TOOL_ITERATIONS);
        assert_eq!(ai_max_tool_iters(&mk(json!({ "maxToolIters": 3 }))), 3);
        // 0 is clamped to 1 so tool use can't be silently disabled.
        assert_eq!(ai_max_tool_iters(&mk(json!({ "maxToolIters": 0 }))), 1);
    }

    #[test]
    fn tool_call_summary_picks_salient_arg() {
        assert_eq!(
            summarize_tool_call("fs__edit-file", &json!({ "path": "src/foo.rs" })),
            "edit-file src/foo.rs"
        );
        assert_eq!(
            summarize_tool_call("shell__run-command", &json!({ "command": "cargo test" })),
            "run-command cargo test"
        );
        // No salient arg → just the action id.
        assert_eq!(summarize_tool_call("fs__list-dir", &json!({})), "list-dir");
    }

    #[test]
    fn output_schema_reads_string_and_trims() {
        let mk = |data: Value| FlowNode {
            id: "n".into(),
            node_type: "ai".into(),
            position: Position { x: 0.0, y: 0.0 },
            data,
        };
        assert_eq!(
            ai_output_schema(&mk(json!({ "outputSchema": "  {\"type\":\"object\"}  " })))
                .as_deref(),
            Some("{\"type\":\"object\"}")
        );
        assert!(ai_output_schema(&mk(json!({ "outputSchema": "   " }))).is_none());
        assert!(ai_output_schema(&mk(json!({}))).is_none());
    }

    #[test]
    fn expect_reads_and_trims() {
        let mk = |data: Value| FlowNode {
            id: "n".into(),
            node_type: "ai".into(),
            position: Position { x: 0.0, y: 0.0 },
            data,
        };
        assert_eq!(
            ai_expect(&mk(json!({ "expect": "  is valid JSON  " }))).as_deref(),
            Some("is valid JSON")
        );
        assert!(ai_expect(&mk(json!({ "expect": "   " }))).is_none());
        assert!(ai_expect(&mk(json!({}))).is_none());
    }

    #[test]
    fn parse_json_lenient_handles_fences_and_prose() {
        assert_eq!(parse_json_lenient(r#"{"a":1}"#), Some(json!({"a":1})));
        assert_eq!(
            parse_json_lenient("```json\n{\"a\":1}\n```"),
            Some(json!({"a":1}))
        );
        assert_eq!(
            parse_json_lenient("Sure! Here you go:\n{\"a\":1}\nHope that helps"),
            Some(json!({"a":1}))
        );
        assert_eq!(parse_json_lenient("not json"), None);
    }

    #[tokio::test]
    async fn dispatcher_emits_tool_call_before_running() {
        use crate::events::CapturingSink;
        let sink = CapturingSink::default();
        // Empty registry: the call errors *after* the observer fires, so the
        // ToolCall event is still captured (it's emitted before adapter lookup).
        let d = AdapterToolDispatcher::new(Arc::new(AdapterRegistry::new()), None).with_observer(
            ToolObserver::new(Arc::new(sink.clone()), "exec-1".into(), "node-1".into()),
        );
        let _ = d.call("fs__edit-file", &json!({ "path": "src/foo.rs" })).await;

        let events = sink.events.lock().unwrap();
        let found = events.iter().find_map(|e| match e {
            ExecutionEvent::ToolCall {
                name,
                summary,
                execution_id,
                node_id,
                ..
            } => Some((
                name.clone(),
                summary.clone(),
                execution_id.clone(),
                node_id.clone(),
            )),
            _ => None,
        });
        assert_eq!(
            found,
            Some((
                "fs__edit-file".into(),
                "edit-file src/foo.rs".into(),
                "exec-1".into(),
                "node-1".into()
            ))
        );
    }

    #[test]
    fn base64_matches_known_vectors() {
        assert_eq!(b64_encode(b""), "");
        assert_eq!(b64_encode(b"f"), "Zg==");
        assert_eq!(b64_encode(b"fo"), "Zm8=");
        assert_eq!(b64_encode(b"foo"), "Zm9v");
        assert_eq!(b64_encode(b"foob"), "Zm9vYg==");
        assert_eq!(b64_encode(b"Man"), "TWFu");
    }

    #[test]
    fn urls_pass_through_paths_encode() {
        assert_eq!(
            resolve_image_ref("https://x/y.png").unwrap(),
            "https://x/y.png"
        );
        assert_eq!(resolve_image_ref("data:image/png;base64,AA").unwrap(), "data:image/png;base64,AA");
        assert!(resolve_image_ref("/no/such/file.png").is_err());
    }

    #[tokio::test]
    async fn staging_defers_writes_and_overlays_reads() {
        use crate::adapter::{Adapter, AdapterDescriptor, AdapterError, NodeOutput};

        // Minimal "fs" adapter that only serves read-file from a temp dir, so
        // the staging dispatcher can compute `before`/overlay without the real
        // flow-adapter-fs crate (which would be a dependency cycle).
        struct DiskReadFs;
        #[async_trait::async_trait]
        impl Adapter for DiskReadFs {
            fn name(&self) -> &str {
                "fs"
            }
            async fn execute(&self, node: &FlowNode) -> Result<NodeOutput, AdapterError> {
                let action = node.data.get("actionId").and_then(|v| v.as_str()).unwrap_or("");
                let root = node.data.get("workspaceRoot").and_then(|v| v.as_str()).unwrap_or("");
                let path = node.data.get("path").and_then(|v| v.as_str()).unwrap_or("");
                if action != "read-file" {
                    return Err(AdapterError::Rejected {
                        adapter: "fs".into(),
                        node_id: node.id.clone(),
                        reason: format!("test adapter only serves read-file, got {action}"),
                    });
                }
                match std::fs::read_to_string(std::path::Path::new(root).join(path)) {
                    Ok(content) => Ok(NodeOutput {
                        status: "succeeded".into(),
                        payload: json!({ "action": "read-file", "path": path, "content": content }),
                    }),
                    Err(e) => Err(AdapterError::Failed {
                        adapter: "fs".into(),
                        node_id: node.id.clone(),
                        reason: e.to_string(),
                    }),
                }
            }
            fn descriptor(&self) -> AdapterDescriptor {
                AdapterDescriptor { name: "fs".into(), summary: String::new(), actions: Vec::new() }
            }
        }

        let root = std::env::temp_dir().join(format!("flow-staging-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(root.join("foo.txt"), "hello").unwrap();

        let mut reg = AdapterRegistry::new();
        reg.register(Arc::new(DiskReadFs));
        let state: StagedEdits = Default::default();
        let d = StagingToolDispatcher::new(
            Arc::new(reg),
            Some(root.to_string_lossy().to_string()),
            state.clone(),
        );

        // write-file → staged, not applied to disk.
        d.call("fs__write-file", &json!({"path":"bar.txt","content":"new"})).await.unwrap();
        assert!(!root.join("bar.txt").exists(), "write must be staged, not applied");

        // read-file of a staged new file is served from the overlay.
        let r = d.call("fs__read-file", &json!({"path":"bar.txt"})).await.unwrap();
        assert_eq!(r["payload"]["content"], "new");

        // edit-file applies to the on-disk content, staged; disk stays put.
        d.call("fs__edit-file", &json!({"path":"foo.txt","oldString":"hello","newString":"world"}))
            .await
            .unwrap();
        let r2 = d.call("fs__read-file", &json!({"path":"foo.txt"})).await.unwrap();
        assert_eq!(r2["payload"]["content"], "world", "overlay reflects the edit");
        assert_eq!(std::fs::read_to_string(root.join("foo.txt")).unwrap(), "hello", "disk untouched");

        // delete-file → staged; a later read errors but the file survives.
        d.call("fs__delete-file", &json!({"path":"foo.txt"})).await.unwrap();
        assert!(d.call("fs__read-file", &json!({"path":"foo.txt"})).await.is_err());
        assert!(root.join("foo.txt").exists(), "delete must be staged, not applied");

        // Net proposed edits, in first-touch order.
        let edits = state.lock().unwrap().edits();
        assert_eq!(edits.len(), 2);
        assert_eq!(edits[0].path, "bar.txt");
        assert_eq!(edits[0].kind, "create");
        assert_eq!(edits[0].after.as_deref(), Some("new"));
        assert_eq!(edits[0].before, None);
        assert_eq!(edits[1].path, "foo.txt");
        assert_eq!(edits[1].kind, "delete");
        assert_eq!(edits[1].before.as_deref(), Some("hello"));
        assert_eq!(edits[1].after, None);

        let _ = std::fs::remove_dir_all(&root);
    }
}