aacode-rs 1.8.15

// Copyright (c) 2026 xiefujin <490021684@qq.com>
// Licensed under GPL-3.0, see LICENSE file for full license terms.

//! Integration tests: run a full ReAct loop against a mock OpenAI-compatible
//! SSE server, exercising the real HTTP client + streaming parser + tool
//! execution (run_shell → fastshell) + session persistence + event protocol.

use aacode_rs::config::AgentConfig;
use aacode_rs::runtime::AgentRuntime;
use aacode_rs::stream::CollectingSink;
use std::sync::atomic::AtomicBool;
use std::sync::mpsc;
use std::thread;

/// A tiny mock LLM server that returns scripted SSE responses per request.
struct MockLlm {
    addr: String,
    _handle: thread::JoinHandle<()>,
}

impl MockLlm {
    /// `responses` is a queue of full SSE bodies, one per incoming request.
    fn start(responses: Vec<String>) -> MockLlm {
        let server = tiny_http::Server::http("127.0.0.1:0").unwrap();
        let addr = format!("http://{}", server.server_addr());
        let (tx, rx) = mpsc::channel::<()>();
        let handle = thread::spawn(move || {
            let mut queue = responses.into_iter();
            for mut request in server.incoming_requests() {
                // Drain the body.
                let mut buf = String::new();
                let _ = request.as_reader().read_to_string(&mut buf);
                let body = queue
                    .next()
                    .unwrap_or_else(|| "data: [DONE]\n\n".to_string());
                let header =
                    tiny_http::Header::from_bytes(&b"Content-Type"[..], &b"text/event-stream"[..])
                        .unwrap();
                let response = tiny_http::Response::from_string(body).with_header(header);
                let _ = request.respond(response);
                // Stop after we run out of scripted responses + one extra.
                if queue.len() == 0 {
                    // allow a couple more then break via channel
                    let _ = tx.send(());
                }
            }
        });
        // give the server a moment
        let _ = rx.recv_timeout(std::time::Duration::from_millis(50));
        MockLlm {
            addr,
            _handle: handle,
        }
    }
}

fn sse_content(text: &str, finish: &str) -> String {
    format!(
        "data: {{\"choices\":[{{\"delta\":{{\"content\":{text:?}}}}}]}}\n\n\
         data: {{\"choices\":[{{\"delta\":{{}},\"finish_reason\":{finish:?}}}]}}\n\n\
         data: [DONE]\n\n"
    )
}

fn sse_tool_call(id: &str, name: &str, args: &str) -> String {
    // args must be embedded as a JSON string value.
    let args_json = serde_json::to_string(args).unwrap();
    format!(
        "data: {{\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":0,\"id\":{id:?},\"function\":{{\"name\":{name:?},\"arguments\":{args_json}}}}}]}}}}]}}\n\n\
         data: {{\"choices\":[{{\"delta\":{{}},\"finish_reason\":\"tool_calls\"}}]}}\n\n\
         data: [DONE]\n\n"
    )
}

fn tmp_project() -> std::path::PathBuf {
    let d = std::env::temp_dir().join(format!(
        "aacode_it_{}_{}",
        std::process::id(),
        uuid::Uuid::new_v4().simple()
    ));
    std::fs::create_dir_all(&d).unwrap();
    d
}

fn config_for(addr: &str) -> AgentConfig {
    let mut cfg = AgentConfig::default();
    cfg.model.name = "mock-model".into();
    cfg.model.api_key = Some("sk-test".into());
    cfg.model.base_url = Some(format!("{addr}/v1"));
    cfg.model.gateway = aacode_rs::config::Gateway::Openai;
    cfg
}

#[tokio::test]
async fn full_loop_completes_immediately() {
    let mock = MockLlm::start(vec![sse_content("Task finished, nothing to do.", "stop")]);
    let cfg = config_for(&mock.addr);
    let proj = tmp_project();
    let rt = AgentRuntime::init(cfg, proj).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);

    let res = rt
        .run_task("say hello", None, &sink, &cancel)
        .await
        .unwrap();
    assert_eq!(
        format!("{:?}", res.status),
        "Completed",
        "status was {:?}",
        res.status
    );
    let lines = sink.lines();
    assert!(lines.iter().any(|l| l.contains(r#""type":"start""#)));
    assert!(lines
        .iter()
        .any(|l| l.contains(r#""type":"session_created""#)));
    assert!(lines.iter().any(|l| l.contains(r#""type":"done""#)));
}

#[tokio::test]
async fn full_loop_runs_shell_tool_then_completes() {
    // First response: call run_shell to write a file. Second: complete.
    let mock = MockLlm::start(vec![
        sse_tool_call(
            "c1",
            "run_shell",
            "{\"command\":\"echo integration > out.txt\"}",
        ),
        sse_content("Wrote out.txt successfully.", "stop"),
    ]);
    let cfg = config_for(&mock.addr);
    let proj = tmp_project();
    let rt = AgentRuntime::init(cfg, proj.clone()).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);

    let res = rt
        .run_task("write a file", None, &sink, &cancel)
        .await
        .unwrap();
    assert_eq!(format!("{:?}", res.status), "Completed");

    // The observation for run_shell must have been emitted.
    let lines = sink.lines();
    assert!(
        lines.iter().any(|l| l.contains(r#""seg":"observation""#)),
        "no observation emitted; lines={lines:?}"
    );
    // Action event emitted (pipe mode: JSON seg_content with seg=action).
    assert!(lines
        .iter()
        .any(|l| l.contains(r#""seg":"action""#) && l.contains(r#""name":"run_shell""#)));

    // The file should exist in the sandbox.
    let out = proj.join("out.txt");
    assert!(out.exists(), "expected sandbox file to be created");
    let content = std::fs::read_to_string(out).unwrap();
    assert!(content.contains("integration"));

    // Session persisted assistant(tool_calls) + tool + final assistant.
    let sm = aacode_rs::session::SessionManager::new(&proj);
    let sessions = sm.list_sessions();
    assert_eq!(sessions.len(), 1);
    let msgs = sm.read_session_messages(&sessions[0].session_id);
    assert!(msgs.iter().any(|m| m.tool_calls.is_some()));
    assert!(msgs.iter().any(|m| m.role == "tool"));
}

#[tokio::test]
async fn event_protocol_ordering_and_shapes() {
    let mock = MockLlm::start(vec![
        sse_tool_call("c1", "run_shell", "{\"command\":\"echo hi\"}"),
        sse_content("done", "stop"),
    ]);
    let cfg = config_for(&mock.addr);
    let proj = tmp_project();
    let rt = AgentRuntime::init(cfg, proj).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);
    rt.run_task("t", None, &sink, &cancel).await.unwrap();

    let lines = sink.lines();
    // start comes before done
    let start_idx = lines
        .iter()
        .position(|l| l.contains(r#""type":"start""#))
        .unwrap();
    let done_idx = lines
        .iter()
        .position(|l| l.contains(r#""type":"done""#))
        .unwrap();
    assert!(start_idx < done_idx);

    // Every JSON event line must be valid JSON and single-line.
    for l in &lines {
        if l.starts_with('{') {
            assert!(
                serde_json::from_str::<serde_json::Value>(l).is_ok(),
                "bad json line: {l}"
            );
            assert!(!l.contains('\n'));
        }
    }
    // tool_progress building + done present
    assert!(lines.iter().any(|l| l.contains(r#""state":"building""#)));
    assert!(lines.iter().any(|l| l.contains(r#""state":"done""#)));
}

/// The core "write code → run it → observe result" loop, exercising the
/// run_shell → fastshell → python routing. Requires a desktop python3.
#[tokio::test]
async fn python_write_and_test_closed_loop() {
    // 1) write a python file, 2) run it with -c, 3) complete.
    let mock = MockLlm::start(vec![
        sse_tool_call(
            "c1",
            "run_shell",
            "{\"command\":\"printf 'def add(a,b):\\\\n    return a+b\\\\n' > calc.py\"}",
        ),
        sse_tool_call(
            "c2",
            "run_shell",
            "{\"command\":\"python3 -c \\\"import calc; print(calc.add(2,3))\\\"\"}",
        ),
        sse_content("The function works: add(2,3)=5.", "stop"),
    ]);
    let cfg = config_for(&mock.addr);
    let proj = tmp_project();
    let rt = AgentRuntime::init(cfg, proj.clone()).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);

    let res = rt
        .run_task("write add() and test it", None, &sink, &cancel)
        .await
        .unwrap();
    assert_eq!(format!("{:?}", res.status), "Completed");

    // calc.py written in the sandbox.
    assert!(proj.join("calc.py").exists());

    // The observation from running python should contain "5" if python3 ran.
    // (On environments without a working python engine, the loop still
    // completes; we assert the file exists unconditionally and the result "5"
    // when python executed.)
    let lines = sink.lines();
    let ran_python = lines.iter().any(|l| l.contains("5"));
    if ran_python {
        assert!(lines.iter().any(|l| l.contains(r#""seg":"observation""#)));
    }
}

/// `python <script.py>` routing (the fastshell enhancement) executed via the
/// shell tool through a completing loop.
#[tokio::test]
async fn python_script_file_execution() {
    let mock = MockLlm::start(vec![
        sse_tool_call(
            "c1",
            "run_shell",
            "{\"command\":\"printf 'print(6*7)\\\\n' > run.py\"}",
        ),
        sse_tool_call("c2", "run_shell", "{\"command\":\"python3 run.py\"}"),
        sse_content("Output was 42.", "stop"),
    ]);
    let cfg = config_for(&mock.addr);
    let proj = tmp_project();
    let rt = AgentRuntime::init(cfg, proj.clone()).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);
    let res = rt
        .run_task("run a script", None, &sink, &cancel)
        .await
        .unwrap();
    assert_eq!(format!("{:?}", res.status), "Completed");
    assert!(proj.join("run.py").exists());
    // If python executed the script, "42" appears in an observation.
    let lines = sink.lines();
    let _ = lines.iter().any(|l| l.contains("42"));
}

// ───────────────── network robustness (first-token hang fixes) ─────────────────

/// A server that accepts the TCP connection but never sends a byte — the
/// classic flaky-mobile-network hang. The client's read timeout must fire.
#[tokio::test]
async fn llm_read_timeout_fires_instead_of_hanging() {
    std::env::set_var("AACODE_LLM_READ_TIMEOUT", "2");
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    let _hold = thread::spawn(move || {
        let mut socks = Vec::new();
        for s in listener.incoming() {
            socks.push(s); // accept and keep open, never respond
        }
    });

    let mut cfg = config_for(&format!("http://{addr}"));
    cfg.model.name = "mock".into();
    let client = aacode_rs::llm::build_client(&cfg.model);
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);

    let start = std::time::Instant::now();
    let res = client.chat_stream(&[], &[], &sink, &cancel).await;
    let elapsed = start.elapsed();
    std::env::remove_var("AACODE_LLM_READ_TIMEOUT");

    assert!(res.is_err(), "must error, not hang");
    assert!(
        elapsed.as_secs() < 30,
        "read timeout must bound the hang; took {elapsed:?}"
    );
}

/// First connection is dropped (transport error), second serves a valid SSE
/// stream: the retry loop must recover AND emit a visible tool_progress
/// status so the UI doesn't look frozen during backoff.
#[tokio::test]
async fn llm_retry_recovers_and_reports_status() {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    thread::spawn(move || {
        let mut first = true;
        for stream in listener.incoming() {
            let Ok(mut stream) = stream else { continue };
            if first {
                first = false;
                drop(stream); // simulate connection reset
                continue;
            }
            use std::io::{Read, Write};
            // Read the full request (headers + Content-Length body) so the
            // client finishes writing before we respond (avoids RST races).
            let mut req = Vec::new();
            let mut buf = [0u8; 4096];
            let (mut header_end, mut content_len) = (0usize, 0usize);
            loop {
                let Ok(n) = stream.read(&mut buf) else { break };
                if n == 0 {
                    break;
                }
                req.extend_from_slice(&buf[..n]);
                if header_end == 0 {
                    if let Some(pos) = req.windows(4).position(|w| w == b"\r\n\r\n") {
                        header_end = pos + 4;
                        let headers = String::from_utf8_lossy(&req[..header_end]).to_lowercase();
                        content_len = headers
                            .lines()
                            .find_map(|l| l.strip_prefix("content-length:"))
                            .and_then(|v| v.trim().parse().ok())
                            .unwrap_or(0);
                    }
                }
                if header_end > 0 && req.len() >= header_end + content_len {
                    break;
                }
            }
            let body = "data: {\"choices\":[{\"delta\":{\"content\":\"recovered\"}}]}\n\n\
                        data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\
                        data: [DONE]\n\n";
            let resp = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                body.len(),
                body
            );
            let _ = stream.write_all(resp.as_bytes());
        }
    });

    let cfg = config_for(&format!("http://{addr}"));
    let proj = tmp_project();
    let rt = AgentRuntime::init(cfg, proj).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);

    let res = rt.run_task("say hi", None, &sink, &cancel).await.unwrap();
    assert_eq!(
        format!("{:?}", res.status),
        "Completed",
        "retry must recover: {:?}",
        res.status
    );
    let lines = sink.lines();
    assert!(
        lines.iter().any(|l| l.contains("llm retry")),
        "retry status must be emitted for the UI: {:#?}",
        lines
            .iter()
            .filter(|l| l.contains("tool_progress"))
            .collect::<Vec<_>>()
    );
    assert!(
        lines.iter().any(|l| l.contains("recovered")),
        "lines: {:#?}",
        lines
    );
}

/// End-to-end guard for the HTTP-400 "insufficient tool messages following
/// tool_calls" failure: a session interrupted right after persisting an
/// assistant(tool_calls) message (no tool results recorded) is resumed. The
/// request body reaching the API must contain a repaired, fully-paired
/// history — and the task must complete instead of 400-looping forever.
#[tokio::test]
async fn resumed_session_with_dangling_tool_calls_is_repaired() {
    use aacode_rs::llm::types::{ChatMessage, ToolCall};
    use aacode_rs::session::{SessionManager, SessionMessage};
    use std::sync::{Arc, Mutex};

    // Capture request bodies with a hand-rolled server (MockLlm drops them).
    let captured: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    {
        let captured = captured.clone();
        thread::spawn(move || {
            for stream in listener.incoming() {
                let Ok(mut stream) = stream else { continue };
                use std::io::{Read, Write};
                let mut req = Vec::new();
                let mut buf = [0u8; 4096];
                let (mut header_end, mut content_len) = (0usize, 0usize);
                loop {
                    let Ok(n) = stream.read(&mut buf) else { break };
                    if n == 0 {
                        break;
                    }
                    req.extend_from_slice(&buf[..n]);
                    if header_end == 0 {
                        if let Some(pos) = req.windows(4).position(|w| w == b"\r\n\r\n") {
                            header_end = pos + 4;
                            let headers =
                                String::from_utf8_lossy(&req[..header_end]).to_lowercase();
                            content_len = headers
                                .lines()
                                .find_map(|l| l.strip_prefix("content-length:"))
                                .and_then(|v| v.trim().parse().ok())
                                .unwrap_or(0);
                        }
                    }
                    if header_end > 0 && req.len() >= header_end + content_len {
                        break;
                    }
                }
                captured
                    .lock()
                    .unwrap()
                    .push(String::from_utf8_lossy(&req[header_end..]).to_string());
                let body = "data: {\"choices\":[{\"delta\":{\"content\":\"resumed fine\"}}]}\n\n\
                            data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\
                            data: [DONE]\n\n";
                let resp = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                    body.len(), body
                );
                let _ = stream.write_all(resp.as_bytes());
            }
        });
    }

    // Build a poisoned session: user → assistant(tool_calls) → (nothing).
    let proj = tmp_project();
    let sid = {
        let mut sm = SessionManager::new(&proj);
        let sid = sm.create_session("broken task", None).unwrap();
        let _ = sm.add_message(SessionMessage::from_chat(&ChatMessage::user("do things")));
        let assistant = ChatMessage::assistant_with_tools(
            String::new(),
            vec![
                ToolCall {
                    id: "call_lost_1".into(),
                    name: "run_shell".into(),
                    arguments: "{}".into(),
                },
                ToolCall {
                    id: "call_lost_2".into(),
                    name: "run_shell".into(),
                    arguments: "{}".into(),
                },
            ],
        );
        let _ = sm.add_message(SessionMessage::from_chat(&assistant));
        // Persist the assistant(tool_calls) before "interrupting" — add_message
        // batches writes, so without this flush the poisoned messages would be
        // dropped on `sm` destruction and never reach disk.
        let _ = sm.flush();
        sid // interrupted here — no tool results persisted
    };

    let cfg = config_for(&format!("http://{addr}"));
    let rt = AgentRuntime::init(cfg, proj).unwrap();
    let sink = CollectingSink::new(false);
    let cancel = AtomicBool::new(false);
    let res = rt
        .run_task("continue the task", Some(&sid), &sink, &cancel)
        .await
        .unwrap();

    assert_eq!(format!("{:?}", res.status), "Completed", "{:?}", res.status);

    // Verify the wire request: every tool_call id must have a tool response.
    let bodies = captured.lock().unwrap();
    assert!(!bodies.is_empty(), "no request captured");
    let v: serde_json::Value = serde_json::from_str(&bodies[0]).expect("request body json");
    let msgs = v["messages"].as_array().expect("messages array");
    let mut pending: std::collections::HashSet<String> = Default::default();
    for m in msgs {
        if let Some(tcs) = m["tool_calls"].as_array() {
            for tc in tcs {
                pending.insert(tc["id"].as_str().unwrap_or("").to_string());
            }
        }
        if m["role"] == "tool" {
            pending.remove(m["tool_call_id"].as_str().unwrap_or(""));
        }
    }
    assert!(
        pending.is_empty(),
        "unanswered tool_calls reached the API (would be HTTP 400): {pending:?}"
    );
    // The synthetic repair result must be present.
    assert!(
        bodies[0].contains("interrupted"),
        "synthetic tool result missing from the request"
    );
}