johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use serde_json::{json, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

const MOCK_RESPONSE: &str =
    r#"{"choices":[{"message":{"role":"assistant","content":"mocked pong"}}]}"#;

async fn mocked_llm_server() -> (String, tokio::task::JoinHandle<String>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();

    let server = tokio::spawn(async move {
        let (mut socket, _) = listener.accept().await.unwrap();
        let mut request = Vec::new();
        let mut body_start = None;
        let mut content_length = 0;

        loop {
            let mut buffer = [0u8; 1024];
            let read = socket.read(&mut buffer).await.unwrap();
            if read == 0 {
                break;
            }
            request.extend_from_slice(&buffer[..read]);

            if body_start.is_none() {
                if let Some(header_end) = find_header_end(&request) {
                    body_start = Some(header_end);
                    content_length = parse_content_length(&request[..header_end]);
                }
            }

            if let Some(body_start) = body_start {
                if request.len() >= body_start + content_length {
                    break;
                }
            }
        }

        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            MOCK_RESPONSE.len(),
            MOCK_RESPONSE,
        );
        socket.write_all(response.as_bytes()).await.unwrap();
        socket.shutdown().await.ok();

        String::from_utf8_lossy(&request).to_string()
    });

    (format!("http://{address}"), server)
}

fn find_header_end(request: &[u8]) -> Option<usize> {
    request
        .windows(4)
        .position(|window| window == b"\r\n\r\n")
        .map(|position| position + 4)
}

fn parse_content_length(headers: &[u8]) -> usize {
    let text = String::from_utf8_lossy(headers);
    text.lines()
        .find_map(|line| {
            let (name, value) = line.split_once(':')?;
            if name.eq_ignore_ascii_case("content-length") {
                value.trim().parse().ok()
            } else {
                None
            }
        })
        .unwrap_or(0)
}

async fn read_http_request(socket: &mut TcpStream) -> String {
    let mut request = Vec::new();
    let mut body_start = None;
    let mut content_length = 0;

    loop {
        let mut buffer = [0u8; 1024];
        let read = socket.read(&mut buffer).await.unwrap();
        if read == 0 {
            break;
        }
        request.extend_from_slice(&buffer[..read]);

        if body_start.is_none() {
            if let Some(header_end) = find_header_end(&request) {
                body_start = Some(header_end);
                content_length = parse_content_length(&request[..header_end]);
            }
        }

        if let Some(body_start) = body_start {
            if request.len() >= body_start + content_length {
                break;
            }
        }
    }

    String::from_utf8_lossy(&request).to_string()
}

async fn write_http_json(socket: &mut TcpStream, body: &str) {
    let response = format!(
        "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
        body.len(),
        body,
    );
    socket.write_all(response.as_bytes()).await.unwrap();
    socket.shutdown().await.ok();
}

async fn mocked_llm_sequence(
    responses: Vec<String>,
) -> (String, tokio::task::JoinHandle<Vec<String>>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();

    let server = tokio::spawn(async move {
        let mut requests = Vec::new();
        for response in responses {
            let (mut socket, _) = listener.accept().await.unwrap();
            requests.push(read_http_request(&mut socket).await);
            write_http_json(&mut socket, &response).await;
        }
        requests
    });

    (format!("http://{address}"), server)
}

async fn loopback_http_target() -> (String, Arc<AtomicBool>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let address = listener.local_addr().unwrap();
    let hit = Arc::new(AtomicBool::new(false));
    let hit_clone = Arc::clone(&hit);

    tokio::spawn(async move {
        let (mut socket, _) = listener.accept().await.unwrap();
        hit_clone.store(true, Ordering::SeqCst);
        let mut buffer = [0u8; 4096];
        let _ = socket.read(&mut buffer).await;
        let body = "fetched-body";
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            body.len(),
            body,
        );
        socket.write_all(response.as_bytes()).await.unwrap();
        socket.shutdown().await.ok();
    });

    (format!("http://{address}/data"), hit)
}

fn session_files(root: &Path) -> Vec<PathBuf> {
    let mut files = Vec::new();
    let Ok(projects) = fs::read_dir(root.join("projects")) else {
        return files;
    };
    for project in projects.flatten() {
        let Ok(entries) = fs::read_dir(project.path()) else {
            continue;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().and_then(|value| value.to_str()) == Some("jsonl") {
                files.push(path);
            }
        }
    }
    files
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn model_and_agent_name_flags_drive_the_cli() {
    let (base_url, server) = mocked_llm_server().await;
    let session_root = tempfile::TempDir::new().unwrap();
    let session_override = format!("claude-code-cli={}", session_root.path().display());

    let mut child = Command::new(env!("CARGO_BIN_EXE_asyncai"))
        .args([
            "--base-url",
            &base_url,
            "--model",
            "flag-model",
            "--agent-name",
            "bob",
            "--session-dir",
            &session_override,
        ])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();

    child.stdin.as_mut().unwrap().write_all(b"ping\n").unwrap();
    drop(child.stdin.take());

    let output = child.wait_with_output().unwrap();
    assert!(
        output.status.success(),
        "cli failed: {}",
        String::from_utf8_lossy(&output.stderr)
    );

    let stdout = String::from_utf8_lossy(&output.stdout);
    let records: Vec<Value> = stdout
        .lines()
        .filter_map(|line| serde_json::from_str(line).ok())
        .collect();

    let agent = records
        .iter()
        .find(|record| record["result"]["event_type"] == "new_agent_created")
        .expect("agent created event");
    assert_eq!(agent["result"]["event_data"]["agent_name"], "bob");

    let message = records
        .iter()
        .find(|record| record["result"]["event_type"] == "new_thread_message")
        .expect("thread message event");
    assert_eq!(
        message["result"]["event_data"]["message_content"]["text"],
        "mocked pong"
    );

    let request_text = server.await.unwrap();
    assert!(request_text.contains("POST /v1/chat/completions"));
    assert!(request_text.contains("\"model\":\"flag-model\""));
    assert!(request_text.contains("\"content\":\"ping\""));
}

#[test]
fn help_lists_model_agent_name_and_atirc() {
    let output = Command::new(env!("CARGO_BIN_EXE_asyncai"))
        .arg("--help")
        .output()
        .unwrap();
    assert!(output.status.success());
    let stdout = String::from_utf8_lossy(&output.stdout);
    assert!(stdout.contains("--model"));
    assert!(stdout.contains("--agent-name"));
    assert!(stdout.contains("--mcp-stdio"));
    assert!(stdout.contains("--monitor"));
    assert!(stdout.contains("atirc"));
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn cli_monitor_captures_stdout_as_task_notifications() {
    let session_root = tempfile::TempDir::new().unwrap();
    let session_override = format!("claude-code-cli={}", session_root.path().display());

    let mut child = Command::new(env!("CARGO_BIN_EXE_asyncai"))
        .args([
            "--session-dir",
            &session_override,
            "--monitor",
            "printf 'monitor-line\\n'",
        ])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();

    child.stdin.as_mut().unwrap().write_all(b"hi\n").unwrap();
    drop(child.stdin.take());

    let output = child.wait_with_output().unwrap();
    assert!(
        output.status.success(),
        "cli failed: {}",
        String::from_utf8_lossy(&output.stderr)
    );

    let files = session_files(session_root.path());
    assert_eq!(files.len(), 1);
    let records: Vec<Value> = fs::read_to_string(&files[0])
        .unwrap()
        .lines()
        .map(|line| serde_json::from_str(line).unwrap())
        .collect();

    let enqueue = records
        .iter()
        .find(|record| {
            record["type"] == "queue-operation"
                && record["operation"] == "enqueue"
                && record["content"]
                    .as_str()
                    .is_some_and(|content| content.contains("monitor-line"))
        })
        .expect("monitor enqueue record");
    let attachment = records
        .iter()
        .find(|record| record["type"] == "attachment")
        .expect("monitor attachment record");
    let remove = records
        .iter()
        .find(|record| record["type"] == "queue-operation" && record["operation"] == "remove")
        .expect("monitor remove record");

    assert!(enqueue.get("parentUuid").is_none());
    assert_eq!(attachment["attachment"]["type"], "queued_command");
    assert_eq!(
        remove["commandUuid"],
        attachment["attachment"]["source_uuid"]
    );
    assert_eq!(
        remove["deliveryId"],
        attachment["attachment"]["delivery_id"]
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cli_mcp_stdio_tool_call_executes_http_get() {
    let (target_url, target_hit) = loopback_http_target().await;
    let tool_call = json!({
        "choices": [{
            "message": {
                "role": "assistant",
                "content": null,
                "tool_calls": [{
                    "id": "call_1",
                    "type": "function",
                    "function": {
                        "name": "run_code",
                        "arguments": json!({
                            "code": format!(
                                "response = await http_get(url={target_url:?})\nresponse['body']"
                            ),
                        })
                        .to_string(),
                    }
                }]
            }
        }]
    })
    .to_string();
    let final_response = json!({
        "choices": [{ "message": { "role": "assistant", "content": "final answer" } }]
    })
    .to_string();

    let (base_url, server) = mocked_llm_sequence(vec![tool_call, final_response]).await;
    let session_root = tempfile::TempDir::new().unwrap();
    let session_override = format!("claude-code-cli={}", session_root.path().display());

    let mut child = Command::new(env!("CARGO_BIN_EXE_asyncai"))
        .args([
            "--base-url",
            &base_url,
            "--model",
            "tool-model",
            "--mcp-stdio",
            env!("CARGO_BIN_EXE_asyncai-mcp-simple"),
            "--code-sandbox",
            "in-process",
            "--session-dir",
            &session_override,
        ])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();

    child
        .stdin
        .as_mut()
        .unwrap()
        .write_all(b"fetch the target\n")
        .unwrap();
    drop(child.stdin.take());

    let output = child.wait_with_output().unwrap();
    assert!(
        output.status.success(),
        "cli failed: {}",
        String::from_utf8_lossy(&output.stderr)
    );

    let stdout = String::from_utf8_lossy(&output.stdout);
    let records: Vec<Value> = stdout
        .lines()
        .filter_map(|line| serde_json::from_str(line).ok())
        .collect();
    let message = records
        .iter()
        .find(|record| record["result"]["event_type"] == "new_thread_message")
        .expect("thread message event");
    assert_eq!(
        message["result"]["event_data"]["message_content"]["text"],
        "final answer"
    );

    assert!(
        target_hit.load(Ordering::SeqCst),
        "MCP tool never called loopback HTTP target"
    );

    let requests = server.await.unwrap();
    assert_eq!(requests.len(), 2);
    assert!(requests[0].contains("\"tools\""));
    assert!(requests[0].contains("\"name\":\"run_code\""));
    assert!(requests[0].contains("async def http_get(*, url: str) -> dict[str, Any]:"));
    assert!(requests[0]
        .contains("Returns: {\\\"body\\\": str, \\\"content_type\\\": str, \\\"status\\\": int}"));
    assert!(!requests[0].contains("\"name\":\"http_get\""));
    assert!(requests[1].contains("\"role\":\"tool\""));
    assert!(requests[1].contains("fetched-body"));

    let first: Value = serde_json::from_str(stdout.lines().next().unwrap()).unwrap();
    assert_eq!(first["work_name"], "session");
    let session_id = first["result"]["session_id"].as_str().unwrap();

    let files = session_files(session_root.path());
    assert_eq!(files.len(), 1, "expected one session transcript");
    assert_eq!(
        files[0].file_stem().and_then(|value| value.to_str()),
        Some(session_id)
    );
    let transcript: Vec<Value> = fs::read_to_string(&files[0])
        .unwrap()
        .lines()
        .map(|line| serde_json::from_str(line).unwrap())
        .collect();
    assert!(transcript.iter().any(|record| {
        record["type"] == "user" && record["message"]["content"] == "fetch the target"
    }));
    assert!(transcript.iter().any(|record| {
        record["type"] == "assistant"
            && record["message"]["content"]
                .as_array()
                .is_some_and(|content| {
                    content
                        .iter()
                        .any(|block| block["type"] == "tool_use" && block["name"] == "run_code")
                })
    }));
    assert!(transcript.iter().any(|record| {
        record["type"] == "user"
            && record["message"]["content"]
                .as_array()
                .is_some_and(|content| {
                    content.iter().any(|block| {
                        block["type"] == "tool_result"
                            && block["content"]
                                .as_str()
                                .is_some_and(|text| text.contains("fetched-body"))
                    })
                })
            && record["toolUseResult"]["nested_calls"][0]["tool"] == "http_get"
            && record["toolUseResult"]["nested_calls"][0]["mutating"] == false
    }));
    assert!(transcript.iter().any(|record| {
        record["type"] == "assistant"
            && record["message"]["content"]
                .as_array()
                .is_some_and(|content| content.iter().any(|block| block["text"] == "final answer"))
    }));
}

fn run_code_call(code: &str) -> String {
    json!({
        "choices": [{
            "message": {
                "role": "assistant",
                "content": null,
                "tool_calls": [{
                    "id": "call_watch",
                    "type": "function",
                    "function": {
                        "name": "run_code",
                        "arguments": json!({ "code": code }).to_string(),
                    }
                }]
            }
        }]
    })
    .to_string()
}

fn final_answer(text: &str) -> String {
    json!({"choices": [{"message": {"role": "assistant", "content": text}}]}).to_string()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cli_watch_notice_becomes_an_agent_turn() {
    let responses = vec![
        run_code_call(
            "await watch_start(name='notes', uri='memo://notes')\nawait memo_set(key='k', value='v')\n'watching'",
        ),
        final_answer("watching notes"),
        final_answer("note k was added"),
    ];
    let (base_url, server) = mocked_llm_sequence(responses).await;
    let session_root = tempfile::TempDir::new().unwrap();
    let session_override = format!("claude-code-cli={}", session_root.path().display());

    let mut child = Command::new(env!("CARGO_BIN_EXE_asyncai"))
        .args([
            "--base-url",
            &base_url,
            "--model",
            "watch-model",
            "--mcp-stdio",
            env!("CARGO_BIN_EXE_asyncai-mcp-simple"),
            "--code-sandbox",
            "in-process",
            "--session-dir",
            &session_override,
        ])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    child
        .stdin
        .as_mut()
        .unwrap()
        .write_all(b"watch the notes and tell me about changes\n")
        .unwrap();

    // The third model request is the watch-driven turn; stdin stays open
    // until it has happened.
    let requests = tokio::time::timeout(std::time::Duration::from_secs(20), server)
        .await
        .expect("watch notice never produced a model request")
        .unwrap();
    drop(child.stdin.take());
    let output = child.wait_with_output().unwrap();
    assert!(
        output.status.success(),
        "cli failed: {}",
        String::from_utf8_lossy(&output.stderr)
    );

    assert_eq!(requests.len(), 3);
    assert!(
        requests[2].contains(r#"<watch-event name=\"notes\" uri=\"memo://notes\">"#),
        "{}",
        requests[2]
    );
    assert!(requests[2].contains("ADDED k"), "{}", requests[2]);
    assert!(
        requests[2].contains("watching notes"),
        "history kept: {}",
        requests[2]
    );

    let stdout = String::from_utf8_lossy(&output.stdout);
    assert!(
        stdout.contains("\"work_name\":\"main.system.notice_stream\""),
        "{stdout}"
    );
    assert!(stdout.contains("note k was added"), "{stdout}");
    let stderr = String::from_utf8_lossy(&output.stderr);
    assert!(stderr.contains("\"event\":\"watch.notice\""), "{stderr}");
    assert!(stderr.contains("\"event\":\"watch.started\""), "{stderr}");

    let files = session_files(session_root.path());
    let transcript = fs::read_to_string(&files[0]).unwrap();
    assert!(
        transcript.contains("watch notes memo://notes"),
        "{transcript}"
    );
    let watch_notification = transcript
        .lines()
        .find(|line| line.contains("watch-notes") && line.contains("task-notification"))
        .expect("watch notification record");
    assert!(
        watch_notification.contains("reply exactly NO_REPLY")
            && !watch_notification.contains("PushNotification"),
        "{watch_notification}"
    );
    let records: Vec<Value> = transcript
        .lines()
        .map(|line| serde_json::from_str(line).unwrap())
        .collect();
    assert!(
        !records.iter().any(|record| record["type"] == "user"
            && record["message"]["content"]
                .as_str()
                .is_some_and(|text| text.contains("watch-event"))),
        "watch notices must not be user records"
    );
}