AkurAI Build
Menu

AkurAI-Build

public

Latest change fe39ce51fc9c202d5411f39fd9a8a2feeca4a85e - Harden test suite: 16→179 hermetic tests, race-free integration harness by Ólafur Búi Ólafsson

use std::io::{self, BufRead, Write};

use anyhow::{Context, Result, bail};
use serde::Deserialize;
use serde_json::{Value, json};

use crate::db::{Database, RepositoryQuery, RunQuery};

const PROTOCOL_VERSION: &str = "2025-06-18";

#[derive(Deserialize, Default)]
#[serde(deny_unknown_fields)]
struct RepositoryArguments {
    #[serde(alias = "q")]
    search: Option<String>,
    visibility: Option<String>,
    #[serde(default = "default_repository_limit")]
    limit: usize,
    #[serde(default)]
    offset: usize,
}

#[derive(Deserialize, Default)]
#[serde(deny_unknown_fields)]
struct RunArguments {
    repo: Option<StringList>,
    status: Option<StringList>,
    git_ref: Option<String>,
    trigger: Option<StringList>,
    #[serde(alias = "q")]
    search: Option<String>,
    #[serde(default = "default_run_limit")]
    limit: usize,
    #[serde(default)]
    offset: usize,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct RunDetailArguments {
    id: i64,
}

#[derive(Deserialize)]
#[serde(untagged)]
enum StringList {
    One(String),
    Many(Vec<String>),
}

impl StringList {
    fn values(self) -> Vec<String> {
        let values = match self {
            Self::One(value) => vec![value],
            Self::Many(values) => values,
        };
        values
            .into_iter()
            .flat_map(|value| {
                value
                    .split(',')
                    .map(str::trim)
                    .filter(|item| !item.is_empty())
                    .map(str::to_owned)
                    .collect::<Vec<_>>()
            })
            .collect()
    }
}

const fn default_repository_limit() -> usize {
    50
}

const fn default_run_limit() -> usize {
    20
}

pub fn serve(database: Database) -> Result<()> {
    let stdin = io::stdin();
    let mut stdout = io::stdout().lock();
    for line in stdin.lock().lines() {
        let line = line.context("read MCP request")?;
        if line.trim().is_empty() {
            continue;
        }
        let response = match serde_json::from_str::<Value>(&line) {
            Ok(request) => handle(&database, &request),
            Err(error) => Some(json_rpc_error(
                Value::Null,
                -32700,
                format!("invalid JSON: {error}"),
            )),
        };
        if let Some(response) = response {
            serde_json::to_writer(&mut stdout, &response)?;
            stdout.write_all(b"\n")?;
            stdout.flush()?;
        }
    }
    Ok(())
}

fn handle(database: &Database, request: &Value) -> Option<Value> {
    let id = request.get("id")?.clone();
    let method = request.get("method").and_then(Value::as_str);
    let Some(method) = method else {
        return Some(json_rpc_error(id, -32600, "missing method"));
    };
    let response = match method {
        "initialize" => json_rpc_result(
            id,
            json!({
                "protocolVersion": PROTOCOL_VERSION,
                "capabilities": {"tools": {"listChanged": false}},
                "serverInfo": {"name": "akurai-build", "version": env!("CARGO_PKG_VERSION")},
                "instructions": "Read-only access to AkurAI Build repositories, runs, run evidence, and Titan workers."
            }),
        ),
        "ping" => json_rpc_result(id, json!({})),
        "tools/list" => json_rpc_result(id, json!({"tools": tools()})),
        "tools/call" => {
            let params = request.get("params").cloned().unwrap_or_else(|| json!({}));
            let name = params.get("name").and_then(Value::as_str);
            let arguments = params
                .get("arguments")
                .cloned()
                .unwrap_or_else(|| json!({}));
            match name {
                Some(name) => match call_tool(database, name, arguments) {
                    Ok(data) => json_rpc_result(id, tool_result(data, false)),
                    Err(error) => json_rpc_result(id, tool_result(json!(error.to_string()), true)),
                },
                None => json_rpc_error(id, -32602, "tool name is required"),
            }
        }
        _ => json_rpc_error(id, -32601, format!("unknown method: {method}")),
    };
    Some(response)
}

fn call_tool(database: &Database, name: &str, arguments: Value) -> Result<Value> {
    match name {
        "akurai_repositories" => {
            let arguments: RepositoryArguments = serde_json::from_value(arguments)?;
            Ok(serde_json::to_value(database.query_repositories(
                &RepositoryQuery {
                    search: arguments.search,
                    visibility: arguments.visibility,
                    limit: arguments.limit,
                    offset: arguments.offset,
                },
            )?)?)
        }
        "akurai_runs" => {
            let arguments: RunArguments = serde_json::from_value(arguments)?;
            Ok(serde_json::to_value(
                database.query_runs(&RunQuery {
                    repository: arguments.repo.map(|value| value.values().join(",")),
                    statuses: arguments.status.map(StringList::values).unwrap_or_default(),
                    git_ref: arguments.git_ref,
                    triggers: arguments
                        .trigger
                        .map(StringList::values)
                        .unwrap_or_default(),
                    search: arguments.search,
                    limit: arguments.limit,
                    offset: arguments.offset,
                })?,
            )?)
        }
        "akurai_run" => {
            let arguments: RunDetailArguments = serde_json::from_value(arguments)?;
            Ok(serde_json::to_value(database.detail(arguments.id)?)?)
        }
        "akurai_workers" => Ok(serde_json::to_value(database.workers()?)?),
        _ => bail!("unknown tool: {name}"),
    }
}

fn tools() -> Vec<Value> {
    vec![
        json!({
            "name": "akurai_repositories",
            "description": "Query registered Git repositories by text and visibility.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "search": {"type": "string", "description": "LIKE search across name, URL, and default branch."},
                    "visibility": {"type": "string", "enum": ["private", "public"]},
                    "limit": {"type": "integer", "minimum": 1, "maximum": 500, "default": 50},
                    "offset": {"type": "integer", "minimum": 0, "maximum": 10000, "default": 0}
                },
                "additionalProperties": false
            }
        }),
        json!({
            "name": "akurai_runs",
            "description": "Query runs across one, many, or all repositories with composable filters.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "repo": string_or_array("Repository name, comma-separated names, or an array."),
                    "status": string_or_array("Status or statuses: queued, running, waiting, succeeded, failed, canceled, interrupted."),
                    "git_ref": {"type": "string", "description": "Exact Git ref."},
                    "trigger": string_or_array("Trigger or triggers: manual, webhook, retry."),
                    "search": {"type": "string", "description": "LIKE search across repository, ref, commit SHA, and retained error."},
                    "limit": {"type": "integer", "minimum": 1, "maximum": 200, "default": 20},
                    "offset": {"type": "integer", "minimum": 0, "maximum": 10000, "default": 0}
                },
                "additionalProperties": false
            }
        }),
        json!({
            "name": "akurai_run",
            "description": "Read one run with jobs, bounded logs, artifacts, and deployments.",
            "inputSchema": {
                "type": "object",
                "properties": {"id": {"type": "integer", "minimum": 1}},
                "required": ["id"],
                "additionalProperties": false
            }
        }),
        json!({
            "name": "akurai_workers",
            "description": "List Titan worker status, capacity, current run, heartbeat, and completed-run count.",
            "inputSchema": {"type": "object", "properties": {}, "additionalProperties": false}
        }),
    ]
}

fn string_or_array(description: &str) -> Value {
    json!({
        "description": description,
        "oneOf": [
            {"type": "string"},
            {"type": "array", "items": {"type": "string"}, "maxItems": 32}
        ]
    })
}

fn tool_result(data: Value, is_error: bool) -> Value {
    let text = if is_error {
        data.as_str().unwrap_or("tool failed").to_owned()
    } else {
        serde_json::to_string_pretty(&data).unwrap_or_else(|_| "null".into())
    };
    json!({"content": [{"type": "text", "text": text}], "isError": is_error})
}

fn json_rpc_result(id: Value, result: Value) -> Value {
    json!({"jsonrpc": "2.0", "id": id, "result": result})
}

fn json_rpc_error(id: Value, code: i64, message: impl Into<String>) -> Value {
    json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message.into()}})
}

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

    const KEY: &str = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f";

    fn db() -> Result<Database> {
        let database = Database::memory(KEY)?;
        database.add_repository("alpha", "https://example.com/alpha.git", "main")?;
        database.add_repository("beta", "https://example.com/beta.git", "main")?;
        database.register_worker("titan-1", "titan", "docker", 1)?;
        Ok(database)
    }

    #[test]
    fn tool_listing_returns_expected_tool_names() -> Result<()> {
        let database = db()?;
        let listed = handle(
            &database,
            &json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}),
        )
        .context("tools/list returned no response")?;
        let tools = listed["result"]["tools"]
            .as_array()
            .context("tools should be an array")?;
        assert_eq!(tools.len(), 4, "should have 4 tools");
        let names: Vec<_> = tools
            .iter()
            .filter_map(|tool| tool["name"].as_str())
            .collect();
        assert!(names.contains(&"akurai_repositories"));
        assert!(names.contains(&"akurai_runs"));
        assert!(names.contains(&"akurai_run"));
        assert!(names.contains(&"akurai_workers"));
        Ok(())
    }

    #[test]
    fn known_read_only_tool_returns_well_formed_output() -> Result<()> {
        let database = db()?;
        // Query repositories.
        let repos_response = handle(
            &database,
            &json!({
                "jsonrpc":"2.0",
                "id":1,
                "method":"tools/call",
                "params":{"name":"akurai_repositories","arguments":{"limit":10,"offset":0}}
            }),
        )
        .context("akurai_repositories returned no response")?;
        assert_eq!(repos_response["result"]["isError"], false);
        let text = repos_response["result"]["content"][0]["text"]
            .as_str()
            .context("repos response has no text")?;
        let repos: Vec<Value> = serde_json::from_str(text)?;
        assert_eq!(repos.len(), 2);

        // Query workers.
        let workers_response = handle(
            &database,
            &json!({
                "jsonrpc":"2.0",
                "id":2,
                "method":"tools/call",
                "params":{"name":"akurai_workers","arguments":{}}
            }),
        )
        .context("akurai_workers returned no response")?;
        assert_eq!(workers_response["result"]["isError"], false);
        let text = workers_response["result"]["content"][0]["text"]
            .as_str()
            .context("workers response has no text")?;
        let workers: Vec<Value> = serde_json::from_str(text)?;
        assert_eq!(workers.len(), 1);
        assert_eq!(workers[0]["id"], "titan-1");
        Ok(())
    }

    #[test]
    fn unknown_tool_errors_cleanly() -> Result<()> {
        let database = db()?;
        let response = handle(
            &database,
            &json!({
                "jsonrpc":"2.0",
                "id":1,
                "method":"tools/call",
                "params":{"name":"nonexistent_tool","arguments":{}}
            }),
        )
        .context("unknown tool returned no response")?;
        assert!(
            response["result"]["isError"] == true,
            "unknown tool should be an error: {response}"
        );
        let text = response["result"]["content"][0]["text"]
            .as_str()
            .unwrap_or("");
        assert!(
            text.contains("unknown tool"),
            "error text should mention unknown tool: {text}"
        );
        Ok(())
    }

    #[test]
    fn missing_tool_name_returns_error() -> Result<()> {
        let database = db()?;
        let response = handle(
            &database,
            &json!({
                "jsonrpc":"2.0",
                "id":1,
                "method":"tools/call",
                "params":{"arguments":{}}
            }),
        )
        .context("missing tool name returned no response")?;
        assert_eq!(response["error"]["code"], -32602);
        Ok(())
    }

    #[test]
    fn unknown_method_returns_error() -> Result<()> {
        let database = db()?;
        let response = handle(
            &database,
            &json!({
                "jsonrpc":"2.0",
                "id":1,
                "method":"nonexistent/method"
            }),
        )
        .context("unknown method returned no response")?;
        assert_eq!(response["error"]["code"], -32601);
        Ok(())
    }

    #[test]
    fn initialize_returns_capabilities() -> Result<()> {
        let database = db()?;
        let response = handle(
            &database,
            &json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","clientInfo":{"name":"test"}}}),
        )
        .context("initialize returned no response")?;
        assert_eq!(response["result"]["protocolVersion"], PROTOCOL_VERSION);
        assert_eq!(response["result"]["serverInfo"]["name"], "akurai-build");
        Ok(())
    }
}