AkurAI-Router / src/claude_cli.rs

Latest change 16baa88e7934816edd8435507847f72fc196163f - OLI-301: claude-cli runner backend + static routes.conf (stale window 0) by Olafur Bui

use std::io::Write;
use std::process::{Command, Stdio};

use crate::config::Config;
use crate::http::{self, Request};
use crate::json::{self, Json};
use crate::util::now_secs;

const ENV_PASSTHROUGH: [&str; 2] = ["PATH", "LANG"];

pub struct CliPrompt {
    pub system: String,
    pub prompt: String,
}

fn message_text(content: Option<&Json>) -> String {
    match content {
        Some(Json::String(text)) => text.clone(),
        Some(Json::Array(parts)) => parts
            .iter()
            .filter_map(|part| part.get_str("text"))
            .collect::<Vec<_>>()
            .join("\n"),
        _ => String::new(),
    }
}

pub fn build_prompt(body: &Json) -> CliPrompt {
    let mut system = Vec::new();
    let mut turns = Vec::new();
    if let Some(Json::Array(messages)) = body.get("messages") {
        for message in messages {
            let text = message_text(message.get("content"));
            if text.trim().is_empty() {
                continue;
            }
            match message.get_str("role").unwrap_or("user") {
                "system" | "developer" => system.push(text),
                "assistant" => turns.push(format!("Assistant: {text}")),
                _ => turns.push(format!("User: {text}")),
            }
        }
    }
    let prompt = if turns.len() == 1 && turns[0].starts_with("User: ") {
        turns[0]["User: ".len()..].to_string()
    } else {
        turns.join("\n\n")
    };
    CliPrompt {
        system: system.join("\n\n"),
        prompt,
    }
}

fn result_event(stream_json: &str) -> Option<Json> {
    stream_json
        .lines()
        .filter_map(|line| json::parse(line).ok())
        .rfind(|event| event.get_str("type") == Some("result"))
}

fn usage_number(usage: Option<&Json>, key: &str) -> u64 {
    match usage.and_then(|u| u.get(key)) {
        Some(Json::Number(n)) => n.parse().unwrap_or(0),
        _ => 0,
    }
}

pub fn stream_json_to_completion(model: &str, stream_json: &str) -> Result<String, String> {
    let event = result_event(stream_json).ok_or("claude CLI produced no result event")?;
    let is_error = event.get_bool("is_error").unwrap_or(false);
    let text = event.get_str("result").unwrap_or("");
    if is_error || event.get_str("terminal_reason") == Some("api_error") {
        return Err(format!("claude CLI error: {text}"));
    }
    let usage = event.get("usage");
    let prompt_tokens = usage_number(usage, "input_tokens")
        + usage_number(usage, "cache_creation_input_tokens")
        + usage_number(usage, "cache_read_input_tokens");
    let completion_tokens = usage_number(usage, "output_tokens");
    let finish = if event.get_str("stop_reason") == Some("max_tokens") {
        "length"
    } else {
        "stop"
    };
    let mut message = Json::object();
    message.set("role", Json::String("assistant".to_string()));
    message.set("content", Json::String(text.to_string()));
    let mut choice = Json::object();
    choice.set("index", Json::Number("0".to_string()));
    choice.set("message", message);
    choice.set("finish_reason", Json::String(finish.to_string()));
    let mut usage_out = Json::object();
    usage_out.set("prompt_tokens", Json::Number(prompt_tokens.to_string()));
    usage_out.set(
        "completion_tokens",
        Json::Number(completion_tokens.to_string()),
    );
    usage_out.set(
        "total_tokens",
        Json::Number((prompt_tokens + completion_tokens).to_string()),
    );
    let mut root = Json::object();
    root.set("id", Json::String(format!("chatcmpl-cli-{}", now_secs())));
    root.set("object", Json::String("chat.completion".to_string()));
    root.set("created", Json::Number(now_secs().to_string()));
    root.set("model", Json::String(model.to_string()));
    root.set("choices", Json::Array(vec![choice]));
    root.set("usage", usage_out);
    Ok(root.stringify())
}

fn run_cli(cfg: &Config, model: &str, prompt: &CliPrompt) -> Result<String, String> {
    let mut cmd = Command::new(&cfg.claude_cli_bin);
    cmd.env_clear()
        .env("HOME", &cfg.claude_cli_home)
        .current_dir(&cfg.data_dir)
        .args(["-p", "--output-format", "stream-json", "--verbose"])
        .args(["--no-session-persistence", "--strict-mcp-config"])
        .args(["--disable-slash-commands", "--tools", ""])
        .args(["--model", model]);
    if !prompt.system.is_empty() {
        cmd.args(["--system-prompt", &prompt.system]);
    }
    for key in ENV_PASSTHROUGH {
        if let Ok(value) = std::env::var(key) {
            cmd.env(key, value);
        }
    }
    let mut child = cmd
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
        .map_err(|e| format!("failed to spawn claude CLI: {e}"))?;
    child
        .stdin
        .take()
        .ok_or("claude CLI stdin unavailable")?
        .write_all(prompt.prompt.as_bytes())
        .map_err(|e| e.to_string())?;
    let out = child.wait_with_output().map_err(|e| e.to_string())?;
    Ok(String::from_utf8_lossy(&out.stdout).to_string())
}

fn handle(req: Request, stream: &mut std::net::TcpStream, cfg: &Config) {
    let authorized = req
        .header("authorization")
        .and_then(|h| h.strip_prefix("Bearer "))
        .is_some_and(|token| token == cfg.claude_cli_token);
    if !authorized {
        let _ = http::send_json(stream, 401, "{\"error\":{\"message\":\"unauthorized\"}}");
        return;
    }
    if req.method != "POST" || req.path != "/v1/chat/completions" {
        let _ = http::send_json(stream, 404, "{\"error\":{\"message\":\"not found\"}}");
        return;
    }
    let Ok(body) = json::parse(&String::from_utf8_lossy(&req.body)) else {
        let _ = http::send_json(stream, 400, "{\"error\":{\"message\":\"invalid JSON\"}}");
        return;
    };
    let model = body
        .get_str("model")
        .unwrap_or(&cfg.default_model)
        .to_string();
    let prompt = build_prompt(&body);
    let result =
        run_cli(cfg, &model, &prompt).and_then(|output| stream_json_to_completion(&model, &output));
    match result {
        Ok(completion) => {
            let _ = http::send_json(stream, 200, &completion);
        }
        Err(message) => {
            let mut err = Json::object();
            err.set("message", Json::String(message));
            let mut root = Json::object();
            root.set("error", err);
            let _ = http::send_json(stream, 502, &root.stringify());
        }
    }
}

pub fn serve_runner(cfg: Config) -> Result<(), String> {
    if cfg.claude_cli_token.len() < 24 {
        return Err("AKURAI_ROUTER_CLAUDE_CLI_TOKEN must be at least 24 characters".to_string());
    }
    let addr = cfg.claude_runner_addr.clone();
    http::serve(&addr, move |req, stream| handle(req, stream, &cfg))
}

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

    #[test]
    fn single_user_message_is_passed_verbatim() {
        let body = json::parse(r#"{"messages":[{"role":"user","content":"hello"}]}"#).unwrap();
        let p = build_prompt(&body);
        assert_eq!(p.prompt, "hello");
        assert_eq!(p.system, "");
    }

    #[test]
    fn system_messages_split_out_and_turns_labelled() {
        let body = json::parse(
            r#"{"messages":[{"role":"system","content":"be brief"},{"role":"user","content":"a"},{"role":"assistant","content":"b"},{"role":"user","content":[{"type":"text","text":"c"}]}]}"#,
        )
        .unwrap();
        let p = build_prompt(&body);
        assert_eq!(p.system, "be brief");
        assert_eq!(p.prompt, "User: a\n\nAssistant: b\n\nUser: c");
    }

    #[test]
    fn result_event_maps_to_chat_completion() {
        let out = "{\"type\":\"system\"}\n{\"type\":\"result\",\"is_error\":false,\"result\":\"Hi!\",\"stop_reason\":\"end_turn\",\"usage\":{\"input_tokens\":9,\"cache_creation_input_tokens\":8000,\"output_tokens\":12}}\n";
        let parsed =
            json::parse(&stream_json_to_completion("claude-haiku-4-5", out).unwrap()).unwrap();
        let choice = match parsed.get("choices") {
            Some(Json::Array(c)) => c[0].clone(),
            _ => panic!("no choices"),
        };
        assert_eq!(
            choice.get("message").unwrap().get_str("content"),
            Some("Hi!")
        );
        assert_eq!(choice.get_str("finish_reason"), Some("stop"));
        assert_eq!(parsed.get_str("model"), Some("claude-haiku-4-5"));
        assert_eq!(
            parsed.get("usage").unwrap().get("total_tokens"),
            Some(&Json::Number("8021".to_string()))
        );
    }

    #[test]
    fn cli_errors_surface_as_errors() {
        let out = "{\"type\":\"result\",\"is_error\":true,\"result\":\"Not logged in\"}";
        assert!(
            stream_json_to_completion("m", out)
                .unwrap_err()
                .contains("Not logged in")
        );
        assert!(stream_json_to_completion("m", "garbage").is_err());
    }
}