johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::sync::{Arc, Mutex};

use async_trait::async_trait;
use asyncai::backend::{AgentBackend, AgentSpec, BackendEvent, RunRequest};
use asyncai::message::{AgentContent, AgentMessage, Role, RunStatus};
use asyncai::openai_backend::OpenAiCompatibleBackend;
use asyncai::tools::{ToolCallOutput, ToolDefinition, ToolError, ToolExecutor};
use asyncai::CodeModeToolExecutor;
use futures::StreamExt;
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) -> Vec<u8> {
    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;
            }
        }
    }

    request
}

async fn write_http_response(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<&'static str>,
) -> (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();
            let request = read_http_request(&mut socket).await;
            requests.push(String::from_utf8_lossy(&request).to_string());
            write_http_response(&mut socket, response).await;
        }
        requests
    });

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

#[derive(Default)]
struct RecordingTools {
    calls: Mutex<Vec<(String, serde_json::Value)>>,
}

#[async_trait]
impl ToolExecutor for RecordingTools {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        Ok(vec![ToolDefinition::new(
            "http_get",
            Some("HTTP GET"),
            serde_json::json!({
                "type": "object",
                "properties": { "url": { "type": "string" } },
                "required": ["url"],
            }),
        )])
    }

    async fn call_tool(
        &self,
        name: &str,
        arguments: serde_json::Value,
    ) -> Result<ToolCallOutput, ToolError> {
        self.calls
            .lock()
            .unwrap()
            .push((name.to_string(), arguments.clone()));
        Ok(ToolCallOutput::text(format!(
            "tool output for {}",
            arguments["url"].as_str().unwrap_or_default()
        )))
    }
}

#[tokio::test]
async fn parses_mocked_openai_compatible_response() {
    let (base_url, server) = mocked_llm_server().await;
    let backend = OpenAiCompatibleBackend::new(base_url, "mock-model");

    let request = RunRequest {
        thread_id: None,
        agent: AgentSpec {
            agent_id: "agent-1".to_string(),
            agent_name: "alice".to_string(),
            instructions: "answer briefly".to_string(),
        },
        messages: vec![AgentMessage {
            role: Role::User,
            content: AgentContent::Text("ping".to_string()),
        }],
    };

    let mut stream = backend.start_run(request).await.unwrap();

    let first = stream.next().await.unwrap().unwrap();
    match first {
        BackendEvent::Message(message) => {
            assert_eq!(message.role, Role::Assistant);
            assert_eq!(
                message.content,
                AgentContent::Text("mocked pong".to_string())
            );
        }
        BackendEvent::Completed(_) => panic!("expected a message first"),
    }

    let second = stream.next().await.unwrap().unwrap();
    match second {
        BackendEvent::Completed(outcome) => {
            assert_eq!(outcome.status, RunStatus::Completed);
            assert_eq!(outcome.messages.len(), 2);
            assert_eq!(
                outcome.final_output,
                Some(AgentContent::Text("mocked pong".to_string()))
            );
        }
        BackendEvent::Message(_) => panic!("expected completion second"),
    }

    assert!(stream.next().await.is_none());

    let request_text = server.await.unwrap();
    assert!(request_text.contains("POST /v1/chat/completions"));
    assert!(request_text.contains("\"model\":\"mock-model\""));
    assert!(request_text.contains("\"content\":\"ping\""));
    assert!(
        request_text.contains("\"max_tokens\":8192"),
        "a cap ends runaway generations: {request_text}"
    );
}

#[tokio::test]
async fn executes_tool_calls_and_returns_final_answer() {
    let tool_call_response = r#"{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call_1","type":"function","function":{"name":"http_get","arguments":"{\"url\":\"http://example.invalid/data\"}"}}]}}]}"#;
    let final_response =
        r#"{"choices":[{"message":{"role":"assistant","content":"final answer"}}]}"#;
    let (base_url, server) = mocked_llm_sequence(vec![tool_call_response, final_response]).await;

    let recorder = Arc::new(RecordingTools::default());
    let backend =
        OpenAiCompatibleBackend::new(base_url, "mock-model").with_tool_executor(recorder.clone());

    let request = RunRequest {
        thread_id: None,
        agent: AgentSpec {
            agent_id: "agent-1".to_string(),
            agent_name: "alice".to_string(),
            instructions: "use tools when useful".to_string(),
        },
        messages: vec![AgentMessage {
            role: Role::User,
            content: AgentContent::Text("fetch it".to_string()),
        }],
    };

    let mut stream = backend.start_run(request).await.unwrap();
    let message = match stream.next().await.unwrap().unwrap() {
        BackendEvent::Message(message) => message,
        BackendEvent::Completed(_) => panic!("expected message first"),
    };
    assert_eq!(
        message.content,
        AgentContent::Text("final answer".to_string())
    );

    match stream.next().await.unwrap().unwrap() {
        BackendEvent::Completed(outcome) => {
            assert_eq!(outcome.status, RunStatus::Completed);
            assert_eq!(
                outcome.final_output,
                Some(AgentContent::Text("final answer".to_string()))
            );
        }
        BackendEvent::Message(_) => panic!("expected completion second"),
    }
    assert!(stream.next().await.is_none());

    let requests = server.await.unwrap();
    assert_eq!(requests.len(), 2);
    assert!(requests[0].contains("\"tools\""));
    assert!(requests[0].contains("http_get"));
    assert!(requests[1].contains("\"role\":\"tool\""));
    assert!(requests[1].contains("\"tool_call_id\":\"call_1\""));
    assert!(
        requests[1].contains("tool output for http://example.invalid/data"),
        "second request: {}",
        requests[1]
    );

    let calls = recorder.calls.lock().unwrap();
    assert_eq!(calls.len(), 1);
    assert_eq!(calls[0].0, "http_get");
    assert_eq!(calls[0].1["url"], "http://example.invalid/data");
}

#[tokio::test]
async fn last_round_withdraws_tools_so_the_model_must_answer() {
    let tool_call_response = r#"{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call_1","type":"function","function":{"name":"http_get","arguments":"{\"url\":\"http://example.invalid/data\"}"}}]}}]}"#;
    let final_response =
        r#"{"choices":[{"message":{"role":"assistant","content":"best effort answer"}}]}"#;
    let mut responses = vec![tool_call_response; 8];
    responses.push(final_response);
    let (base_url, server) = mocked_llm_sequence(responses).await;

    let recorder = Arc::new(RecordingTools::default());
    let backend =
        OpenAiCompatibleBackend::new(base_url, "mock-model").with_tool_executor(recorder.clone());
    let request = RunRequest {
        thread_id: None,
        agent: AgentSpec {
            agent_id: "agent-1".to_string(),
            agent_name: "alice".to_string(),
            instructions: "use tools".to_string(),
        },
        messages: vec![AgentMessage {
            role: Role::User,
            content: AgentContent::Text("loop forever".to_string()),
        }],
    };

    let mut stream = backend.start_run(request).await.unwrap();
    match stream.next().await.unwrap().unwrap() {
        BackendEvent::Message(message) => assert_eq!(
            message.content,
            AgentContent::Text("best effort answer".to_string())
        ),
        BackendEvent::Completed(_) => panic!("expected message first"),
    }

    let requests = server.await.unwrap();
    assert_eq!(requests.len(), 9);
    assert!(requests[7].contains("\"tools\""));
    assert!(!requests[8].contains("\"tools\""), "{}", requests[8]);
    assert!(requests[8].contains("Tool budget exhausted"));
    assert!(
        !requests[8].contains("\"role\":\"system\",\"content\":\"Tool budget"),
        "chat templates reject late system messages"
    );
    assert_eq!(recorder.calls.lock().unwrap().len(), 8);
}

fn run_code_response(call_id: &str, code: &str) -> &'static str {
    let response = serde_json::json!({
        "choices": [{
            "message": {
                "role": "assistant",
                "content": null,
                "tool_calls": [{
                    "id": call_id,
                    "type": "function",
                    "function": {
                        "name": "run_code",
                        "arguments": serde_json::json!({"code": code}).to_string(),
                    },
                }],
            },
        }],
    });
    Box::leak(response.to_string().into_boxed_str())
}

#[tokio::test]
async fn code_mode_runs_python_and_keeps_state_across_rounds() {
    let first = run_code_response(
        "call_1",
        "data = await http_get(url='http://example.invalid/data')\nlen(data)",
    );
    let second = run_code_response("call_2", "data.upper()");
    let final_response =
        r#"{"choices":[{"message":{"role":"assistant","content":"final answer"}}]}"#;
    let (base_url, server) = mocked_llm_sequence(vec![first, second, final_response]).await;

    let recorder = Arc::new(RecordingTools::default());
    let code_mode = CodeModeToolExecutor::new(recorder.clone()).await.unwrap();
    let backend = OpenAiCompatibleBackend::new(base_url, "mock-model")
        .with_tool_executor(Arc::new(code_mode));

    let request = RunRequest {
        thread_id: None,
        agent: AgentSpec {
            agent_id: "agent-1".to_string(),
            agent_name: "alice".to_string(),
            instructions: "use tools when useful".to_string(),
        },
        messages: vec![AgentMessage {
            role: Role::User,
            content: AgentContent::Text("fetch it".to_string()),
        }],
    };

    let mut stream = backend.start_run(request).await.unwrap();
    match stream.next().await.unwrap().unwrap() {
        BackendEvent::Message(message) => assert_eq!(
            message.content,
            AgentContent::Text("final answer".to_string())
        ),
        BackendEvent::Completed(_) => panic!("expected message first"),
    }

    let requests = server.await.unwrap();
    assert_eq!(requests.len(), 3);
    assert!(requests[0].contains("\"name\":\"run_code\""));
    assert!(
        !requests[0].contains("\"name\":\"http_get\""),
        "MCP tools must not be exposed directly"
    );
    assert!(requests[0].contains("async def http_get(*, url: str) -> Any:"));
    assert!(requests[1].contains("\"tool_call_id\":\"call_1\""));
    let tool_output = "tool output for http://example.invalid/data";
    assert!(
        requests[1].contains(&format!("\"content\":\"{}\"", tool_output.len())),
        "{}",
        requests[1]
    );
    assert!(
        requests[2].contains("TOOL OUTPUT FOR HTTP://EXAMPLE.INVALID/DATA"),
        "REPL state was not kept across rounds: {}",
        requests[2]
    );

    let calls = recorder.calls.lock().unwrap();
    assert_eq!(calls.len(), 1, "second round must reuse the variable");
}