supercode-core 0.1.0

//! Drive the agent loop with a mock provider — no network — to prove the
//! tool-execution cycle and session continuation work end to end.

use std::sync::atomic::{AtomicUsize, Ordering};

use async_trait::async_trait;
use supercode::session::Session;
use supercode::{
    Agent, ChatMessage, ChatRequest, Config, FunctionCall, Provider, Role, ToolCall, Usage,
};

/// First turn: ask for a `list_dir` tool call. Second turn: answer with text
/// that echoes back what the tool returned (so we can assert the tool actually
/// ran and its output was fed back into context).
struct ScriptedProvider {
    calls: AtomicUsize,
}

#[async_trait]
impl Provider for ScriptedProvider {
    async fn complete(
        &self,
        req: &ChatRequest,
        _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
    ) -> supercode::Result<(ChatMessage, Usage)> {
        let n = self.calls.fetch_add(1, Ordering::SeqCst);
        if n == 0 {
            // The tool list must have been advertised.
            assert!(req.tools.iter().any(|t| t.name == "list_dir"));
            let call = ChatMessage {
                role: Role::Assistant,
                content: None,
                content_parts: None,
                tool_calls: Some(vec![ToolCall {
                    id: "call_1".into(),
                    kind: "function".into(),
                    function: FunctionCall {
                        name: "list_dir".into(),
                        arguments: "{}".into(),
                    },
                }]),
                tool_call_id: None,
                name: None,
                metadata: Default::default(),
            };
            Ok((call, Usage::default()))
        } else {
            // The previous message must be the tool result we fed back.
            let last = req.messages.last().unwrap();
            assert_eq!(last.role, Role::Tool);
            let listing = last.content.clone().unwrap_or_default();
            Ok((
                ChatMessage::assistant(format!("saw: {listing}")),
                Usage::default(),
            ))
        }
    }
}

fn temp_dir(tag: &str) -> std::path::PathBuf {
    let dir = std::env::temp_dir().join(format!("supercode-{tag}-{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    dir
}

#[tokio::test]
async fn agent_runs_a_tool_then_answers() {
    let dir = temp_dir("agentloop");
    std::fs::write(dir.join("marker.txt"), "hi").unwrap();

    let config = Config::builder().cwd(dir.clone()).build();
    let provider = Box::new(ScriptedProvider {
        calls: AtomicUsize::new(0),
    });
    let mut agent = Agent::with_provider(config, provider);

    let reply = agent.send("what files are here?").await.unwrap();

    // The final answer reflects the tool output, proving the loop ran the tool
    // and fed the result back to the model.
    assert!(reply.contains("marker.txt"), "reply was: {reply}");

    // History shape: system, user, assistant(tool_call), tool, assistant.
    let roles: Vec<Role> = agent.history().iter().map(|m| m.role).collect();
    assert_eq!(
        roles,
        vec![
            Role::System,
            Role::User,
            Role::Assistant,
            Role::Tool,
            Role::Assistant,
        ]
    );

    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn disabled_tools_are_not_advertised() {
    use std::sync::{Arc, Mutex};

    let dir = temp_dir("disabled");
    let config = Config::builder()
        .cwd(dir.clone())
        .disable_tool("bash")
        .tool_description("read_file", "Custom read description")
        .build();

    // A provider that records the tool set it was offered into shared handles.
    struct Recorder {
        names: Arc<Mutex<Vec<String>>>,
        read_desc: Arc<Mutex<String>>,
    }
    #[async_trait]
    impl Provider for Recorder {
        async fn complete(
            &self,
            req: &ChatRequest,
            _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
        ) -> supercode::Result<(ChatMessage, Usage)> {
            *self.names.lock().unwrap() = req.tools.iter().map(|t| t.name.clone()).collect();
            for t in &req.tools {
                if t.name == "read_file" {
                    *self.read_desc.lock().unwrap() = t.description.clone();
                }
            }
            Ok((ChatMessage::assistant("ok"), Usage::default()))
        }
    }

    let names = Arc::new(Mutex::new(Vec::new()));
    let read_desc = Arc::new(Mutex::new(String::new()));
    let provider = Box::new(Recorder {
        names: names.clone(),
        read_desc: read_desc.clone(),
    });
    let mut agent = Agent::with_provider(config, provider);
    agent.send("hi").await.unwrap();

    let offered = names.lock().unwrap().clone();
    assert!(
        !offered.contains(&"bash".to_string()),
        "bash should be disabled"
    );
    assert!(offered.contains(&"read_file".to_string()));
    assert_eq!(
        read_desc.lock().unwrap().as_str(),
        "Custom read description"
    );

    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn load_session_seeds_history_for_continuation() {
    let fixture =
        std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/codex_session.jsonl");
    let session = Session::load(&fixture).unwrap();
    let loaded = session.messages.len();

    let config = Config::builder().build();
    let provider = Box::new(ScriptedProvider {
        calls: AtomicUsize::new(1), // skip straight to a text answer
    });
    let mut agent = Agent::with_provider(config, provider);
    agent.load_session(session);

    // System prompt + every loaded message.
    assert_eq!(agent.history().len(), 1 + loaded);
    assert_eq!(agent.history()[0].role, Role::System);
}

/// Provider that always asks for a `bash` tool call once, then answers.
struct BashOnce {
    calls: AtomicUsize,
}
#[async_trait]
impl Provider for BashOnce {
    async fn complete(
        &self,
        req: &ChatRequest,
        _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
    ) -> supercode::Result<(ChatMessage, Usage)> {
        let n = self.calls.fetch_add(1, Ordering::SeqCst);
        if n == 0 {
            let call = ChatMessage {
                role: Role::Assistant,
                content: None,
                content_parts: None,
                tool_calls: Some(vec![ToolCall {
                    id: "c1".into(),
                    kind: "function".into(),
                    function: FunctionCall {
                        name: "bash".into(),
                        arguments: "{\"command\":\"echo ran\"}".into(),
                    },
                }]),
                tool_call_id: None,
                name: None,
                metadata: Default::default(),
            };
            Ok((call, Usage::default()))
        } else {
            // Echo what the tool returned so the test can see if it executed.
            let last = req.messages.last().unwrap();
            Ok((
                ChatMessage::assistant(last.content.clone().unwrap_or_default()),
                Usage::default(),
            ))
        }
    }
}

#[tokio::test]
async fn approval_denied_blocks_tool_execution() {
    use supercode::ApprovalPolicy;
    let dir = temp_dir("approval");
    // Untrusted policy + a handler that denies everything.
    let config = Config::builder()
        .cwd(dir.clone())
        .approval(ApprovalPolicy::Untrusted)
        .approval_handler(Box::new(|_call| false))
        .build();
    let mut agent = Agent::with_provider(
        config,
        Box::new(BashOnce {
            calls: AtomicUsize::new(0),
        }),
    );
    let reply = agent.send("run it").await.unwrap();
    // The tool did NOT run; the model saw a "not approved" result.
    assert!(reply.contains("not approved"), "reply: {reply}");
    assert!(
        !reply.contains("ran"),
        "tool must not have executed: {reply}"
    );
    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn approval_handler_allows_and_allowlist_skips() {
    use std::sync::{Arc, Mutex};
    use supercode::ApprovalPolicy;
    let dir = temp_dir("approval2");

    // OnRequest + bash on the allowlist => handler is never consulted, tool runs.
    let consulted = Arc::new(Mutex::new(0usize));
    let c2 = consulted.clone();
    let config = Config::builder()
        .cwd(dir.clone())
        .approval(ApprovalPolicy::OnRequest)
        .auto_approve_tool("bash")
        .approval_handler(Box::new(move |_| {
            *c2.lock().unwrap() += 1;
            true
        }))
        .build();
    let mut agent = Agent::with_provider(
        config,
        Box::new(BashOnce {
            calls: AtomicUsize::new(0),
        }),
    );
    let reply = agent.send("run it").await.unwrap();
    assert!(
        reply.contains("ran"),
        "allowlisted tool should run: {reply}"
    );
    assert_eq!(
        *consulted.lock().unwrap(),
        0,
        "handler must not be consulted for allowlisted tool"
    );
    std::fs::remove_dir_all(&dir).ok();
}

/// Always asks for a tool call (never stops on its own), reporting output tokens.
struct AlwaysToolCall {
    calls: AtomicUsize,
}
#[async_trait]
impl Provider for AlwaysToolCall {
    async fn complete(
        &self,
        _req: &ChatRequest,
        _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
    ) -> supercode::Result<(ChatMessage, Usage)> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        let msg = ChatMessage {
            role: Role::Assistant,
            content: None,
            content_parts: None,
            tool_calls: Some(vec![ToolCall {
                id: "c".into(),
                kind: "function".into(),
                function: FunctionCall {
                    name: "list_dir".into(),
                    arguments: "{}".into(),
                },
            }]),
            tool_call_id: None,
            name: None,
            metadata: Default::default(),
        };
        let usage = Usage {
            prompt_tokens: 0,
            completion_tokens: 100,
            total_tokens: 100,
        };
        Ok((msg, usage))
    }
}

#[tokio::test]
async fn output_token_budget_stops_the_loop() {
    let dir = temp_dir("budget");
    // Budget of 50 output tokens; each turn reports 100 → loop stops after turn 1
    // instead of running to max_iterations (which would be 25 calls).
    let config = Config::builder()
        .cwd(dir.clone())
        .max_iterations(25)
        .max_total_output_tokens(50)
        .build();
    let provider = std::sync::Arc::new(AlwaysToolCall {
        calls: AtomicUsize::new(0),
    });
    // Re-expose call count via a thin wrapper.
    struct Shared(std::sync::Arc<AlwaysToolCall>);
    #[async_trait]
    impl Provider for Shared {
        async fn complete(
            &self,
            req: &ChatRequest,
            d: &(dyn for<'a> Fn(&'a str) + Send + Sync),
        ) -> supercode::Result<(ChatMessage, Usage)> {
            self.0.complete(req, d).await
        }
    }
    let mut agent = Agent::with_provider(config, Box::new(Shared(provider.clone())));
    let res = agent.send("go").await;
    assert!(
        res.is_ok(),
        "budget should stop cleanly, not error: {res:?}"
    );
    assert_eq!(
        provider.calls.load(Ordering::SeqCst),
        1,
        "budget must stop after one model turn"
    );
    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn auto_loads_project_context_files() {
    let dir = temp_dir("ctx");
    std::fs::write(dir.join("CLAUDE.md"), "PROJECT_RULE: always be terse").unwrap();
    let extra = temp_dir("ctx-extra");
    std::fs::write(extra.join("AGENTS.md"), "EXTRA_ROOT_NOTE here").unwrap();

    let config = Config::builder()
        .cwd(dir.clone())
        .add_dir(extra.clone())
        .project_context(true)
        .system_prompt("BASE PROMPT")
        .build();
    let agent = Agent::with_provider(
        config,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(1),
        }),
    );
    let sys = agent.history()[0].content.clone().unwrap_or_default();
    assert!(sys.contains("BASE PROMPT"));
    assert!(
        sys.contains("PROJECT_RULE: always be terse"),
        "CLAUDE.md not loaded: {sys}"
    );
    assert!(
        sys.contains("EXTRA_ROOT_NOTE here"),
        "extra-root AGENTS.md not loaded: {sys}"
    );

    // Disabled by default: no context files loaded.
    let config = Config::builder()
        .cwd(dir.clone())
        .system_prompt("BASE")
        .build();
    let agent = Agent::with_provider(
        config,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(1),
        }),
    );
    assert!(!agent.history()[0]
        .content
        .clone()
        .unwrap_or_default()
        .contains("PROJECT_RULE"));

    std::fs::remove_dir_all(&dir).ok();
    std::fs::remove_dir_all(&extra).ok();
}

#[tokio::test]
async fn transcript_persistence_and_checkpoint_rewind() {
    let dir = temp_dir("persist");
    std::fs::write(dir.join("marker.txt"), "hi").unwrap();
    let config = Config::builder().cwd(dir.clone()).build();
    let mut agent = Agent::with_provider(
        config,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(0),
        }),
    );

    let cp = agent.checkpoint(); // before any turn (just the system msg)
    let reply = agent.send("what files?").await.unwrap();
    assert!(reply.contains("marker.txt"));
    let grown = agent.history().len();
    assert!(grown > cp);

    // Persist + reload into a fresh agent → identical history.
    let path = dir.join("session.jsonl");
    agent.save_transcript(&path).unwrap();
    let config2 = Config::builder().cwd(dir.clone()).build();
    let mut agent2 = Agent::with_provider(
        config2,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(1),
        }),
    );
    agent2.load_transcript(&path).unwrap();
    assert_eq!(agent2.history().len(), grown);
    assert_eq!(
        agent2.history()[0].content,
        agent.history()[0].content,
        "system message preserved across save/load"
    );

    // Rewind discards everything after the checkpoint.
    agent.rewind_to(cp);
    assert_eq!(agent.history().len(), cp);

    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn session_store_lifecycle() {
    use supercode::SessionStore;
    let dir = temp_dir("store");
    let store = SessionStore::open(&dir).unwrap();

    store
        .save("s1", "First session", "{\"role\":\"user\"}\n")
        .unwrap();
    store.save("s2", "Second", "{\"role\":\"user\"}\n").unwrap();
    let names: Vec<String> = store.list().into_iter().map(|i| i.name).collect();
    assert_eq!(names, vec!["s1", "s2"]);
    assert_eq!(
        store.list().iter().find(|i| i.name == "s1").unwrap().title,
        "First session"
    );

    // Title rename.
    store.set_title("s1", "Renamed").unwrap();
    assert_eq!(
        store.list().iter().find(|i| i.name == "s1").unwrap().title,
        "Renamed"
    );

    // Archive moves it but keeps it listed as archived + still loadable.
    store.archive("s1").unwrap();
    let s1 = store.list().into_iter().find(|i| i.name == "s1").unwrap();
    assert!(s1.archived);
    assert!(store.load("s1").unwrap().contains("user"));

    // Delete removes it entirely.
    store.delete("s2").unwrap();
    assert!(!store.list().iter().any(|i| i.name == "s2"));

    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn pre_and_post_tool_hooks() {
    use std::sync::{Arc, Mutex};
    let dir = temp_dir("hooks");
    std::fs::write(dir.join("marker.txt"), "hi").unwrap();

    // post-hook records every (tool, is_error); pre-hook blocks "bash".
    let log = Arc::new(Mutex::new(Vec::<String>::new()));
    let l2 = log.clone();
    let config = Config::builder()
        .cwd(dir.clone())
        .pre_tool_hook(Box::new(|name, _args| {
            if name == "bash" {
                Some("bash disabled by hook".into())
            } else {
                None
            }
        }))
        .post_tool_hook(Box::new(move |name, _out, is_err| {
            l2.lock().unwrap().push(format!("{name}:{is_err}"));
        }))
        .build();
    // ScriptedProvider calls list_dir (allowed) then answers.
    let mut agent = Agent::with_provider(
        config,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(0),
        }),
    );
    let reply = agent.send("list").await.unwrap();
    assert!(reply.contains("marker.txt"));
    // post-hook saw the list_dir call (not an error).
    assert!(
        log.lock().unwrap().iter().any(|e| e == "list_dir:false"),
        "{:?}",
        log.lock().unwrap()
    );

    // Now prove the pre-hook blocks bash via BashOnce.
    let log2 = Arc::new(Mutex::new(Vec::<String>::new()));
    let l3 = log2.clone();
    let config = Config::builder()
        .cwd(dir.clone())
        .pre_tool_hook(Box::new(|name, _| {
            if name == "bash" {
                Some("nope".into())
            } else {
                None
            }
        }))
        .post_tool_hook(Box::new(move |n, _, e| {
            l3.lock().unwrap().push(format!("{n}:{e}"))
        }))
        .build();
    let mut agent = Agent::with_provider(
        config,
        Box::new(BashOnce {
            calls: AtomicUsize::new(0),
        }),
    );
    let reply = agent.send("run").await.unwrap();
    assert!(reply.contains("blocked by pre-tool hook"), "{reply}");
    // post-hook should NOT have fired for the blocked bash (it never executed).
    assert!(!log2.lock().unwrap().iter().any(|e| e.starts_with("bash:")));

    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn slash_prompt_expansion_and_compaction() {
    let dir = temp_dir("prompts");
    // Custom prompt + built-in code-review.
    let config = Config::builder()
        .cwd(dir.clone())
        .prompt("greet", "Say hello to {args} politely.")
        .build();
    let agent = Agent::with_provider(
        config,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(1),
        }),
    );
    assert_eq!(
        agent.expand_prompt("/greet Alice"),
        "Say hello to Alice politely."
    );
    assert!(agent
        .expand_prompt("/code-review the auth module")
        .contains("correctness bugs"));
    assert!(agent
        .expand_prompt("/code-review the auth module")
        .contains("the auth module"));
    assert_eq!(agent.expand_prompt("plain message"), "plain message"); // non-slash untouched
    assert_eq!(agent.expand_prompt("/unknown x"), "/unknown x"); // unknown command untouched

    // Compaction: build a long history via a saved transcript, then compact.
    let mut lines = String::new();
    lines.push_str(&serde_json::to_string(&supercode::ChatMessage::system("SYS")).unwrap());
    lines.push('\n');
    for i in 0..10 {
        lines.push_str(
            &serde_json::to_string(&supercode::ChatMessage::user(format!("msg {i}"))).unwrap(),
        );
        lines.push('\n');
    }
    let tpath = dir.join("hist.jsonl");
    std::fs::write(&tpath, lines).unwrap();
    let config = Config::builder()
        .cwd(dir.clone())
        .compact_after_messages(6)
        .build();
    let mut agent = Agent::with_provider(
        config,
        Box::new(ScriptedProvider {
            calls: AtomicUsize::new(1),
        }),
    );
    agent.load_transcript(&tpath).unwrap();
    let before = agent.history().len();
    let compacted = agent.maybe_compact();
    assert!(
        compacted,
        "should compact when over threshold (had {before})"
    );
    assert!(agent.history().len() < before);
    assert_eq!(agent.history()[0].role, supercode::Role::System); // system preserved
    assert!(agent.history()[1]
        .content
        .as_deref()
        .unwrap_or("")
        .contains("compacted"));
    // Most recent message retained.
    assert!(agent
        .history()
        .last()
        .unwrap()
        .content
        .as_deref()
        .unwrap_or("")
        .contains("msg 9"));
    std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn config_profiles_from_file_and_output_format() {
    use supercode::Config;
    let dir = temp_dir("profiles");
    let cfg = r#"{ "profiles": {
        "fast": {"model": "openai/gpt-5", "effort": "low", "sandbox": "read_only", "approval": "untrusted"},
        "careful": {"model": "anthropic/claude-opus-4-8", "effort": "high", "project_context": true}
    }}"#;
    let path = dir.join("supercode.json");
    std::fs::write(&path, cfg).unwrap();

    let fast = Config::from_profile_file(&path, "fast").unwrap().build();
    assert_eq!(fast.model, "openai/gpt-5");
    assert_eq!(fast.effort.as_deref(), Some("low"));
    assert_eq!(fast.sandbox, supercode::SandboxPolicy::ReadOnly);
    assert_eq!(fast.approval, supercode::ApprovalPolicy::Untrusted);

    let careful = Config::from_profile_file(&path, "careful").unwrap().build();
    assert_eq!(careful.model, "anthropic/claude-opus-4-8");
    assert!(careful.load_project_context);
    // Unset fields keep their defaults.
    assert_eq!(careful.sandbox, supercode::SandboxPolicy::DangerFullAccess);

    assert!(Config::from_profile_file(&path, "missing").is_err());

    // Output formatting.
    assert_eq!(supercode::format_reply("hi", false), "hi");
    assert_eq!(supercode::format_reply("hi", true), "{\"result\":\"hi\"}");

    std::fs::remove_dir_all(&dir).ok();
}

/// A trivial provider that just answers with fixed text (no tools).
struct SaysProvider(String);
#[async_trait]
impl Provider for SaysProvider {
    async fn complete(
        &self,
        _req: &ChatRequest,
        _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
    ) -> supercode::Result<(ChatMessage, Usage)> {
        Ok((ChatMessage::assistant(self.0.clone()), Usage::default()))
    }
}

#[tokio::test]
async fn subagent_shares_transport_and_runs() {
    let dir = temp_dir("subagent");
    let config = Config::builder().cwd(dir.clone()).build();
    let agent = Agent::with_provider(config, Box::new(SaysProvider("PARENT".into())));
    // Subagent shares the same provider (Arc) → same canned answer, fresh convo.
    let out = agent
        .run_subagent("You are a worker.", "do the thing")
        .await
        .unwrap();
    assert_eq!(out, "PARENT");
    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn background_task_runs_and_returns_agent() {
    let dir = temp_dir("bg");
    let config = Config::builder().cwd(dir.clone()).build();
    let agent = Agent::with_provider(config, Box::new(SaysProvider("DONE".into())));
    let handle = agent.run_in_background("go do it");
    let (agent, result) = handle.await.unwrap();
    assert_eq!(result.unwrap(), "DONE");
    // The returned agent retains the conversation.
    assert!(agent
        .history()
        .iter()
        .any(|m| m.content.as_deref() == Some("DONE")));
    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn send_with_files_injects_contents() {
    let dir = temp_dir("files");
    std::fs::write(dir.join("notes.txt"), "IMPORTANT_NOTE_42").unwrap();
    std::fs::write(dir.join("img.bin"), [0u8, 159, 146, 150]).unwrap(); // invalid utf8
    let config = Config::builder().cwd(dir.clone()).build();
    let mut agent = Agent::with_provider(config, Box::new(SaysProvider("ok".into())));

    agent
        .send_with_files(
            "review these",
            &[dir.join("notes.txt"), dir.join("img.bin")],
        )
        .await
        .unwrap();

    // The user message carries the file contents (text inline; binary noted).
    let user = agent
        .history()
        .iter()
        .find(|m| m.role == supercode::Role::User)
        .unwrap();
    let c = user.content.clone().unwrap();
    assert!(c.contains("review these"));
    assert!(c.contains("IMPORTANT_NOTE_42"), "text file injected: {c}");
    assert!(
        c.contains("binary content omitted"),
        "binary file noted: {c}"
    );
    std::fs::remove_dir_all(&dir).ok();
}

#[tokio::test]
async fn send_with_images_builds_multimodal_message() {
    use supercode::ChatMessage;
    // The user message serializes `content` as a multimodal array on the wire.
    let m = ChatMessage::user_with_images(
        "what is in this image?",
        &[
            "https://example.com/a.png".to_string(),
            "data:image/png;base64,AAAA".to_string(),
        ],
    );
    let wire: serde_json::Value = serde_json::to_value(&m).unwrap();
    assert_eq!(wire["role"], "user");
    assert!(
        wire["content"].is_array(),
        "content must be a multimodal array"
    );
    assert_eq!(wire["content"][0]["type"], "text");
    assert_eq!(wire["content"][0]["text"], "what is in this image?");
    assert_eq!(wire["content"][1]["type"], "image_url");
    assert_eq!(
        wire["content"][1]["image_url"]["url"],
        "https://example.com/a.png"
    );
    assert_eq!(
        wire["content"][2]["image_url"]["url"],
        "data:image/png;base64,AAAA"
    );

    // Round-trips through Deserialize (transcript persistence).
    let back: ChatMessage = serde_json::from_value(wire).unwrap();
    assert!(back.content_parts.is_some() && back.content.is_none());

    // A plain text message still serializes content as a string.
    let t = serde_json::to_value(ChatMessage::user("hi")).unwrap();
    assert!(t["content"].is_string());

    // Agent path: the provider receives the multimodal message.
    let dir = temp_dir("vision");
    let config = Config::builder().cwd(dir.clone()).build();
    struct CapturingProvider(std::sync::Arc<std::sync::Mutex<Option<ChatRequest>>>);
    #[async_trait]
    impl Provider for CapturingProvider {
        async fn complete(
            &self,
            req: &ChatRequest,
            _d: &(dyn for<'a> Fn(&'a str) + Send + Sync),
        ) -> supercode::Result<(ChatMessage, Usage)> {
            *self.0.lock().unwrap() = Some(req.clone());
            Ok((ChatMessage::assistant("a cat"), Usage::default()))
        }
    }
    let captured = std::sync::Arc::new(std::sync::Mutex::new(None));
    let mut agent = Agent::with_provider(config, Box::new(CapturingProvider(captured.clone())));
    let reply = agent
        .send_with_images("describe", &["data:image/png;base64,XYZ".into()])
        .await
        .unwrap();
    assert_eq!(reply, "a cat");
    let req = captured.lock().unwrap().clone().unwrap();
    let user = req
        .messages
        .iter()
        .find(|m| m.role == supercode::Role::User)
        .unwrap();
    assert!(
        user.content_parts.is_some(),
        "the user turn carried image parts to the provider"
    );
    std::fs::remove_dir_all(&dir).ok();
}

/// A provider that always returns a plain text reply, reporting a fixed number
/// of completion tokens — to prove the agent accumulates output tokens across
/// sends (the figure the bench reports).
struct TokenReporter(u64);

#[async_trait]
impl Provider for TokenReporter {
    async fn complete(
        &self,
        _req: &ChatRequest,
        _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
    ) -> supercode::Result<(ChatMessage, Usage)> {
        let usage = Usage {
            prompt_tokens: 0,
            completion_tokens: self.0,
            total_tokens: self.0,
        };
        Ok((ChatMessage::assistant("done"), usage))
    }
}

#[tokio::test]
async fn agent_accumulates_output_tokens_across_sends() {
    let mut agent = Agent::with_provider(Config::builder().build(), Box::new(TokenReporter(7)));
    assert_eq!(agent.total_output_tokens(), 0);
    agent.send("one").await.unwrap();
    assert_eq!(agent.total_output_tokens(), 7);
    agent.send("two").await.unwrap();
    assert_eq!(
        agent.total_output_tokens(),
        14,
        "tokens accumulate across sends"
    );
}

/// A provider that calls a tool whose output is huge, then summarizes — to
/// prove oversized tool output is capped before it re-enters the context.
struct BigToolThenDone {
    calls: AtomicUsize,
}

#[async_trait]
impl Provider for BigToolThenDone {
    async fn complete(
        &self,
        req: &ChatRequest,
        _on_delta: &(dyn for<'a> Fn(&'a str) + Send + Sync),
    ) -> supercode::Result<(ChatMessage, Usage)> {
        let n = self.calls.fetch_add(1, Ordering::SeqCst);
        if n == 0 {
            let mut msg = ChatMessage::assistant("");
            msg.tool_calls = Some(vec![ToolCall {
                id: "c1".into(),
                kind: "function".into(),
                function: FunctionCall {
                    name: "list_dir".into(),
                    arguments: "{}".into(),
                },
            }]);
            Ok((msg, Usage::default()))
        } else {
            // The tool result we got back must have been capped.
            let tool_msg = req
                .messages
                .iter()
                .rev()
                .find(|m| m.role == Role::Tool)
                .unwrap();
            let len = tool_msg.content.as_deref().unwrap_or("").len();
            assert!(len <= 5000, "tool output should be capped, got {len} bytes");
            Ok((ChatMessage::assistant("done"), Usage::default()))
        }
    }
}

#[tokio::test]
async fn oversized_tool_output_is_capped() {
    // A tool that returns a megabyte of text.
    struct Big;
    #[async_trait]
    impl supercode::tools::Tool for Big {
        fn name(&self) -> &str {
            "list_dir"
        }
        fn description(&self) -> &str {
            "x"
        }
        fn parameters(&self) -> serde_json::Value {
            serde_json::json!({"type":"object"})
        }
        async fn execute(
            &self,
            _a: serde_json::Value,
            _c: &supercode::tools::ToolContext,
        ) -> supercode::Result<String> {
            Ok("A".repeat(1_000_000))
        }
    }
    let config = Config::builder().max_tool_output_bytes(4096).build();
    let mut reg = supercode::tools::ToolRegistry::new();
    reg.register(Big);
    let mut agent = Agent::with_parts(
        config,
        Box::new(BigToolThenDone {
            calls: AtomicUsize::new(0),
        }),
        reg,
    );
    let reply = agent.send("go").await.unwrap();
    assert_eq!(reply, "done");
}