claude_agent/

inference.rs

// Layer 4: THE BRAIN CORE — Dual Model Inference
// Left hemisphere: Opus 4.6 (primary). Right hemisphere: Haiku 4.5 (auxiliary).
// Crystalline bridge: Prompt Cache (8.3x read ratio).
//
// Two backends:
//   1. CLI mode (default) — uses `claude -p` from Claude Code subscription
//   2. API mode — direct POST to api.anthropic.com (requires API credits)

use anyhow::Result;
use reqwest::Client;
use tokio::process::Command;

use crate::types::*;

pub const PRIMARY: ModelConfig = ModelConfig {
    id: "claude-opus-4-6",
    max_tokens: 16384,
};

pub const AUXILIARY: ModelConfig = ModelConfig {
    id: "claude-haiku-4-5-20251001",
    max_tokens: 8192,
};

const API_URL: &str = "https://api.anthropic.com/v1/messages";
const API_VERSION: &str = "2023-06-01";

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum InferenceBackend {
    Cli,  // claude -p (Claude Max subscription)
    Api,  // direct API (requires credits)
}

pub struct InferenceEngine {
    client: Client,
    api_key: Option<String>,
    backend: InferenceBackend,
}

impl InferenceEngine {
    pub fn new(api_key: Option<&str>, backend: InferenceBackend) -> Self {
        Self {
            client: Client::new(),
            api_key: api_key.map(String::from),
            backend,
        }
    }

    /// Streaming chat — displays text tokens as they arrive, returns accumulated response
    pub async fn chat_stream(
        &self,
        request: &InferenceRequest,
        on_text: &mut dyn FnMut(&str),
    ) -> Result<InferenceResponse> {
        match self.backend {
            InferenceBackend::Cli => self.chat_stream_cli(request, on_text).await,
            InferenceBackend::Api => self.chat_stream_api(request, on_text).await,
        }
    }

    /// API streaming — SSE from api.anthropic.com
    async fn chat_stream_api(
        &self,
        request: &InferenceRequest,
        on_text: &mut dyn FnMut(&str),
    ) -> Result<InferenceResponse> {
        let api_key = self.api_key.as_deref()
            .ok_or_else(|| anyhow::anyhow!("API backend requires ANTHROPIC_API_KEY"))?;

        let mut body = serde_json::to_value(request)?;
        body.as_object_mut().unwrap().insert("stream".to_string(), serde_json::Value::Bool(true));

        let response = self.client
            .post(API_URL)
            .header("x-api-key", api_key)
            .header("anthropic-version", API_VERSION)
            .header("content-type", "application/json")
            .json(&body)
            .send()
            .await?;

        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            anyhow::bail!("API error {status}: {body}");
        }

        parse_sse_stream(response, on_text).await
    }

    /// CLI streaming — incremental stdout from `claude -p --output-format stream-json`
    async fn chat_stream_cli(
        &self,
        request: &InferenceRequest,
        on_text: &mut dyn FnMut(&str),
    ) -> Result<InferenceResponse> {
        let prompt = extract_last_user_text(&request.messages);
        if prompt.is_empty() {
            anyhow::bail!("No user message to send");
        }

        let mut args = vec![
            "-p".to_string(),
            "--output-format".to_string(),
            "stream-json".to_string(),
            "--verbose".to_string(),
        ];

        if let Some(ref system) = request.system {
            args.push("--append-system-prompt".to_string());
            args.push(system.clone());
        }

        if let Some(ref tools) = request.tools {
            let tool_names: Vec<String> = tools.iter().map(|t| {
                match t.name.as_str() {
                    "bash" => "Bash".to_string(),
                    "read" => "Read".to_string(),
                    "edit" => "Edit".to_string(),
                    "write" => "Write".to_string(),
                    "glob" => "Glob".to_string(),
                    "grep" => "Grep".to_string(),
                    "web_fetch" => "WebFetch".to_string(),
                    other => other.to_string(),
                }
            }).collect();
            args.push("--allowedTools".to_string());
            args.push(tool_names.join(","));
        }

        // Prompt goes after -- to separate from flags
        args.push("--".to_string());
        args.push(prompt);

        use tokio::io::{AsyncBufReadExt, BufReader};

        let mut child = Command::new("claude")
            .args(&args)
            .stdin(std::process::Stdio::null())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped())
            .env_remove("CLAUDECODE")
            .env_remove("CLAUDE_CODE_ENTRYPOINT")
            .env_remove("ANTHROPIC_API_KEY")
            .spawn()?;

        let stdout = child.stdout.take()
            .ok_or_else(|| anyhow::anyhow!("Failed to capture CLI stdout"))?;
        let reader = BufReader::new(stdout);
        let mut lines = reader.lines();

        let mut result_event: Option<serde_json::Value> = None;
        let mut full_text = String::new();

        while let Ok(Some(line)) = lines.next_line().await {
            if line.trim().is_empty() { continue; }

            let event: serde_json::Value = match serde_json::from_str(&line) {
                Ok(v) => v,
                Err(_) => continue,
            };

            let event_type = event.get("type").and_then(|v| v.as_str()).unwrap_or("");

            match event_type {
                "content_block_delta" => {
                    if let Some(delta) = event.get("delta") {
                        if let Some(text) = delta.get("text").and_then(|v| v.as_str()) {
                            on_text(text);
                            full_text.push_str(text);
                        }
                    }
                }
                "assistant" => {
                    // stream-json wraps the response in {"type":"assistant","message":{...}}
                    if let Some(message) = event.get("message") {
                        if let Some(content) = message.get("content").and_then(|c| c.as_array()) {
                            for block in content {
                                if let Some(text) = block.get("text").and_then(|t| t.as_str()) {
                                    on_text(text);
                                    full_text.push_str(text);
                                }
                            }
                        }
                        // Check for billing errors
                        if let Some(err) = event.get("error").and_then(|e| e.as_str()) {
                            if err == "billing_error" {
                                let msg = full_text.clone();
                                if !msg.is_empty() {
                                    anyhow::bail!("{msg}");
                                }
                            }
                        }
                    }
                }
                "result" => {
                    // Check for billing/auth errors in result
                    if event.get("is_error").and_then(|v| v.as_bool()).unwrap_or(false) {
                        let msg = event.get("result").and_then(|v| v.as_str()).unwrap_or("Unknown error");
                        if msg.contains("Credit balance") || msg.contains("billing") || msg.contains("auth") {
                            anyhow::bail!("{msg}. Check your Claude subscription or API credits.");
                        }
                    }
                    result_event = Some(event);
                }
                _ => {}
            }
        }

        let _status = child.wait().await?;

        // Extract usage from result event
        let usage = result_event.as_ref().map(|r| {
            let u = r.get("usage").unwrap_or(&serde_json::Value::Null);
            Usage {
                input_tokens: u.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0),
                output_tokens: u.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0),
                cache_read_input_tokens: u.get("cache_read_input_tokens").and_then(|v| v.as_u64()).unwrap_or(0),
                cache_creation_input_tokens: u.get("cache_creation_input_tokens").and_then(|v| v.as_u64()).unwrap_or(0),
            }
        }).unwrap_or_default();

        let cost = result_event.as_ref()
            .and_then(|r| r.get("total_cost_usd"))
            .and_then(|v| v.as_f64());

        let duration_ms = result_event.as_ref()
            .and_then(|r| r.get("duration_api_ms"))
            .and_then(|v| v.as_u64())
            .unwrap_or(0);

        if !full_text.is_empty() {
            Ok(InferenceResponse {
                content: vec![ContentBlock::Text { text: full_text }],
                stop_reason: Some("end_turn".to_string()),
                usage,
                model: "cli-stream".to_string(),
                cli_meta: Some(CliMeta {
                    cost_usd: cost.unwrap_or(0.0),
                    duration_ms,
                    duration_api_ms: result_event.as_ref()
                        .and_then(|r| r.get("duration_api_ms"))
                        .and_then(|v| v.as_u64())
                        .unwrap_or(0),
                    num_turns: result_event.as_ref()
                        .and_then(|r| r.get("num_turns"))
                        .and_then(|v| v.as_u64())
                        .unwrap_or(1),
                }),
            })
        } else if let Some(ref r) = result_event {
            // Result exists but no text — check for error
            let result_text = r.get("result").and_then(|v| v.as_str()).unwrap_or("");
            if r.get("is_error").and_then(|v| v.as_bool()).unwrap_or(false) {
                anyhow::bail!("Claude CLI error: {result_text}");
            }
            Ok(InferenceResponse {
                content: vec![ContentBlock::Text { text: result_text.to_string() }],
                stop_reason: Some("end_turn".to_string()),
                usage,
                model: "cli-stream".to_string(),
                cli_meta: None,
            })
        } else {
            anyhow::bail!("CLI stream ended with no output")
        }
    }

    pub fn build_request(
        &self,
        messages: &[Message],
        system: Option<&str>,
        tools: &[ToolDefinition],
        model: Option<&str>,
    ) -> InferenceRequest {
        InferenceRequest {
            model: model.unwrap_or(PRIMARY.id).to_string(),
            max_tokens: PRIMARY.max_tokens,
            messages: messages.to_vec(),
            system: system.map(String::from),
            tools: if tools.is_empty() {
                None
            } else {
                Some(tools.to_vec())
            },
        }
    }

    /// Auto-route: pick Opus for complex tasks, Haiku for simple ones.
    /// Returns the model ID to use based on the last user message.
    pub fn auto_route(&self, messages: &[Message]) -> &'static str {
        let last_text = extract_last_user_text(messages);
        let complexity = estimate_complexity(&last_text);
        if complexity >= ComplexityLevel::High {
            PRIMARY.id
        } else {
            AUXILIARY.id
        }
    }
}

fn extract_last_user_text(messages: &[Message]) -> String {
    for msg in messages.iter().rev() {
        if matches!(msg.role, Role::User) {
            match &msg.content {
                MessageContent::Text(t) => return t.clone(),
                MessageContent::Blocks(blocks) => {
                    for block in blocks {
                        if let ContentBlock::Text { text } = block {
                            return text.clone();
                        }
                    }
                }
            }
        }
    }
    String::new()
}

/// Parse an SSE stream from the Anthropic Messages API
async fn parse_sse_stream(
    response: reqwest::Response,
    on_text: &mut dyn FnMut(&str),
) -> Result<InferenceResponse> {
    use futures_util::StreamExt;

    let mut stream = response.bytes_stream();
    let mut buffer = String::new();

    let mut content_blocks: Vec<ContentBlock> = Vec::new();
    let mut current_text = String::new();
    let mut current_tool_id = String::new();
    let mut current_tool_name = String::new();
    let mut current_tool_input = String::new();
    let mut in_tool = false;
    let mut stop_reason: Option<String> = None;
    let mut model = String::new();
    let mut usage = Usage::default();

    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        buffer.push_str(&String::from_utf8_lossy(&chunk));

        while let Some(pos) = buffer.find("\n\n") {
            let event_block = buffer[..pos].to_string();
            buffer = buffer[pos + 2..].to_string();

            let data = event_block.lines()
                .find(|l| l.starts_with("data: "))
                .map(|l| &l[6..]);

            let Some(data) = data else { continue };
            if data == "[DONE]" { continue; }

            let event: serde_json::Value = match serde_json::from_str(data) {
                Ok(v) => v,
                Err(_) => continue,
            };

            let etype = event.get("type").and_then(|v| v.as_str()).unwrap_or("");

            match etype {
                "message_start" => {
                    if let Some(msg) = event.get("message") {
                        model = msg.get("model").and_then(|v| v.as_str()).unwrap_or("").to_string();
                        if let Some(u) = msg.get("usage") {
                            usage.input_tokens = u.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
                            usage.cache_read_input_tokens = u.get("cache_read_input_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
                            usage.cache_creation_input_tokens = u.get("cache_creation_input_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
                        }
                    }
                }
                "content_block_start" => {
                    if let Some(cb) = event.get("content_block") {
                        match cb.get("type").and_then(|v| v.as_str()) {
                            Some("text") => { in_tool = false; }
                            Some("tool_use") => {
                                in_tool = true;
                                current_tool_id = cb.get("id").and_then(|v| v.as_str()).unwrap_or("").to_string();
                                current_tool_name = cb.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string();
                                current_tool_input.clear();
                            }
                            _ => {}
                        }
                    }
                }
                "content_block_delta" => {
                    if let Some(delta) = event.get("delta") {
                        match delta.get("type").and_then(|v| v.as_str()) {
                            Some("text_delta") => {
                                if let Some(text) = delta.get("text").and_then(|v| v.as_str()) {
                                    on_text(text);
                                    current_text.push_str(text);
                                }
                            }
                            Some("input_json_delta") => {
                                if let Some(json) = delta.get("partial_json").and_then(|v| v.as_str()) {
                                    current_tool_input.push_str(json);
                                }
                            }
                            _ => {}
                        }
                    }
                }
                "content_block_stop" => {
                    if !in_tool && !current_text.is_empty() {
                        content_blocks.push(ContentBlock::Text { text: current_text.clone() });
                        current_text.clear();
                    }
                    if in_tool && !current_tool_name.is_empty() {
                        let input = serde_json::from_str(&current_tool_input)
                            .unwrap_or(serde_json::Value::Object(Default::default()));
                        content_blocks.push(ContentBlock::ToolUse {
                            id: current_tool_id.clone(),
                            name: current_tool_name.clone(),
                            input,
                        });
                        current_tool_name.clear();
                        current_tool_input.clear();
                        in_tool = false;
                    }
                }
                "message_delta" => {
                    if let Some(delta) = event.get("delta") {
                        stop_reason = delta.get("stop_reason")
                            .and_then(|v| v.as_str())
                            .map(String::from);
                    }
                    if let Some(u) = event.get("usage") {
                        usage.output_tokens = u.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0);
                    }
                }
                _ => {}
            }
        }
    }

    if !current_text.is_empty() {
        content_blocks.push(ContentBlock::Text { text: current_text });
    }

    Ok(InferenceResponse {
        content: content_blocks,
        stop_reason,
        usage,
        model,
        cli_meta: None,
    })
}

// ─── Multi-Model Routing ────────────────────────────────────────

#[derive(Debug, PartialEq, PartialOrd)]
enum ComplexityLevel {
    Low,    // Simple questions, greetings, one-liners
    Medium, // Moderate tasks, single-file edits
    High,   // Multi-file changes, architecture, debugging
}

fn estimate_complexity(text: &str) -> ComplexityLevel {
    let lower = text.to_lowercase();
    let words: Vec<&str> = text.split_whitespace().collect();
    let word_count = words.len();

    // High complexity signals
    let high_signals = [
        "refactor", "architect", "design", "migrate", "implement",
        "debug", "investigate", "analyze", "review", "security",
        "optimize", "performance", "deploy", "infrastructure",
        "multiple files", "multi-file", "across the codebase",
    ];
    let has_high = high_signals.iter().any(|s| lower.contains(s));

    // Low complexity signals
    let low_signals = [
        "what is", "how do", "explain", "hello", "hi ", "thanks",
        "what's", "define", "list", "show me", "tell me",
    ];
    let has_low = low_signals.iter().any(|s| lower.contains(s));

    // Code blocks or file paths suggest medium+
    let has_code = text.contains("```") || text.contains("fn ") || text.contains("def ");
    let has_paths = text.contains('/') && text.contains('.');

    if has_high || word_count > 50 || (has_code && has_paths) {
        ComplexityLevel::High
    } else if has_low && word_count < 15 && !has_code {
        ComplexityLevel::Low
    } else {
        ComplexityLevel::Medium
    }
}